Island: Replace AWSCommandRunner with start_infection_monkey_agent()

The AWSCommandRunner is a subclass of CmdRunner, which attempts to make
it easy to run commands on AWS nodes asynchronously. There are some
issues with the implementation, including unnecessary complexity and a
circular dependency between the CmdRunner and Cmd classes. A simple
function that runs a single command on a single instance is a simpler
solution. It can be run with a thread worker pool if asynchronicity is
required.

monkey/monkey_island/cc/services/aws/aws_command_runner.py

@@ -1,39 +1,103 @@
 import logging
 import time
 
+import botocore
+
+from common.utils import Timer
+
+REMOTE_COMMAND_TIMEOUT = 5
+STATUS_CHECK_SLEEP_TIME = 1
+
 logger = logging.getLogger(__name__)
 
 
-class AWSCommandRunner():
+def start_infection_monkey_agent(
+    aws_client: botocore.client.BaseClient, target_instance_id: str, target_os: str, island_ip: str
+):
     """
-    Class for running commands on a remote AWS machine
+    Run a command on a remote AWS instance
     """
+    command = _get_run_agent_command(target_os, island_ip)
+    command_id = _run_command_async(aws_client, target_instance_id, target_os, command)
+    _wait_for_command_to_complete(aws_client, target_instance_id, command_id)
 
-    def __init__(self, is_linux, instance_id, region=None):
-        self.instance_id = instance_id
-        self.region = region
-        self.ssm = aws_service.get_client("ssm", region)
 
-    def query_command(self, command_id):
-        time.sleep(2)
-        return self.ssm.get_command_invocation(CommandId=command_id, InstanceId=self.instance_id)
+def _get_run_agent_command(target_os: str, island_ip: str):
+    if target_os == "linux":
+        return _get_run_monkey_cmd_linux_line(island_ip)
 
-    def get_command_result(self, command_info):
-        return AwsCmdResult(command_info)
+    return _get_run_monkey_cmd_windows_line(island_ip)
 
-    def get_command_status(self, command_info):
-        if command_info["Status"] == "InProgress":
-            return CmdStatus.IN_PROGRESS
-        elif command_info["Status"] == "Success":
-            return CmdStatus.SUCCESS
-        else:
-            return CmdStatus.FAILURE
 
-    def run_command_async(self, command_line):
-        doc_name = "AWS-RunShellScript" if self.is_linux else "AWS-RunPowerShellScript"
-        command_res = self.ssm.send_command(
-            DocumentName=doc_name,
-            Parameters={"commands": [command_line]},
-            InstanceIds=[self.instance_id],
+def _get_run_monkey_cmd_linux_line(island_ip):
+    binary_name = "monkey-linux-64"
+
+    download_url = f"https://{island_ip}:5000/api/agent/download/linux"
+    download_cmd = f"wget --no-check-certificate {download_url} -O {binary_name}"
+
+    chmod_cmd = f"chmod +x {binary_name}"
+    run_agent_cmd = f"./{binary_name} m0nk3y -s {island_ip}:5000"
+
+    return f"{download_cmd}; {chmod_cmd}; {run_agent_cmd}"
+
+
+def _get_run_monkey_cmd_windows_line(island_ip):
+    agent_exe_path = r".\\monkey.exe"
+
+    ignore_ssl_errors_cmd = (
+        "[System.Net.ServicePointManager]::ServerCertificateValidationCallback = {$true}"
     )
-        return command_res["Command"]["CommandId"]
+
+    download_url = f"https://{island_ip}:5000/api/agent/download/windows"
+    download_cmd = (
+        f"(New-Object System.Net.WebClient).DownloadFile('{download_url}', '{agent_exe_path}')"
+    )
+
+    run_agent_cmd = (
+        f"Start-Process -FilePath '{agent_exe_path}' -ArgumentList 'm0nk3y -s {island_ip}:5000'"
+    )
+
+    return f"{ignore_ssl_errors_cmd}; {download_cmd}; {run_agent_cmd};"
+
+
+def _run_command_async(
+    aws_client: botocore.client.BaseClient, target_instance_id: str, target_os: str, command: str
+):
+    doc_name = "AWS-RunShellScript" if target_os == "linux" else "AWS-RunPowerShellScript"
+
+    logger.debug(f'Running command on {target_instance_id} -- {doc_name}: "{command}"')
+    command_response = aws_client.ssm.send_command(
+        DocumentName=doc_name,
+        Parameters={"commands": [command]},
+        InstanceIds=[target_instance_id],
+    )
+
+    command_id = command_response["CommandId"]
+    logger.debug(
+        f"Started command on AWS instance {target_instance_id} with command ID {command_id}"
+    )
+
+    return command_id
+
+
+def _wait_for_command_to_complete(
+    aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str
+):
+    timer = Timer()
+    timer.set(REMOTE_COMMAND_TIMEOUT)
+
+    while not timer.is_expired():
+        sleep_time = min((timer.time_remaining - STATUS_CHECK_SLEEP_TIME), STATUS_CHECK_SLEEP_TIME)
+        time.sleep(sleep_time)
+
+        command_status = aws_client.ssm.get_command_invocation(
+            CommandId=command_id, InstanceId=target_instance_id
+        )["Status"]
+        logger.debug(f"Command {command_id} status: {command_status}")
+
+        if command_status == "Success":
+            break
+
+        if command_status != "InProgress":
+            # TODO: Create an exception for this occasion and raise it with useful information.
+            raise Exception("COMMAND FAILED")

monkey/monkey_island/cc/services/aws/aws_service.py

@@ -58,9 +58,9 @@ class AWSService:
         :raises: botocore.exceptions.ClientError if can't describe local instance information.
         :return: All visible instances from this instance
         """
-        local_ssm_client = boto3.client("ssm", self.island_aws_instance.region)
+        ssm_client = boto3.client("ssm", self.island_aws_instance.region)
         try:
-            response = local_ssm_client.describe_instance_information()
+            response = ssm_client.describe_instance_information()
             return response[INSTANCE_INFORMATION_LIST_KEY]
         except botocore.exceptions.ClientError as err:
             logger.warning("AWS client error while trying to get manage dinstances: {err}")