From 653bfbd24bf0c062357544b8a4f747f6e08568cf Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 15:25:56 -0400 Subject: [PATCH] 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. --- .../cc/services/aws/aws_command_runner.py | 116 ++++++++++++++---- .../cc/services/aws/aws_service.py | 4 +- 2 files changed, 92 insertions(+), 28 deletions(-) diff --git a/monkey/monkey_island/cc/services/aws/aws_command_runner.py b/monkey/monkey_island/cc/services/aws/aws_command_runner.py index b07a74e55..9f5da8c87 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/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], - ) - return command_res["Command"]["CommandId"] +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}" + ) + + 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") diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index 1a8dec455..4f52583ab 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/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}")