Try for a graceful shutdown before forcing it

pwncollege · Sep 24, 2024 · d9bd4de · d9bd4de
1 parent 9f76a13
commit d9bd4de
Showing 1 changed file with 15 additions and 8 deletions.
diff --git a/dojo_plugin/utils/mac_docker.py b/dojo_plugin/utils/mac_docker.py
@@ -44,7 +44,7 @@ def __init__(self, hostname=None, username=None, key_filename=None, guest_key_fi
     def close(self):
         pass  # No persistent connection to close
 
-    def _ssh_exec(self, command, only_stdout=True, exception_on_fail=True, input=None, capture_output=True):
+    def _ssh_exec(self, command, only_stdout=True, exception_on_fail=True, input=None, capture_output=True, timeout_seconds=None):
         ssh_command = ['ssh',
                        "-a",
                        "-o", "StrictHostKeychecking=no",
@@ -68,7 +68,7 @@ def _ssh_exec(self, command, only_stdout=True, exception_on_fail=True, input=Non
             stdout_loc = None
             stderr_loc = None
 
-        result = subprocess.run(ssh_command, stdout=stdout_loc, stderr=stderr_loc, input=input)
+        result = subprocess.run(ssh_command, stdout=stdout_loc, stderr=stderr_loc, input=input, timeout=timeout_seconds)
         if result.returncode != 0:
             if exception_on_fail:
                 error_msg = result.stdout.strip()
@@ -138,10 +138,17 @@ def __init__(self, client, vm_info):
 
     def remove(self, force=True):
         # Kill the VM
-        command = f'{MAC_GUEST_CONTROL_FILE} kill-vm {self.id}'
-        exitcode, output = self.client._ssh_exec(command)
-        if b'Error' in output:
-            raise Exception(f'Error removing container: {output}')
+
+        # first try to shutdown the VM
+        try:
+            self.exec_run("shutdown -h now", "0", timeout_seconds=10)
+        except subprocess.TimeoutExpired:
+            # if that didn't work, kill it
+            if force:
+                command = f'{MAC_GUEST_CONTROL_FILE} kill-vm {self.id}'
+                exitcode, output = self.client._ssh_exec(command)
+                if b'Error' in output:
+                    raise Exception(f'Error removing container: {output}')
 
     def wait(self, condition='removed'):
         # Wait until the VM is removed
@@ -161,7 +168,7 @@ def start(self):
         pass
 
     # returns exit_code, output
-    def exec_run(self, cmd, user=None, input=None, capture_output=True, **kwargs):
+    def exec_run(self, cmd, user=None, input=None, capture_output=True, timeout_seconds=None, **kwargs):
         # SSH to the VM's IP address and run the command
         if user == "0" or user == None:
             # they want to run the command as root
@@ -170,7 +177,7 @@ def exec_run(self, cmd, user=None, input=None, capture_output=True, **kwargs):
             # they want to run the command as hacker
             cmd = f"exec sudo su - hacker -c {shlex.quote(cmd)}"
         command = f"{MAC_GUEST_CONTROL_FILE} exec {self.id} {shlex.quote(cmd)}"
-        exitcode, output = self.client._ssh_exec(command, only_stdout=False, exception_on_fail=False, input=input, capture_output=capture_output)
+        exitcode, output = self.client._ssh_exec(command, only_stdout=False, exception_on_fail=False, input=input, capture_output=capture_output, timeout_seconds=timeout_seconds)
         return exitcode, output
 
     # execve shell