From 32b4ed66e8f98a54dae68b9d88e99002ea27771b Mon Sep 17 00:00:00 2001
From: Joshua Boniface <joshua@boniface.me>
Date: Wed, 23 Oct 2019 23:37:42 -0400
Subject: [PATCH] Add VM disable command

Closes #53
---
 client-cli/pvc.py   | 20 ++++++++++++++++++++
 client-common/vm.py | 20 +++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/client-cli/pvc.py b/client-cli/pvc.py
index 217d2645..a1f0fa20 100755
--- a/client-cli/pvc.py
+++ b/client-cli/pvc.py
@@ -510,6 +510,25 @@ def vm_stop(domain):
     retcode, retmsg = pvc_vm.stop_vm(zk_conn, domain)
     cleanup(retcode, retmsg, zk_conn)
 
+###############################################################################
+# pvc vm disable
+###############################################################################
+@click.command(name='disable', short_help='Mark a virtual machine as disabled.')
+@click.argument(
+    'domain'
+)
+def vm_disable(domain):
+    """
+    Prevent stopped virtual machine DOMAIN from being counted towards cluster health status. DOMAIN may be a UUID or name.
+
+    Use this option for VM that are stopped intentionally or long-term and which should not impact cluster health if stopped. A VM can be started directly from disable state.
+    """
+
+    # Open a Zookeeper connection
+    zk_conn = pvc_common.startZKConnection(zk_host)
+    retcode, retmsg = pvc_vm.disable_vm(zk_conn, domain)
+    cleanup(retcode, retmsg, zk_conn)
+
 ###############################################################################
 # pvc vm move
 ###############################################################################
@@ -1833,6 +1852,7 @@ cli_vm.add_command(vm_start)
 cli_vm.add_command(vm_restart)
 cli_vm.add_command(vm_shutdown)
 cli_vm.add_command(vm_stop)
+cli_vm.add_command(vm_disable)
 cli_vm.add_command(vm_move)
 cli_vm.add_command(vm_migrate)
 cli_vm.add_command(vm_unmigrate)
diff --git a/client-common/vm.py b/client-common/vm.py
index 33d30e2a..16de1148 100644
--- a/client-common/vm.py
+++ b/client-common/vm.py
@@ -398,6 +398,23 @@ def stop_vm(zk_conn, domain):
 
     return True, 'Forcibly stopping VM "{}".'.format(domain)
 
+def disable_vm(zk_conn, domain):
+    # Validate that VM exists in cluster
+    dom_uuid = getDomainUUID(zk_conn, domain)
+    if not dom_uuid:
+        common.stopZKConnection(zk_conn)
+        return False, 'ERROR: Could not find VM "{}" in the cluster!'.format(domain)
+
+    # Get state and verify we're OK to proceed
+    current_state = zkhandler.readdata(zk_conn, '/domains/{}/state'.format(dom_uuid))
+    if current_state != 'stop':
+        return False, 'ERROR: VM "{}" must be stopped before disabling!'.format(domain)
+
+    # Set the VM to start
+    zkhandler.writedata(zk_conn, {'/domains/{}/state'.format(dom_uuid): 'disabled'})
+
+    return True, 'Marked VM "{}" as disabled.'.format(domain)
+
 def move_vm(zk_conn, domain, target_node):
     # Validate that VM exists in cluster
     dom_uuid = getDomainUUID(zk_conn, domain)
@@ -623,7 +640,7 @@ def get_list(zk_conn, node, state, limit, is_fuzzy=True):
             return False, 'Specified node "{}" is invalid.'.format(node)
 
     if state:
-        valid_states = [ 'start', 'restart', 'shutdown', 'stop', 'failed', 'migrate', 'unmigrate' ]
+        valid_states = [ 'start', 'restart', 'shutdown', 'stop', 'disabled', 'failed', 'migrate', 'unmigrate' ]
         if not state in valid_states:
             return False, 'VM state "{}" is not valid.'.format(state)
 
@@ -709,6 +726,7 @@ def format_info(zk_conn, domain_information, long_output):
         'restart': ansiprint.yellow(),
         'shutdown': ansiprint.yellow(),
         'stop': ansiprint.red(),
+        'disabled': ansiprint.blue(),
         'failed': ansiprint.red(),
         'migrate': ansiprint.blue(),
         'unmigrate': ansiprint.blue()