From f11ac5cc2bfc251eb820f372a1703ef170aef6e9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 12:59:28 -0400 Subject: [PATCH 01/34] Island: Move AWSService to service/aws/ --- monkey/monkey_island/cc/services/__init__.py | 6 +++++- monkey/monkey_island/cc/services/aws/__init__.py | 1 + monkey/monkey_island/cc/services/{ => aws}/aws_service.py | 0 3 files changed, 6 insertions(+), 1 deletion(-) create mode 100644 monkey/monkey_island/cc/services/aws/__init__.py rename monkey/monkey_island/cc/services/{ => aws}/aws_service.py (100%) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 7d96ee5c4..f50990ce3 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -4,4 +4,8 @@ from .directory_file_storage_service import DirectoryFileStorageService from .authentication.authentication_service import AuthenticationService from .authentication.json_file_user_datastore import JsonFileUserDatastore -from .aws_service import AWSService +from .aws import AWSService + +# TODO: This is a temporary import to keep some tests passing. Remove it before merging #1928 to +# develop. +from .aws import aws_service diff --git a/monkey/monkey_island/cc/services/aws/__init__.py b/monkey/monkey_island/cc/services/aws/__init__.py new file mode 100644 index 000000000..ff20c89a2 --- /dev/null +++ b/monkey/monkey_island/cc/services/aws/__init__.py @@ -0,0 +1 @@ +from .aws_service import AWSService diff --git a/monkey/monkey_island/cc/services/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py similarity index 100% rename from monkey/monkey_island/cc/services/aws_service.py rename to monkey/monkey_island/cc/services/aws/aws_service.py From 83f230c5a14c52319336ba5a1f11950d02129e7e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 13:00:55 -0400 Subject: [PATCH 02/34] Island: Move aws_cmd_runner.py -> aws_command_runner.py --- .../aws_cmd_runner.py => services/aws/aws_command_runner.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename monkey/monkey_island/cc/{server_utils/aws_cmd_runner.py => services/aws/aws_command_runner.py} (100%) diff --git a/monkey/monkey_island/cc/server_utils/aws_cmd_runner.py b/monkey/monkey_island/cc/services/aws/aws_command_runner.py similarity index 100% rename from monkey/monkey_island/cc/server_utils/aws_cmd_runner.py rename to monkey/monkey_island/cc/services/aws/aws_command_runner.py From a7dcc73a3da0469612c0bf1ea64bae0117becd65 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 13:01:34 -0400 Subject: [PATCH 03/34] Island: Rename AwsCmdRunner -> AWSCommandRunner --- .../monkey_island/cc/services/aws/aws_command_runner.py | 8 +------- monkey/monkey_island/cc/services/remote_run_aws.py | 4 ++-- 2 files changed, 3 insertions(+), 9 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 16c959197..b07a74e55 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -1,21 +1,15 @@ import logging import time -from common.cmd.cmd_runner import CmdRunner -from common.cmd.cmd_status import CmdStatus -from monkey_island.cc.server_utils.aws_cmd_result import AwsCmdResult -from monkey_island.cc.services import aws_service - logger = logging.getLogger(__name__) -class AwsCmdRunner(CmdRunner): +class AWSCommandRunner(): """ Class for running commands on a remote AWS machine """ def __init__(self, is_linux, instance_id, region=None): - super(AwsCmdRunner, self).__init__(is_linux) self.instance_id = instance_id self.region = region self.ssm = aws_service.get_client("ssm", region) diff --git a/monkey/monkey_island/cc/services/remote_run_aws.py b/monkey/monkey_island/cc/services/remote_run_aws.py index c3219171c..02c3b79fc 100644 --- a/monkey/monkey_island/cc/services/remote_run_aws.py +++ b/monkey/monkey_island/cc/services/remote_run_aws.py @@ -2,7 +2,7 @@ import logging from common.cmd.cmd import Cmd from common.cmd.cmd_runner import CmdRunner -from monkey_island.cc.server_utils.aws_cmd_runner import AwsCmdRunner +from monkey_island.cc.services.aws.aws_command_runner import AWSCommandRunner logger = logging.getLogger(__name__) @@ -40,7 +40,7 @@ class RemoteRunAwsService: @staticmethod def _run_aws_cmd_async(instance_id, is_linux, cmd_line): - cmd_runner = AwsCmdRunner(is_linux, instance_id) + cmd_runner = AWSCommandRunner(is_linux, instance_id) return Cmd(cmd_runner, cmd_runner.run_command_async(cmd_line)) @staticmethod From c24bb1024e290981b1de50620b4ff05498e4040e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 13:25:02 -0400 Subject: [PATCH 04/34] Agent: Move Timer to common/ --- monkey/common/utils/__init__.py | 1 + monkey/{infection_monkey => common}/utils/timer.py | 0 monkey/infection_monkey/exploit/log4shell.py | 2 +- monkey/infection_monkey/exploit/sshexec.py | 2 +- monkey/infection_monkey/master/automated_master.py | 2 +- monkey/infection_monkey/network_scanning/tcp_scanner.py | 2 +- .../telemetry/messengers/batching_telemetry_messenger.py | 2 +- monkey/infection_monkey/tunnel.py | 2 +- monkey/infection_monkey/utils/decorators.py | 2 +- .../unit_tests/{infection_monkey => common}/utils/test_timer.py | 2 +- .../tests/unit_tests/infection_monkey/utils/test_decorators.py | 2 +- 11 files changed, 10 insertions(+), 9 deletions(-) rename monkey/{infection_monkey => common}/utils/timer.py (100%) rename monkey/tests/unit_tests/{infection_monkey => common}/utils/test_timer.py (97%) diff --git a/monkey/common/utils/__init__.py b/monkey/common/utils/__init__.py index e69de29bb..57725fa62 100644 --- a/monkey/common/utils/__init__.py +++ b/monkey/common/utils/__init__.py @@ -0,0 +1 @@ +from .timer import Timer diff --git a/monkey/infection_monkey/utils/timer.py b/monkey/common/utils/timer.py similarity index 100% rename from monkey/infection_monkey/utils/timer.py rename to monkey/common/utils/timer.py diff --git a/monkey/infection_monkey/exploit/log4shell.py b/monkey/infection_monkey/exploit/log4shell.py index e967ee6cb..dc8fe0c66 100644 --- a/monkey/infection_monkey/exploit/log4shell.py +++ b/monkey/infection_monkey/exploit/log4shell.py @@ -3,6 +3,7 @@ import time from pathlib import PurePath from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT +from common.utils import Timer from infection_monkey.exploit.log4shell_utils import ( LINUX_EXPLOIT_TEMPLATE_PATH, WINDOWS_EXPLOIT_TEMPLATE_PATH, @@ -21,7 +22,6 @@ from infection_monkey.network.tools import get_interface_to_target from infection_monkey.utils.commands import build_monkey_commandline from infection_monkey.utils.monkey_dir import get_monkey_dir_path from infection_monkey.utils.threading import interruptible_iter -from infection_monkey.utils.timer import Timer logger = logging.getLogger(__name__) diff --git a/monkey/infection_monkey/exploit/sshexec.py b/monkey/infection_monkey/exploit/sshexec.py index 2b13b719e..3c4c46641 100644 --- a/monkey/infection_monkey/exploit/sshexec.py +++ b/monkey/infection_monkey/exploit/sshexec.py @@ -5,6 +5,7 @@ from pathlib import PurePath import paramiko from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT +from common.utils import Timer from common.utils.attack_utils import ScanStatus from common.utils.exceptions import FailedExploitationError from infection_monkey.exploit.HostExploiter import HostExploiter @@ -17,7 +18,6 @@ from infection_monkey.telemetry.attack.t1222_telem import T1222Telem from infection_monkey.utils.brute_force import generate_identity_secret_pairs from infection_monkey.utils.commands import build_monkey_commandline from infection_monkey.utils.threading import interruptible_iter -from infection_monkey.utils.timer import Timer logger = logging.getLogger(__name__) SSH_PORT = 22 diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py index b0ffa56fa..ac0429e1d 100644 --- a/monkey/infection_monkey/master/automated_master.py +++ b/monkey/infection_monkey/master/automated_master.py @@ -3,6 +3,7 @@ import threading import time from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple +from common.utils import Timer from infection_monkey.credential_store import ICredentialsStore from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.i_master import IMaster @@ -13,7 +14,6 @@ from infection_monkey.telemetry.credentials_telem import CredentialsTelem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger from infection_monkey.telemetry.post_breach_telem import PostBreachTelem from infection_monkey.utils.threading import create_daemon_thread, interruptible_iter -from infection_monkey.utils.timer import Timer from . import Exploiter, IPScanner, Propagator from .option_parsing import custom_pba_is_enabled diff --git a/monkey/infection_monkey/network_scanning/tcp_scanner.py b/monkey/infection_monkey/network_scanning/tcp_scanner.py index d0c6e3e7a..17c72f606 100644 --- a/monkey/infection_monkey/network_scanning/tcp_scanner.py +++ b/monkey/infection_monkey/network_scanning/tcp_scanner.py @@ -4,9 +4,9 @@ import socket import time from typing import Iterable, Mapping, Tuple +from common.utils import Timer from infection_monkey.i_puppet import PortScanData, PortStatus from infection_monkey.network.tools import BANNER_READ, DEFAULT_TIMEOUT, tcp_port_to_service -from infection_monkey.utils.timer import Timer logger = logging.getLogger(__name__) diff --git a/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py index 9dc051666..88c9d7d13 100644 --- a/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py +++ b/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py @@ -2,10 +2,10 @@ import queue import threading from typing import Dict +from common.utils import Timer from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem from infection_monkey.telemetry.i_telem import ITelem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger -from infection_monkey.utils.timer import Timer DEFAULT_PERIOD = 5 WAKES_PER_PERIOD = 4 diff --git a/monkey/infection_monkey/tunnel.py b/monkey/infection_monkey/tunnel.py index 26368bff6..b4ca3b517 100644 --- a/monkey/infection_monkey/tunnel.py +++ b/monkey/infection_monkey/tunnel.py @@ -4,11 +4,11 @@ import struct import time from threading import Event, Thread +from common.utils import Timer from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_free_tcp_port, local_ips from infection_monkey.network.tools import check_tcp_port, get_interface_to_target from infection_monkey.transport.base import get_last_serve_time -from infection_monkey.utils.timer import Timer logger = logging.getLogger(__name__) diff --git a/monkey/infection_monkey/utils/decorators.py b/monkey/infection_monkey/utils/decorators.py index 31ac0661b..363f29796 100644 --- a/monkey/infection_monkey/utils/decorators.py +++ b/monkey/infection_monkey/utils/decorators.py @@ -1,7 +1,7 @@ import threading from functools import wraps -from .timer import Timer +from common.utils import Timer def request_cache(ttl: float): diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_timer.py b/monkey/tests/unit_tests/common/utils/test_timer.py similarity index 97% rename from monkey/tests/unit_tests/infection_monkey/utils/test_timer.py rename to monkey/tests/unit_tests/common/utils/test_timer.py index b5291cc0e..43a503cdc 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_timer.py +++ b/monkey/tests/unit_tests/common/utils/test_timer.py @@ -2,7 +2,7 @@ import time import pytest -from infection_monkey.utils.timer import Timer +from common.utils import Timer @pytest.fixture diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py b/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py index 5fe9a3881..7731a3a23 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py @@ -3,8 +3,8 @@ from unittest.mock import MagicMock import pytest +from common.utils import Timer from infection_monkey.utils.decorators import request_cache -from infection_monkey.utils.timer import Timer class MockTimer(Timer): From 653bfbd24bf0c062357544b8a4f747f6e08568cf Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 15:25:56 -0400 Subject: [PATCH 05/34] 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}") From 144506c32d586125c5258e76121b9a37fb86c6a1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 15:43:52 -0400 Subject: [PATCH 06/34] Island: Implement AWSService._run_agent_on_managed_instance() --- .../cc/services/aws/aws_command_runner.py | 3 ++ .../cc/services/aws/aws_service.py | 29 +++++++++++++++---- 2 files changed, 27 insertions(+), 5 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 9f5da8c87..6795868d7 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -11,6 +11,7 @@ STATUS_CHECK_SLEEP_TIME = 1 logger = logging.getLogger(__name__) +# TODO: Make sure the return type is compatible with what RemoteRun is expecting. Add typehint. def start_infection_monkey_agent( aws_client: botocore.client.BaseClient, target_instance_id: str, target_os: str, island_ip: str ): @@ -21,6 +22,8 @@ def start_infection_monkey_agent( 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) + # TODO: Return result + def _get_run_agent_command(target_os: str, island_ip: str): if target_os == "linux": diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index 4f52583ab..db0ac5468 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/monkey/monkey_island/cc/services/aws/aws_service.py @@ -6,6 +6,8 @@ import botocore from common.aws.aws_instance import AWSInstance +from .aws_command_runner import start_infection_monkey_agent + INSTANCE_INFORMATION_LIST_KEY = "InstanceInformationList" INSTANCE_ID_KEY = "InstanceId" COMPUTER_NAME_KEY = "ComputerName" @@ -66,12 +68,29 @@ class AWSService: logger.warning("AWS client error while trying to get manage dinstances: {err}") raise err - def run_agent_on_managed_instances(self, instance_ids: Iterable[str]): - for id_ in instance_ids: - self._run_agent_on_managed_instance(id_) + # TODO: Determine the return type + def run_agents_on_managed_instances( + self, instances: Iterable[Mapping[str, str]], island_ip: str + ): + """ + Run an agent on one or more managed AWS instances. + :param instances: An iterable of instances that the agent will be run on + :param island_ip: The IP address of the Island to pass to the new agents + :return: Mapping with 'instance_id' as a key the agent's status as a value + """ - def _run_agent_on_managed_instance(self, instance_id: str): - pass + results = [] + # TODO: Use threadpool or similar to run these in parallel (daemon threads) + for i in instances: + results.append( + self._run_agent_on_managed_instance(i["instance_id"], i["os"], island_ip) + ) + + return results + + def _run_agent_on_managed_instance(self, instance_id: str, os: str, island_ip: str): + ssm_client = boto3.client("ssm", self.island_aws_instance.region) + return start_infection_monkey_agent(ssm_client, instance_id, os, island_ip) def _filter_relevant_instance_info(raw_managed_instances_info: Sequence[Mapping[str, Any]]): From 528ca76c3268d6eff27251fbcb5053c458a0be5e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 15:52:01 -0400 Subject: [PATCH 07/34] Island: Use AWSService in RemoteRun resource --- monkey/monkey_island/cc/resources/remote_run.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index f918c9253..d1c1149b5 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -5,8 +5,7 @@ from botocore.exceptions import ClientError, NoCredentialsError from flask import jsonify, make_response, request from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.services import aws_service -from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService +from monkey_island.cc.services import AWSService CLIENT_ERROR_FORMAT = ( "ClientError, error message: '{}'. Probably, the IAM role that has been associated with the " @@ -19,20 +18,24 @@ NO_CREDS_ERROR_FORMAT = ( class RemoteRun(flask_restful.Resource): + def __init__(self, aws_service: AWSService): + self._aws_service = aws_service + def run_aws_monkeys(self, request_body): instances = request_body.get("instances") island_ip = request_body.get("island_ip") - return RemoteRunAwsService.run_aws_monkeys(instances, island_ip) + + return self._aws_service.run_agents_on_managed_instances(instances, island_ip) @jwt_required def get(self): action = request.args.get("action") if action == "list_aws": - is_aws = aws_service.is_on_aws() + is_aws = self._aws_service.island_is_running_on_aws() resp = {"is_aws": is_aws} if is_aws: try: - resp["instances"] = aws_service.get_instances() + resp["instances"] = self._aws_service.get_managed_instances() except NoCredentialsError as e: resp["error"] = NO_CREDS_ERROR_FORMAT.format(e) return jsonify(resp) From 68f31db03ad03122213faa113588760c79819542 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 16:23:35 -0400 Subject: [PATCH 08/34] UT: Add StubDIContainer --- monkey/common/__init__.py | 2 +- monkey/tests/common/__init__.py | 1 + monkey/tests/common/stub_di_container.py | 20 +++++++++++++++++++ .../cc/resources/test_pba_file_download.py | 4 ++-- .../cc/resources/test_pba_file_upload.py | 4 ++-- 5 files changed, 26 insertions(+), 5 deletions(-) create mode 100644 monkey/tests/common/__init__.py create mode 100644 monkey/tests/common/stub_di_container.py diff --git a/monkey/common/__init__.py b/monkey/common/__init__.py index 5f2cba2f1..01be39cf4 100644 --- a/monkey/common/__init__.py +++ b/monkey/common/__init__.py @@ -1 +1 @@ -from .di_container import DIContainer +from .di_container import DIContainer, UnregisteredTypeError diff --git a/monkey/tests/common/__init__.py b/monkey/tests/common/__init__.py new file mode 100644 index 000000000..b50d9a82f --- /dev/null +++ b/monkey/tests/common/__init__.py @@ -0,0 +1 @@ +from .stub_di_container import StubDIContainer diff --git a/monkey/tests/common/stub_di_container.py b/monkey/tests/common/stub_di_container.py new file mode 100644 index 000000000..bed434145 --- /dev/null +++ b/monkey/tests/common/stub_di_container.py @@ -0,0 +1,20 @@ +from typing import Any, Sequence, Type, TypeVar +from unittest.mock import MagicMock + +from common import DIContainer, UnregisteredTypeError + +T = TypeVar("T") + + +class StubDIContainer(DIContainer): + def resolve(self, type_: Type[T]) -> T: + try: + return super().resolve(type_) + except UnregisteredTypeError: + return MagicMock() + + def resolve_dependencies(self, type_: Type[T]) -> Sequence[Any]: + try: + return super().resolve_dependencies(type_) + except UnregisteredTypeError: + return MagicMock() diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index d0a8ec48f..e05ab3362 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py @@ -2,8 +2,8 @@ import io from typing import BinaryIO import pytest +from tests.common import StubDIContainer -from common import DIContainer from monkey_island.cc.services import FileRetrievalError, IFileStorageService FILE_NAME = "test_file" @@ -32,7 +32,7 @@ class MockFileStorageService(IFileStorageService): @pytest.fixture def flask_client(build_flask_client, tmp_path): - container = DIContainer() + container = StubDIContainer() container.register(IFileStorageService, MockFileStorageService) with build_flask_client(container) as flask_client: diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py index da1eb2c02..9cbaa50d8 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -2,9 +2,9 @@ import io from typing import BinaryIO import pytest +from tests.common import StubDIContainer from tests.utils import raise_ -from common import DIContainer from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE from monkey_island.cc.services import FileRetrievalError, IFileStorageService @@ -65,7 +65,7 @@ def file_storage_service(): @pytest.fixture def flask_client(build_flask_client, file_storage_service): - container = DIContainer() + container = StubDIContainer() container.register_instance(IFileStorageService, file_storage_service) with build_flask_client(container) as flask_client: From fe1e8ccb75ec28100e2eae202064a5d81bb075da Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 16:27:07 -0400 Subject: [PATCH 09/34] Island: Remove disused RemoteRunAwsService --- .../cc/services/remote_run_aws.py | 82 ------------------- 1 file changed, 82 deletions(-) delete mode 100644 monkey/monkey_island/cc/services/remote_run_aws.py diff --git a/monkey/monkey_island/cc/services/remote_run_aws.py b/monkey/monkey_island/cc/services/remote_run_aws.py deleted file mode 100644 index 02c3b79fc..000000000 --- a/monkey/monkey_island/cc/services/remote_run_aws.py +++ /dev/null @@ -1,82 +0,0 @@ -import logging - -from common.cmd.cmd import Cmd -from common.cmd.cmd_runner import CmdRunner -from monkey_island.cc.services.aws.aws_command_runner import AWSCommandRunner - -logger = logging.getLogger(__name__) - - -class RemoteRunAwsService: - @staticmethod - def run_aws_monkeys(instances, island_ip): - """ - Runs monkeys on the given instances - :param instances: List of instances to run on - :param island_ip: IP of island the monkey will communicate with - :return: Dictionary with instance ids as keys, and True/False as values if succeeded or not - """ - return CmdRunner.run_multiple_commands( - instances, - lambda instance: RemoteRunAwsService._run_aws_monkey_cmd_async( - instance["instance_id"], - RemoteRunAwsService._is_linux(instance["os"]), - island_ip, - ), - lambda _, result: result.is_success, - ) - - @staticmethod - def _run_aws_monkey_cmd_async(instance_id, is_linux, island_ip): - """ - Runs a monkey remotely using AWS - :param instance_id: Instance ID of target - :param is_linux: Whether target is linux - :param island_ip: IP of the island which the instance will try to connect to - :return: Cmd - """ - cmd_text = RemoteRunAwsService._get_run_monkey_cmd_line(is_linux, island_ip) - return RemoteRunAwsService._run_aws_cmd_async(instance_id, is_linux, cmd_text) - - @staticmethod - def _run_aws_cmd_async(instance_id, is_linux, cmd_line): - cmd_runner = AWSCommandRunner(is_linux, instance_id) - return Cmd(cmd_runner, cmd_runner.run_command_async(cmd_line)) - - @staticmethod - def _is_linux(os): - return "linux" == os - - @staticmethod - def _get_run_monkey_cmd_linux_line(island_ip): - return ( - r"wget --no-check-certificate https://" - + island_ip - + r":5000/api/agent/download/linux " - + r"-O monkey-linux-64" - + r"; chmod +x monkey-linux-64" - + r"; ./monkey-linux-64" - + r" m0nk3y -s " - + island_ip - + r":5000" - ) - - @staticmethod - def _get_run_monkey_cmd_windows_line(island_ip): - return ( - r"[System.Net.ServicePointManager]::ServerCertificateValidationCallback = {" - r"$true}; (New-Object System.Net.WebClient).DownloadFile('https://" - + island_ip - + r":5000/api/agent/download/windows'" - + r"'.\\monkey.exe'); " - r";Start-Process -FilePath '.\\monkey.exe' " - r"-ArgumentList 'm0nk3y -s " + island_ip + r":5000'; " - ) - - @staticmethod - def _get_run_monkey_cmd_line(is_linux, island_ip): - return ( - RemoteRunAwsService._get_run_monkey_cmd_linux_line(island_ip) - if is_linux - else RemoteRunAwsService._get_run_monkey_cmd_windows_line(island_ip) - ) From 4f3a150f5ad1b1fdcb8561802405516b1825ff70 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 9 May 2022 16:30:24 -0400 Subject: [PATCH 10/34] Common: Remove disused common.cmd.* --- monkey/common/cmd/__init__.py | 0 monkey/common/cmd/aws/__init__.py | 0 monkey/common/cmd/cmd.py | 8 -- monkey/common/cmd/cmd_result.py | 10 -- monkey/common/cmd/cmd_runner.py | 154 ------------------------------ monkey/common/cmd/cmd_status.py | 7 -- 6 files changed, 179 deletions(-) delete mode 100644 monkey/common/cmd/__init__.py delete mode 100644 monkey/common/cmd/aws/__init__.py delete mode 100644 monkey/common/cmd/cmd.py delete mode 100644 monkey/common/cmd/cmd_result.py delete mode 100644 monkey/common/cmd/cmd_runner.py delete mode 100644 monkey/common/cmd/cmd_status.py diff --git a/monkey/common/cmd/__init__.py b/monkey/common/cmd/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/monkey/common/cmd/aws/__init__.py b/monkey/common/cmd/aws/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/monkey/common/cmd/cmd.py b/monkey/common/cmd/cmd.py deleted file mode 100644 index 6daa970a6..000000000 --- a/monkey/common/cmd/cmd.py +++ /dev/null @@ -1,8 +0,0 @@ -class Cmd(object): - """ - Class representing a command - """ - - def __init__(self, cmd_runner, cmd_id): - self.cmd_runner = cmd_runner - self.cmd_id = cmd_id diff --git a/monkey/common/cmd/cmd_result.py b/monkey/common/cmd/cmd_result.py deleted file mode 100644 index 85179a053..000000000 --- a/monkey/common/cmd/cmd_result.py +++ /dev/null @@ -1,10 +0,0 @@ -class CmdResult(object): - """ - Class representing a command result - """ - - def __init__(self, is_success, status_code=None, stdout=None, stderr=None): - self.is_success = is_success - self.status_code = status_code - self.stdout = stdout - self.stderr = stderr diff --git a/monkey/common/cmd/cmd_runner.py b/monkey/common/cmd/cmd_runner.py deleted file mode 100644 index 38d22b0a1..000000000 --- a/monkey/common/cmd/cmd_runner.py +++ /dev/null @@ -1,154 +0,0 @@ -import logging -import time -from abc import abstractmethod - -from common.cmd.cmd_result import CmdResult -from common.cmd.cmd_status import CmdStatus - -logger = logging.getLogger(__name__) - - -class CmdRunner(object): - """ - Interface for running commands on a remote machine - - Since these classes are a bit complex, I provide a list of common terminology and formats: - * command line - a command line. e.g. 'echo hello' - * command - represent a single command which was already run. Always of type Cmd - * command id - any unique identifier of a command which was already run - * command result - represents the result of running a command. Always of type CmdResult - * command status - represents the current status of a command. Always of type CmdStatus - * command info - Any consistent structure representing additional information of a command - which was already run - * instance - a machine that commands will be run on. Can be any dictionary with 'instance_id' - as a field - * instance_id - any unique identifier of an instance (machine). Can be of any format - """ - - # Default command timeout in seconds - DEFAULT_TIMEOUT = 5 - # Time to sleep when waiting on commands. - WAIT_SLEEP_TIME = 1 - - def __init__(self, is_linux): - self.is_linux = is_linux - - @staticmethod - def run_multiple_commands(instances, inst_to_cmd, inst_n_cmd_res_to_res): - """ - Run multiple commands on various instances - :param instances: List of instances. - :param inst_to_cmd: Function which receives an instance, runs a command asynchronously - and returns Cmd - :param inst_n_cmd_res_to_res: Function which receives an instance and CmdResult - and returns a parsed result (of any format) - :return: Dictionary with 'instance_id' as key and parsed result as value - """ - command_instance_dict = {} - - for instance in instances: - command = inst_to_cmd(instance) - command_instance_dict[command] = instance - - instance_results = {} - command_result_pairs = CmdRunner.wait_commands(list(command_instance_dict.keys())) - for command, result in command_result_pairs: - instance = command_instance_dict[command] - instance_results[instance["instance_id"]] = inst_n_cmd_res_to_res(instance, result) - - return instance_results - - @abstractmethod - def run_command_async(self, command_line): - """ - Runs the given command on the remote machine asynchronously. - :param command_line: The command line to run - :return: Command ID (in any format) - """ - raise NotImplementedError() - - @staticmethod - def wait_commands(commands, timeout=DEFAULT_TIMEOUT): - """ - Waits on all commands up to given timeout - :param commands: list of commands (of type Cmd) - :param timeout: Timeout in seconds for command. - :return: commands and their results (tuple of Command and CmdResult) - """ - init_time = time.time() - curr_time = init_time - - results = [] - # TODO: Use timer.Timer - while (curr_time - init_time < timeout) and (len(commands) != 0): - for command in list( - commands - ): # list(commands) clones the list. We do so because we remove items inside - CmdRunner._process_command(command, commands, results, True) - - time.sleep(CmdRunner.WAIT_SLEEP_TIME) - curr_time = time.time() - - for command in list(commands): - CmdRunner._process_command(command, commands, results, False) - - for command, result in results: - if not result.is_success: - logger.error( - f"The command with id: {str(command.cmd_id)} failed. " - f"Status code: {str(result.status_code)}" - ) - - return results - - @abstractmethod - def query_command(self, command_id): - """ - Queries the already run command for more info - :param command_id: The command ID to query - :return: Command info (in any format) - """ - raise NotImplementedError() - - @abstractmethod - def get_command_result(self, command_info): - """ - Gets the result of the already run command - :param command_info: The command info of the command to get the result of - :return: CmdResult - """ - raise NotImplementedError() - - @abstractmethod - def get_command_status(self, command_info): - """ - Gets the status of the already run command - :param command_info: The command info of the command to get the result of - :return: CmdStatus - """ - raise NotImplementedError() - - @staticmethod - def _process_command(command, commands, results, should_process_only_finished): - """ - Removes the command from the list, processes its result and appends to results - :param command: Command to process. Must be in commands. - :param commands: List of unprocessed commands. - :param results: List of command results. - :param should_process_only_finished: If True, processes only if command finished. - :return: None - """ - c_runner = command.cmd_runner - c_id = command.cmd_id - try: - command_info = c_runner.query_command(c_id) - if (not should_process_only_finished) or c_runner.get_command_status( - command_info - ) != CmdStatus.IN_PROGRESS: - commands.remove(command) - results.append((command, c_runner.get_command_result(command_info))) - except Exception: - logger.exception("Exception while querying command: `%s`", str(c_id)) - if not should_process_only_finished: - commands.remove(command) - results.append((command, CmdResult(False))) diff --git a/monkey/common/cmd/cmd_status.py b/monkey/common/cmd/cmd_status.py deleted file mode 100644 index 6a9bbae71..000000000 --- a/monkey/common/cmd/cmd_status.py +++ /dev/null @@ -1,7 +0,0 @@ -from enum import Enum - - -class CmdStatus(Enum): - IN_PROGRESS = 0 - SUCCESS = 1 - FAILURE = 2 From c655aaffe9c4607dc2803fbb7392abc8d958e8c4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 08:04:33 -0400 Subject: [PATCH 11/34] UT: Move test_aws_service.py to services/aws/ --- .../monkey_island/cc/services/{ => aws}/test_aws_service.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename monkey/tests/unit_tests/monkey_island/cc/services/{ => aws}/test_aws_service.py (100%) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_service.py similarity index 100% rename from monkey/tests/unit_tests/monkey_island/cc/services/test_aws_service.py rename to monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_service.py From bfeea19c88155700db46a747f2cb6bf6a4921b8b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 08:47:40 -0400 Subject: [PATCH 12/34] Island: Fix broken AWS queries in aws_command_runner --- monkey/monkey_island/cc/services/aws/aws_command_runner.py | 6 +++--- 1 file changed, 3 insertions(+), 3 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 6795868d7..88eb7e3b4 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -69,13 +69,13 @@ def _run_command_async( 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( + command_response = aws_client.send_command( DocumentName=doc_name, Parameters={"commands": [command]}, InstanceIds=[target_instance_id], ) - command_id = command_response["CommandId"] + command_id = command_response["Command"]["CommandId"] logger.debug( f"Started command on AWS instance {target_instance_id} with command ID {command_id}" ) @@ -93,7 +93,7 @@ def _wait_for_command_to_complete( 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( + command_status = aws_client.get_command_invocation( CommandId=command_id, InstanceId=target_instance_id )["Status"] logger.debug(f"Command {command_id} status: {command_status}") From 79eb584c5d3943dfd0dd12d546fe5994bfc445a7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 08:48:51 -0400 Subject: [PATCH 13/34] Island: Skip try_add_aws_exporter_to_manager() This is causing the island to crash at the moment and will be completely removed later in #1928. --- monkey/monkey_island/cc/services/reporting/exporter_init.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/reporting/exporter_init.py b/monkey/monkey_island/cc/services/reporting/exporter_init.py index bb2d568b9..5eb640ee7 100644 --- a/monkey/monkey_island/cc/services/reporting/exporter_init.py +++ b/monkey/monkey_island/cc/services/reporting/exporter_init.py @@ -9,7 +9,7 @@ logger = logging.getLogger(__name__) def populate_exporter_list(): manager = ReportExporterManager() - try_add_aws_exporter_to_manager(manager) + # try_add_aws_exporter_to_manager(manager) if len(manager.get_exporters_list()) != 0: logger.debug( From 27f8195be5d87585e7091272be1573dc15c1aa60 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 09:47:46 -0400 Subject: [PATCH 14/34] UT: Add unit tests for aws_command_runner --- .../cc/services/aws/aws_command_runner.py | 4 +- .../services/aws/test_aws_command_runner.py | 220 ++++++++++++++++++ 2 files changed, 223 insertions(+), 1 deletion(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py 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 88eb7e3b4..7983cf28a 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -7,6 +7,8 @@ from common.utils import Timer REMOTE_COMMAND_TIMEOUT = 5 STATUS_CHECK_SLEEP_TIME = 1 +LINUX_DOCUMENT_NAME = "AWS-RunShellScript" +WINDOWS_DOCUMENT_NAME = "AWS-RunPowerShellScript" logger = logging.getLogger(__name__) @@ -66,7 +68,7 @@ def _get_run_monkey_cmd_windows_line(island_ip): 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" + doc_name = LINUX_DOCUMENT_NAME if target_os == "linux" else WINDOWS_DOCUMENT_NAME logger.debug(f'Running command on {target_instance_id} -- {doc_name}: "{command}"') command_response = aws_client.send_command( diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py new file mode 100644 index 000000000..c86fe06d8 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -0,0 +1,220 @@ +from unittest.mock import MagicMock + +import pytest + +from monkey_island.cc.services.aws.aws_command_runner import ( + LINUX_DOCUMENT_NAME, + WINDOWS_DOCUMENT_NAME, + start_infection_monkey_agent, +) + +INSTANCE_ID = "BEEFFACE" +ISLAND_IP = "127.0.0.1" +""" + "commands": [ + "wget --no-check-certificate " + "https://172.31.32.78:5000/api/agent/download/linux " + "-O monkey-linux-64; chmod +x " + "monkey-linux-64; ./monkey-linux-64 " + "m0nk3y -s 172.31.32.78:5000" + ] + """ + + +@pytest.fixture +def send_command_response(): + return { + "Command": { + "CloudWatchOutputConfig": { + "CloudWatchLogGroupName": "", + "CloudWatchOutputEnabled": False, + }, + "CommandId": "fe3cf24f-71b7-42b9-93ca-e27c34dd0581", + "CompletedCount": 0, + "DocumentName": "AWS-RunShellScript", + "DocumentVersion": "$DEFAULT", + "InstanceIds": ["i-0b62d6f0b0d9d7e77"], + "OutputS3Region": "eu-central-1", + "Parameters": {"commands": []}, + "Status": "Pending", + "StatusDetails": "Pending", + "TargetCount": 1, + "Targets": [], + "TimeoutSeconds": 3600, + }, + "ResponseMetadata": { + "HTTPHeaders": { + "connection": "keep-alive", + "content-length": "973", + "content-type": "application/x-amz-json-1.1", + "date": "Tue, 10 May 2022 12:35:49 GMT", + "server": "Server", + "x-amzn-requestid": "110b1563-aaf0-4e09-bd23-2db465822be7", + }, + "HTTPStatusCode": 200, + "RequestId": "110b1563-aaf0-4e09-bd23-2db465822be7", + "RetryAttempts": 0, + }, + } + + +@pytest.fixture +def in_progress_response(): + return { + "CloudWatchOutputConfig": {"CloudWatchLogGroupName": "", "CloudWatchOutputEnabled": False}, + "CommandId": "a5332cc6-0f9f-48e6-826a-d4bd7cabc2ee", + "Comment": "", + "DocumentName": "AWS-RunShellScript", + "DocumentVersion": "$DEFAULT", + "ExecutionEndDateTime": "", + "InstanceId": "i-0b62d6f0b0d9d7e77", + "PluginName": "aws:runShellScript", + "ResponseCode": -1, + "StandardErrorContent": "", + "StandardErrorUrl": "", + "StandardOutputContent": "", + "StandardOutputUrl": "", + "Status": "InProgress", + "StatusDetails": "InProgress", + } + + +@pytest.fixture +def success_response(): + return { + "CloudWatchOutputConfig": {"CloudWatchLogGroupName": "", "CloudWatchOutputEnabled": False}, + "CommandId": "a5332cc6-0f9f-48e6-826a-d4bd7cabc2ee", + "Comment": "", + "DocumentName": "AWS-RunShellScript", + "DocumentVersion": "$DEFAULT", + "ExecutionEndDateTime": "", + "InstanceId": "i-0b62d6f0b0d9d7e77", + "PluginName": "aws:runShellScript", + "ResponseCode": -1, + "StandardErrorContent": "", + "StandardErrorUrl": "", + "StandardOutputContent": "", + "StandardOutputUrl": "", + "Status": "Success", + "StatusDetails": "Success", + } + + +@pytest.fixture +def error_response(): + return { + "CloudWatchOutputConfig": {"CloudWatchLogGroupName": "", "CloudWatchOutputEnabled": False}, + "CommandId": "a5332cc6-0f9f-48e6-826a-d4bd7cabc2ee", + "Comment": "", + "DocumentName": "AWS-RunShellScript", + "DocumentVersion": "$DEFAULT", + "ExecutionEndDateTime": "", + "InstanceId": "i-0b62d6f0b0d9d7e77", + "PluginName": "aws:runShellScript", + "ResponseCode": -1, + "StandardErrorContent": "ERROR RUNNING COMMAND", + "StandardErrorUrl": "", + "StandardOutputContent": "", + "StandardOutputUrl": "", + # NOTE: "Error" is technically not a valid value for this field, but we want to test that + # anything other than "Success" and "InProgress" is treated as an error. This is + # simpler than testing all of the different possible error cases. + "Status": "Error", + "StatusDetails": "Error", + } + + +@pytest.fixture(autouse=True) +def patch_timeouts(monkeypatch): + monkeypatch.setattr( + "monkey_island.cc.services.aws.aws_command_runner.REMOTE_COMMAND_TIMEOUT", 0.03 + ) + monkeypatch.setattr( + "monkey_island.cc.services.aws.aws_command_runner.STATUS_CHECK_SLEEP_TIME", 0.01 + ) + + +@pytest.fixture +def successful_mock_client(send_command_response, success_response): + aws_client = MagicMock() + aws_client.send_command = MagicMock(return_value=send_command_response) + aws_client.get_command_invocation = MagicMock(return_value=success_response) + + return aws_client + + +def test_correct_instance_id(successful_mock_client, send_command_response, success_response): + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) + + successful_mock_client.send_command.assert_called_once() + assert successful_mock_client.send_command.call_args.kwargs["InstanceIds"] == [INSTANCE_ID] + + +def test_linux_doc_name(successful_mock_client, send_command_response, success_response): + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) + + successful_mock_client.send_command.assert_called_once() + assert ( + successful_mock_client.send_command.call_args.kwargs["DocumentName"] == LINUX_DOCUMENT_NAME + ) + + +def test_windows_doc_name(successful_mock_client, send_command_response, success_response): + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) + + successful_mock_client.send_command.assert_called_once() + assert ( + successful_mock_client.send_command.call_args.kwargs["DocumentName"] + == WINDOWS_DOCUMENT_NAME + ) + + +def test_linux_command(successful_mock_client, send_command_response, success_response): + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) + + successful_mock_client.send_command.assert_called_once() + assert ( + "wget" in successful_mock_client.send_command.call_args.kwargs["Parameters"]["commands"][0] + ) + + +def test_windows_command(successful_mock_client, send_command_response, success_response): + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) + + successful_mock_client.send_command.assert_called_once() + assert ( + "DownloadFile" + in successful_mock_client.send_command.call_args.kwargs["Parameters"]["commands"][0] + ) + + +def test_in_progress_no_timeout(send_command_response, in_progress_response, success_response): + aws_client = MagicMock() + aws_client.send_command = MagicMock(return_value=send_command_response) + aws_client.get_command_invocation = MagicMock( + side_effect=[in_progress_response, in_progress_response, success_response] + ) + + # If this test fails, an exception will be raised + start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + + +# TODO: Address this test case +""" +def test_in_progress_timeout(send_command_response, in_progress_response): + aws_client = MagicMock() + aws_client.send_command = MagicMock(return_value=send_command_response) + aws_client.get_command_invocation = MagicMock(return_value=in_progress_response) + + with pytest.raises(Exception): + start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) +""" + + +def test_failed_command(send_command_response, error_response): + aws_client = MagicMock() + aws_client.send_command = MagicMock(return_value=send_command_response) + aws_client.get_command_invocation = MagicMock(return_value=error_response) + + with pytest.raises(Exception): + start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) From 109ea87196317ed099cfca747eb50ffdb5b6fca2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 09:47:58 -0400 Subject: [PATCH 15/34] Island: Fix negative sleep time bug in aws_command_runner --- monkey/monkey_island/cc/services/aws/aws_command_runner.py | 3 +-- 1 file changed, 1 insertion(+), 2 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 7983cf28a..78dfdc21c 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -92,8 +92,7 @@ def _wait_for_command_to_complete( 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) + time.sleep(STATUS_CHECK_SLEEP_TIME) command_status = aws_client.get_command_invocation( CommandId=command_id, InstanceId=target_instance_id From e5285f2f78c28baa1881a81d300c1fba4242adf8 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Tue, 10 May 2022 16:57:36 +0300 Subject: [PATCH 16/34] Island: Add custom error and sketch out AWS command results --- .../cc/services/aws/aws_command_runner.py | 27 ++++++++++++++++--- 1 file changed, 24 insertions(+), 3 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 78dfdc21c..b6e9aa16e 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -85,6 +85,10 @@ def _run_command_async( return command_id +class AWSCommandError(Exception): + pass + + def _wait_for_command_to_complete( aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str ): @@ -94,9 +98,11 @@ def _wait_for_command_to_complete( while not timer.is_expired(): time.sleep(STATUS_CHECK_SLEEP_TIME) - command_status = aws_client.get_command_invocation( + command_result = aws_client.get_command_invocation( CommandId=command_id, InstanceId=target_instance_id - )["Status"] + ) + command_status = command_result["Status"] + logger.debug(f"Command {command_id} status: {command_status}") if command_status == "Success": @@ -104,4 +110,19 @@ def _wait_for_command_to_complete( if command_status != "InProgress": # TODO: Create an exception for this occasion and raise it with useful information. - raise Exception("COMMAND FAILED") + raise AWSCommandError( + f"AWS command failed." f" Command invocation contents: {command_result}" + ) + + +def _fetch_command_results( + aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str +): + command_results = aws_client.ssm.get_command_invocation( + CommandId=command_id, InstanceId=target_instance_id + ) + # TODO: put these into a dataclass and return + # self.is_successful(command_info, True) + # command_results["ResponseCode"] + # command_results["StandardOutputContent"] + # command_results["StandardErrorContent"] From 2804ba9b07a455dad3d29ce9ad781465186fc5ca Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 13:00:39 -0400 Subject: [PATCH 17/34] Island: Return AWSCommandResults from start_infection_monkey_agent() --- .../monkey_island/cc/services/aws/__init__.py | 1 + .../cc/services/aws/aws_command_runner.py | 76 ++++++++++++------- .../cc/services/aws/aws_service.py | 11 +-- .../services/aws/test_aws_command_runner.py | 42 ++++++---- vulture_allowlist.py | 2 + 5 files changed, 83 insertions(+), 49 deletions(-) diff --git a/monkey/monkey_island/cc/services/aws/__init__.py b/monkey/monkey_island/cc/services/aws/__init__.py index ff20c89a2..78ed07bc5 100644 --- a/monkey/monkey_island/cc/services/aws/__init__.py +++ b/monkey/monkey_island/cc/services/aws/__init__.py @@ -1 +1,2 @@ from .aws_service import AWSService +from .aws_command_runner import AWSCommandResults, AWSCommandStatus 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 b6e9aa16e..69ada244c 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -1,5 +1,7 @@ import logging import time +from dataclasses import dataclass +from enum import Enum, auto import botocore @@ -13,18 +15,34 @@ WINDOWS_DOCUMENT_NAME = "AWS-RunPowerShellScript" logger = logging.getLogger(__name__) -# TODO: Make sure the return type is compatible with what RemoteRun is expecting. Add typehint. +class AWSCommandStatus(Enum): + SUCCESS = auto() + IN_PROGRESS = auto() + ERROR = auto() + + +@dataclass(frozen=True) +class AWSCommandResults: + response_code: int + stdout: str + stderr: str + status: AWSCommandStatus + + @property + def success(self): + return self.status == AWSCommandStatus.SUCCESS + + def start_infection_monkey_agent( aws_client: botocore.client.BaseClient, target_instance_id: str, target_os: str, island_ip: str -): +) -> AWSCommandResults: """ 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) - # TODO: Return result + return _wait_for_command_to_complete(aws_client, target_instance_id, command_id) def _get_run_agent_command(target_os: str, island_ip: str): @@ -85,44 +103,44 @@ def _run_command_async( return command_id -class AWSCommandError(Exception): - pass - - def _wait_for_command_to_complete( aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str -): +) -> AWSCommandResults: timer = Timer() timer.set(REMOTE_COMMAND_TIMEOUT) while not timer.is_expired(): time.sleep(STATUS_CHECK_SLEEP_TIME) - command_result = aws_client.get_command_invocation( - CommandId=command_id, InstanceId=target_instance_id - ) - command_status = command_result["Status"] + command_results = _fetch_command_results(aws_client, target_instance_id, command_id) + logger.debug(f"Command {command_id} status: {command_results.status.name}") - logger.debug(f"Command {command_id} status: {command_status}") + if command_results.status != AWSCommandStatus.IN_PROGRESS: + return command_results - if command_status == "Success": - break - - if command_status != "InProgress": - # TODO: Create an exception for this occasion and raise it with useful information. - raise AWSCommandError( - f"AWS command failed." f" Command invocation contents: {command_result}" - ) + return command_results def _fetch_command_results( aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str -): - command_results = aws_client.ssm.get_command_invocation( +) -> AWSCommandResults: + command_results = aws_client.get_command_invocation( CommandId=command_id, InstanceId=target_instance_id ) - # TODO: put these into a dataclass and return - # self.is_successful(command_info, True) - # command_results["ResponseCode"] - # command_results["StandardOutputContent"] - # command_results["StandardErrorContent"] + command_status = command_results["Status"] + logger.debug(f"Command {command_id} status: {command_status}") + + aws_command_result_status = None + if command_status == "Success": + aws_command_result_status = AWSCommandStatus.SUCCESS + elif command_status == "InProgress": + aws_command_result_status = AWSCommandStatus.IN_PROGRESS + else: + aws_command_result_status = AWSCommandStatus.ERROR + + return AWSCommandResults( + command_results["ResponseCode"], + command_results["StandardOutputContent"], + command_results["StandardErrorContent"], + aws_command_result_status, + ) diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index db0ac5468..cf4acd89d 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/monkey/monkey_island/cc/services/aws/aws_service.py @@ -6,7 +6,7 @@ import botocore from common.aws.aws_instance import AWSInstance -from .aws_command_runner import start_infection_monkey_agent +from .aws_command_runner import AWSCommandResults, start_infection_monkey_agent INSTANCE_INFORMATION_LIST_KEY = "InstanceInformationList" INSTANCE_ID_KEY = "InstanceId" @@ -68,15 +68,14 @@ class AWSService: logger.warning("AWS client error while trying to get manage dinstances: {err}") raise err - # TODO: Determine the return type def run_agents_on_managed_instances( self, instances: Iterable[Mapping[str, str]], island_ip: str - ): + ) -> Sequence[AWSCommandResults]: """ Run an agent on one or more managed AWS instances. :param instances: An iterable of instances that the agent will be run on :param island_ip: The IP address of the Island to pass to the new agents - :return: Mapping with 'instance_id' as a key the agent's status as a value + :return: A sequence of AWSCommandResults """ results = [] @@ -88,7 +87,9 @@ class AWSService: return results - def _run_agent_on_managed_instance(self, instance_id: str, os: str, island_ip: str): + def _run_agent_on_managed_instance( + self, instance_id: str, os: str, island_ip: str + ) -> AWSCommandResults: ssm_client = boto3.client("ssm", self.island_aws_instance.region) return start_infection_monkey_agent(ssm_client, instance_id, os, island_ip) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py index c86fe06d8..3386772ed 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -5,6 +5,8 @@ import pytest from monkey_island.cc.services.aws.aws_command_runner import ( LINUX_DOCUMENT_NAME, WINDOWS_DOCUMENT_NAME, + AWSCommandResults, + AWSCommandStatus, start_infection_monkey_agent, ) @@ -143,14 +145,14 @@ def successful_mock_client(send_command_response, success_response): return aws_client -def test_correct_instance_id(successful_mock_client, send_command_response, success_response): +def test_correct_instance_id(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) successful_mock_client.send_command.assert_called_once() assert successful_mock_client.send_command.call_args.kwargs["InstanceIds"] == [INSTANCE_ID] -def test_linux_doc_name(successful_mock_client, send_command_response, success_response): +def test_linux_doc_name(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) successful_mock_client.send_command.assert_called_once() @@ -159,7 +161,7 @@ def test_linux_doc_name(successful_mock_client, send_command_response, success_r ) -def test_windows_doc_name(successful_mock_client, send_command_response, success_response): +def test_windows_doc_name(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) successful_mock_client.send_command.assert_called_once() @@ -169,7 +171,7 @@ def test_windows_doc_name(successful_mock_client, send_command_response, success ) -def test_linux_command(successful_mock_client, send_command_response, success_response): +def test_linux_command(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) successful_mock_client.send_command.assert_called_once() @@ -178,7 +180,7 @@ def test_linux_command(successful_mock_client, send_command_response, success_re ) -def test_windows_command(successful_mock_client, send_command_response, success_response): +def test_windows_command(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) successful_mock_client.send_command.assert_called_once() @@ -188,27 +190,24 @@ def test_windows_command(successful_mock_client, send_command_response, success_ ) -def test_in_progress_no_timeout(send_command_response, in_progress_response, success_response): +def test_multiple_status_queries(send_command_response, in_progress_response, success_response): aws_client = MagicMock() aws_client.send_command = MagicMock(return_value=send_command_response) aws_client.get_command_invocation = MagicMock( side_effect=[in_progress_response, in_progress_response, success_response] ) - # If this test fails, an exception will be raised - start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + assert command_results.status == AWSCommandStatus.SUCCESS -# TODO: Address this test case -""" def test_in_progress_timeout(send_command_response, in_progress_response): aws_client = MagicMock() aws_client.send_command = MagicMock(return_value=send_command_response) aws_client.get_command_invocation = MagicMock(return_value=in_progress_response) - with pytest.raises(Exception): - start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) -""" + command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + assert command_results.status == AWSCommandStatus.IN_PROGRESS def test_failed_command(send_command_response, error_response): @@ -216,5 +215,18 @@ def test_failed_command(send_command_response, error_response): aws_client.send_command = MagicMock(return_value=send_command_response) aws_client.get_command_invocation = MagicMock(return_value=error_response) - with pytest.raises(Exception): - start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + assert command_results.status == AWSCommandStatus.ERROR + + +@pytest.mark.parametrize( + "status, success", + [ + (AWSCommandStatus.SUCCESS, True), + (AWSCommandStatus.IN_PROGRESS, False), + (AWSCommandStatus.ERROR, False), + ], +) +def test_command_resuls_status(status, success): + results = AWSCommandResults(0, "", "", status) + assert results.success == success diff --git a/vulture_allowlist.py b/vulture_allowlist.py index c8c1378d4..cec60a0c9 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -167,3 +167,5 @@ _.instance_name # unused attribute (monkey/common/cloud/azure/azure_instance.py _.instance_name # unused attribute (monkey/common/cloud/azure/azure_instance.py:64) GCPHandler # unused function (envs/monkey_zoo/blackbox/test_blackbox.py:57) architecture # unused variable (monkey/infection_monkey/exploit/caching_agent_repository.py:25) + +response_code # unused variable (monkey/monkey_island/cc/services/aws/aws_command_runner.py:26) From 487d1c55afb923670daa4505f221d31b6c3df3fa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 13:23:50 -0400 Subject: [PATCH 18/34] Common: Add queue_to_list() --- monkey/common/utils/code_utils.py | 15 +++++++++++++ .../common/utils/test_code_utils.py | 22 +++++++++++++++++++ 2 files changed, 37 insertions(+) create mode 100644 monkey/tests/unit_tests/common/utils/test_code_utils.py diff --git a/monkey/common/utils/code_utils.py b/monkey/common/utils/code_utils.py index 810fbb0d7..251ce9375 100644 --- a/monkey/common/utils/code_utils.py +++ b/monkey/common/utils/code_utils.py @@ -1,3 +1,7 @@ +import queue +from typing import Any, List + + class abstractstatic(staticmethod): __slots__ = () @@ -15,3 +19,14 @@ class Singleton(type): if cls not in cls._instances: cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs) return cls._instances[cls] + + +def queue_to_list(q: queue.Queue) -> List[Any]: + list_ = [] + try: + while True: + list_.append(q.get_nowait()) + except queue.Empty: + pass + + return list_ diff --git a/monkey/tests/unit_tests/common/utils/test_code_utils.py b/monkey/tests/unit_tests/common/utils/test_code_utils.py new file mode 100644 index 000000000..411b07a63 --- /dev/null +++ b/monkey/tests/unit_tests/common/utils/test_code_utils.py @@ -0,0 +1,22 @@ +from queue import Queue + +from common.utils.code_utils import queue_to_list + + +def test_empty_queue_to_empty_list(): + q = Queue() + + list_ = queue_to_list(q) + + assert len(list_) == 0 + + +def test_queue_to_list(): + expected_list = [8, 6, 7, 5, 3, 0, 9] + q = Queue() + for i in expected_list: + q.put(i) + + list_ = queue_to_list(q) + + assert list_ == expected_list From 60229f4a6594d1ca3b623004510f096a948a7de4 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 13:27:06 -0400 Subject: [PATCH 19/34] Island: Run multiple AWS commands concurrently --- .../monkey_island/cc/resources/remote_run.py | 4 ++- .../cc/services/aws/aws_service.py | 27 +++++++++++++------ 2 files changed, 22 insertions(+), 9 deletions(-) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index d1c1149b5..f872fc07b 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -1,4 +1,5 @@ import json +from typing import Sequence import flask_restful from botocore.exceptions import ClientError, NoCredentialsError @@ -6,6 +7,7 @@ from flask import jsonify, make_response, request from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services import AWSService +from monkey_island.cc.services.aws import AWSCommandResults CLIENT_ERROR_FORMAT = ( "ClientError, error message: '{}'. Probably, the IAM role that has been associated with the " @@ -21,7 +23,7 @@ class RemoteRun(flask_restful.Resource): def __init__(self, aws_service: AWSService): self._aws_service = aws_service - def run_aws_monkeys(self, request_body): + def run_aws_monkeys(self, request_body) -> Sequence[AWSCommandResults]: instances = request_body.get("instances") island_ip = request_body.get("island_ip") diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index cf4acd89d..8518f864d 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/monkey/monkey_island/cc/services/aws/aws_service.py @@ -1,10 +1,13 @@ import logging +from queue import Queue +from threading import Thread from typing import Any, Iterable, Mapping, Sequence import boto3 import botocore from common.aws.aws_instance import AWSInstance +from common.utils.code_utils import queue_to_list from .aws_command_runner import AWSCommandResults, start_infection_monkey_agent @@ -78,20 +81,28 @@ class AWSService: :return: A sequence of AWSCommandResults """ - results = [] - # TODO: Use threadpool or similar to run these in parallel (daemon threads) + results_queue = Queue() + command_threads = [] for i in instances: - results.append( - self._run_agent_on_managed_instance(i["instance_id"], i["os"], island_ip) + command_threads.append( + Thread( + target=self._run_agent_on_managed_instance, + args=(results_queue, i["instance_id"], i["os"], island_ip), + daemon=True, + ) ) - return results + for thread in command_threads: + thread.join() + + return queue_to_list(results_queue) def _run_agent_on_managed_instance( - self, instance_id: str, os: str, island_ip: str - ) -> AWSCommandResults: + self, results_queue: Queue, instance_id: str, os: str, island_ip: str + ): ssm_client = boto3.client("ssm", self.island_aws_instance.region) - return start_infection_monkey_agent(ssm_client, instance_id, os, island_ip) + command_results = start_infection_monkey_agent(ssm_client, instance_id, os, island_ip) + results_queue.put(command_results) def _filter_relevant_instance_info(raw_managed_instances_info: Sequence[Mapping[str, Any]]): From 94fd2a26d91535f9462ac4c197eb0e3176dbaf75 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 10 May 2022 13:37:44 -0400 Subject: [PATCH 20/34] Island: Remove disused AwsCmdResult This was replaced by monkey_island.services.aws.AWSCommandResults. --- .../cc/server_utils/aws_cmd_result.py | 29 ------------------- 1 file changed, 29 deletions(-) delete mode 100644 monkey/monkey_island/cc/server_utils/aws_cmd_result.py diff --git a/monkey/monkey_island/cc/server_utils/aws_cmd_result.py b/monkey/monkey_island/cc/server_utils/aws_cmd_result.py deleted file mode 100644 index 0a6c5f3cc..000000000 --- a/monkey/monkey_island/cc/server_utils/aws_cmd_result.py +++ /dev/null @@ -1,29 +0,0 @@ -from common.cmd.cmd_result import CmdResult - - -class AwsCmdResult(CmdResult): - """ - Class representing an AWS command result - """ - - def __init__(self, command_info): - super(AwsCmdResult, self).__init__( - self.is_successful(command_info, True), - command_info["ResponseCode"], - command_info["StandardOutputContent"], - command_info["StandardErrorContent"], - ) - self.command_info = command_info - - @staticmethod - def is_successful(command_info, is_timeout=False): - """ - Determines whether the command was successful. If it timed out and was still in progress, - we assume it worked. - :param command_info: Command info struct (returned by ssm.get_command_invocation) - :param is_timeout: Whether the given command timed out - :return: True if successful, False otherwise. - """ - return (command_info["Status"] == "Success") or ( - is_timeout and (command_info["Status"] == "InProgress") - ) From aa159a424044c2bb893a541fe1bfc251ffa340ad Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 11 May 2022 10:20:27 +0300 Subject: [PATCH 21/34] Island: Improve aws_command_runner.py readability a bit --- .../cc/services/aws/aws_command_runner.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 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 69ada244c..6e88c0e21 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -42,7 +42,8 @@ def start_infection_monkey_agent( command = _get_run_agent_command(target_os, island_ip) command_id = _run_command_async(aws_client, target_instance_id, target_os, command) - return _wait_for_command_to_complete(aws_client, target_instance_id, command_id) + _wait_for_command_to_complete(aws_client, target_instance_id, command_id) + return _fetch_command_results(aws_client, target_instance_id, command_id) def _get_run_agent_command(target_os: str, island_ip: str): @@ -105,7 +106,7 @@ def _run_command_async( def _wait_for_command_to_complete( aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str -) -> AWSCommandResults: +): timer = Timer() timer.set(REMOTE_COMMAND_TIMEOUT) @@ -116,9 +117,7 @@ def _wait_for_command_to_complete( logger.debug(f"Command {command_id} status: {command_results.status.name}") if command_results.status != AWSCommandStatus.IN_PROGRESS: - return command_results - - return command_results + return def _fetch_command_results( @@ -130,7 +129,6 @@ def _fetch_command_results( command_status = command_results["Status"] logger.debug(f"Command {command_id} status: {command_status}") - aws_command_result_status = None if command_status == "Success": aws_command_result_status = AWSCommandStatus.SUCCESS elif command_status == "InProgress": From 30fb57c37fb434f30672a7f07adc2c3183c42c0f Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 11 May 2022 11:10:45 +0300 Subject: [PATCH 22/34] Island: Fix a bug where agent run threads won't start --- monkey/monkey_island/cc/services/aws/aws_service.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index 8518f864d..39ac8b49b 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/monkey/monkey_island/cc/services/aws/aws_service.py @@ -92,6 +92,8 @@ class AWSService: ) ) + [thread.start() for thread in command_threads] + for thread in command_threads: thread.join() From 0b5a507f386535b5371bd45c2a6ee2f29cd20aea Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 11 May 2022 11:46:13 +0300 Subject: [PATCH 23/34] Island: Make AWSCommandStatus enum json serializable --- monkey/monkey_island/cc/resources/remote_run.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index f872fc07b..8f02ee066 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -54,8 +54,20 @@ class RemoteRun(flask_restful.Resource): resp = {} if body.get("type") == "aws": result = self.run_aws_monkeys(body) + result = self._get_encodable_results(result) resp["result"] = result return jsonify(resp) # default action return make_response({"error": "Invalid action"}, 500) + + @staticmethod + def _get_encodable_results(results: Sequence[AWSCommandResults]) -> str: + results_copy = [] + for result in results: + results_copy.append( + AWSCommandResults( + result.response_code, result.stdout, result.stderr, result.status.name.lower() + ) + ) + return results_copy From b8006a62747c696e9427df953b076da1011f0fae Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 11 May 2022 14:32:39 +0300 Subject: [PATCH 24/34] Island: Add instance ID to AWSCommandResults Instance id tells us where the command was launched --- monkey/monkey_island/cc/resources/remote_run.py | 6 +++++- monkey/monkey_island/cc/services/aws/aws_command_runner.py | 2 ++ .../cc/services/aws/test_aws_command_runner.py | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index 8f02ee066..df4b808e2 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -67,7 +67,11 @@ class RemoteRun(flask_restful.Resource): for result in results: results_copy.append( AWSCommandResults( - result.response_code, result.stdout, result.stderr, result.status.name.lower() + result.instance_id, + result.response_code, + result.stdout, + result.stderr, + result.status.name.lower(), ) ) return results_copy 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 6e88c0e21..9a85ba49a 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -23,6 +23,7 @@ class AWSCommandStatus(Enum): @dataclass(frozen=True) class AWSCommandResults: + instance_id: str response_code: int stdout: str stderr: str @@ -137,6 +138,7 @@ def _fetch_command_results( aws_command_result_status = AWSCommandStatus.ERROR return AWSCommandResults( + target_instance_id, command_results["ResponseCode"], command_results["StandardOutputContent"], command_results["StandardErrorContent"], diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py index 3386772ed..f04932b93 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -228,5 +228,5 @@ def test_failed_command(send_command_response, error_response): ], ) def test_command_resuls_status(status, success): - results = AWSCommandResults(0, "", "", status) + results = AWSCommandResults(INSTANCE_ID, 0, "", "", status) assert results.success == success From 680dbca574cfcaccbd9cbb662b8d7e0033712fb2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 07:56:15 -0400 Subject: [PATCH 25/34] UT: Use only Python3.7 features in test_aws_command_runner --- .../services/aws/test_aws_command_runner.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py index f04932b93..1076e3a4d 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -149,45 +149,40 @@ def test_correct_instance_id(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) successful_mock_client.send_command.assert_called_once() - assert successful_mock_client.send_command.call_args.kwargs["InstanceIds"] == [INSTANCE_ID] + call_args_kwargs = successful_mock_client.send_command.call_args[1] + assert call_args_kwargs["InstanceIds"] == [INSTANCE_ID] def test_linux_doc_name(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) successful_mock_client.send_command.assert_called_once() - assert ( - successful_mock_client.send_command.call_args.kwargs["DocumentName"] == LINUX_DOCUMENT_NAME - ) + call_args_kwargs = successful_mock_client.send_command.call_args[1] + assert call_args_kwargs["DocumentName"] == LINUX_DOCUMENT_NAME def test_windows_doc_name(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) successful_mock_client.send_command.assert_called_once() - assert ( - successful_mock_client.send_command.call_args.kwargs["DocumentName"] - == WINDOWS_DOCUMENT_NAME - ) + call_args_kwargs = successful_mock_client.send_command.call_args[1] + assert call_args_kwargs["DocumentName"] == WINDOWS_DOCUMENT_NAME def test_linux_command(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) successful_mock_client.send_command.assert_called_once() - assert ( - "wget" in successful_mock_client.send_command.call_args.kwargs["Parameters"]["commands"][0] - ) + call_args_kwargs = successful_mock_client.send_command.call_args[1] + assert "wget" in call_args_kwargs["Parameters"]["commands"][0] def test_windows_command(successful_mock_client): start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) successful_mock_client.send_command.assert_called_once() - assert ( - "DownloadFile" - in successful_mock_client.send_command.call_args.kwargs["Parameters"]["commands"][0] - ) + call_args_kwargs = successful_mock_client.send_command.call_args[1] + assert "DownloadFile" in call_args_kwargs["Parameters"]["commands"][0] def test_multiple_status_queries(send_command_response, in_progress_response, success_response): From 825c7b9ecf8d29a05b7de54cb36b2ce91057ff24 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 08:04:25 -0400 Subject: [PATCH 26/34] Island: Refactor logic to start threads in AWSService --- .../monkey_island/cc/services/aws/aws_service.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index 39ac8b49b..5a6151580 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/monkey/monkey_island/cc/services/aws/aws_service.py @@ -84,15 +84,13 @@ class AWSService: results_queue = Queue() command_threads = [] for i in instances: - command_threads.append( - Thread( - target=self._run_agent_on_managed_instance, - args=(results_queue, i["instance_id"], i["os"], island_ip), - daemon=True, - ) + t = Thread( + target=self._run_agent_on_managed_instance, + args=(results_queue, i["instance_id"], i["os"], island_ip), + daemon=True, ) - - [thread.start() for thread in command_threads] + t.start() + command_threads.append(t) for thread in command_threads: thread.join() From d1d960abc20caecd71382a56cd39d41cd4136fe3 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 08:09:56 -0400 Subject: [PATCH 27/34] UT: Fix faulty logic in test_multiple_status_queries The test is checking that some calls return "in progress" and subsequent calls return success. Using itertools.repeat allows all future calls to return success. --- .../monkey_island/cc/services/aws/test_aws_command_runner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py index 1076e3a4d..dd1877a6a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -1,3 +1,4 @@ +from itertools import chain, repeat from unittest.mock import MagicMock import pytest @@ -189,7 +190,7 @@ def test_multiple_status_queries(send_command_response, in_progress_response, su aws_client = MagicMock() aws_client.send_command = MagicMock(return_value=send_command_response) aws_client.get_command_invocation = MagicMock( - side_effect=[in_progress_response, in_progress_response, success_response] + side_effect=chain([in_progress_response, in_progress_response], repeat(success_response)) ) command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) From 3702ecbc8c4327b60f8cd0df155f4b5b97347f9e Mon Sep 17 00:00:00 2001 From: vakarisz Date: Wed, 11 May 2022 15:15:15 +0300 Subject: [PATCH 28/34] UI: Fix AWSInstanceTable.js to show status of run commands --- .../RunMonkeyPage/RunOnAWS/AWSInstanceTable.js | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/ui/src/components/pages/RunMonkeyPage/RunOnAWS/AWSInstanceTable.js b/monkey/monkey_island/cc/ui/src/components/pages/RunMonkeyPage/RunOnAWS/AWSInstanceTable.js index cf792f7b8..f9e90831e 100644 --- a/monkey/monkey_island/cc/ui/src/components/pages/RunMonkeyPage/RunOnAWS/AWSInstanceTable.js +++ b/monkey/monkey_island/cc/ui/src/components/pages/RunMonkeyPage/RunOnAWS/AWSInstanceTable.js @@ -70,10 +70,11 @@ function AWSInstanceTable(props) { let color = 'inherit'; if (r) { let instId = r.original.instance_id; + let runResult = getRunResults(instId); if (isSelected(instId)) { color = '#ffed9f'; - } else if (Object.prototype.hasOwnProperty.call(props.results, instId)) { - color = props.results[instId] ? '#00f01b' : '#f00000' + } else if (runResult) { + color = runResult.status === "error" ? '#f00000' : '#00f01b' } } @@ -82,6 +83,15 @@ function AWSInstanceTable(props) { }; } + function getRunResults(instanceId) { + for(let result of props.results){ + if (result.instance_id === instanceId){ + return result + } + } + return false + } + return (
Date: Wed, 11 May 2022 08:18:40 -0400 Subject: [PATCH 29/34] Island: Add timeout parameter to start_infection_monkey_agent() --- .../cc/services/aws/aws_command_runner.py | 13 ++++++---- .../cc/services/aws/aws_service.py | 15 ++++++++--- .../services/aws/test_aws_command_runner.py | 26 +++++++++++-------- 3 files changed, 34 insertions(+), 20 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 9a85ba49a..e896e2cfb 100644 --- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py +++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py @@ -7,7 +7,6 @@ import botocore from common.utils import Timer -REMOTE_COMMAND_TIMEOUT = 5 STATUS_CHECK_SLEEP_TIME = 1 LINUX_DOCUMENT_NAME = "AWS-RunShellScript" WINDOWS_DOCUMENT_NAME = "AWS-RunPowerShellScript" @@ -35,7 +34,11 @@ class AWSCommandResults: def start_infection_monkey_agent( - aws_client: botocore.client.BaseClient, target_instance_id: str, target_os: str, island_ip: str + aws_client: botocore.client.BaseClient, + target_instance_id: str, + target_os: str, + island_ip: str, + timeout: float, ) -> AWSCommandResults: """ Run a command on a remote AWS instance @@ -43,7 +46,7 @@ def start_infection_monkey_agent( 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) + _wait_for_command_to_complete(aws_client, target_instance_id, command_id, timeout) return _fetch_command_results(aws_client, target_instance_id, command_id) @@ -106,10 +109,10 @@ def _run_command_async( def _wait_for_command_to_complete( - aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str + aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str, timeout: float ): timer = Timer() - timer.set(REMOTE_COMMAND_TIMEOUT) + timer.set(timeout) while not timer.is_expired(): time.sleep(STATUS_CHECK_SLEEP_TIME) diff --git a/monkey/monkey_island/cc/services/aws/aws_service.py b/monkey/monkey_island/cc/services/aws/aws_service.py index 5a6151580..92a0e8bf5 100644 --- a/monkey/monkey_island/cc/services/aws/aws_service.py +++ b/monkey/monkey_island/cc/services/aws/aws_service.py @@ -11,6 +11,7 @@ from common.utils.code_utils import queue_to_list from .aws_command_runner import AWSCommandResults, start_infection_monkey_agent +DEFAULT_REMOTE_COMMAND_TIMEOUT = 5 INSTANCE_INFORMATION_LIST_KEY = "InstanceInformationList" INSTANCE_ID_KEY = "InstanceId" COMPUTER_NAME_KEY = "ComputerName" @@ -72,12 +73,16 @@ class AWSService: raise err def run_agents_on_managed_instances( - self, instances: Iterable[Mapping[str, str]], island_ip: str + self, + instances: Iterable[Mapping[str, str]], + island_ip: str, + timeout: float = DEFAULT_REMOTE_COMMAND_TIMEOUT, ) -> Sequence[AWSCommandResults]: """ Run an agent on one or more managed AWS instances. :param instances: An iterable of instances that the agent will be run on :param island_ip: The IP address of the Island to pass to the new agents + :param timeout: The maximum number of seconds to wait for the agents to start :return: A sequence of AWSCommandResults """ @@ -86,7 +91,7 @@ class AWSService: for i in instances: t = Thread( target=self._run_agent_on_managed_instance, - args=(results_queue, i["instance_id"], i["os"], island_ip), + args=(results_queue, i["instance_id"], i["os"], island_ip, timeout), daemon=True, ) t.start() @@ -98,10 +103,12 @@ class AWSService: return queue_to_list(results_queue) def _run_agent_on_managed_instance( - self, results_queue: Queue, instance_id: str, os: str, island_ip: str + self, results_queue: Queue, instance_id: str, os: str, island_ip: str, timeout: float ): ssm_client = boto3.client("ssm", self.island_aws_instance.region) - command_results = start_infection_monkey_agent(ssm_client, instance_id, os, island_ip) + command_results = start_infection_monkey_agent( + ssm_client, instance_id, os, island_ip, timeout + ) results_queue.put(command_results) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py index dd1877a6a..aa4cfdb4b 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -11,6 +11,7 @@ from monkey_island.cc.services.aws.aws_command_runner import ( start_infection_monkey_agent, ) +TIMEOUT = 0.03 INSTANCE_ID = "BEEFFACE" ISLAND_IP = "127.0.0.1" """ @@ -129,9 +130,6 @@ def error_response(): @pytest.fixture(autouse=True) def patch_timeouts(monkeypatch): - monkeypatch.setattr( - "monkey_island.cc.services.aws.aws_command_runner.REMOTE_COMMAND_TIMEOUT", 0.03 - ) monkeypatch.setattr( "monkey_island.cc.services.aws.aws_command_runner.STATUS_CHECK_SLEEP_TIME", 0.01 ) @@ -147,7 +145,7 @@ def successful_mock_client(send_command_response, success_response): def test_correct_instance_id(successful_mock_client): - start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP, TIMEOUT) successful_mock_client.send_command.assert_called_once() call_args_kwargs = successful_mock_client.send_command.call_args[1] @@ -155,7 +153,7 @@ def test_correct_instance_id(successful_mock_client): def test_linux_doc_name(successful_mock_client): - start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP, TIMEOUT) successful_mock_client.send_command.assert_called_once() call_args_kwargs = successful_mock_client.send_command.call_args[1] @@ -163,7 +161,7 @@ def test_linux_doc_name(successful_mock_client): def test_windows_doc_name(successful_mock_client): - start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP, TIMEOUT) successful_mock_client.send_command.assert_called_once() call_args_kwargs = successful_mock_client.send_command.call_args[1] @@ -171,7 +169,7 @@ def test_windows_doc_name(successful_mock_client): def test_linux_command(successful_mock_client): - start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP) + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "linux", ISLAND_IP, TIMEOUT) successful_mock_client.send_command.assert_called_once() call_args_kwargs = successful_mock_client.send_command.call_args[1] @@ -179,7 +177,7 @@ def test_linux_command(successful_mock_client): def test_windows_command(successful_mock_client): - start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP) + start_infection_monkey_agent(successful_mock_client, INSTANCE_ID, "windows", ISLAND_IP, TIMEOUT) successful_mock_client.send_command.assert_called_once() call_args_kwargs = successful_mock_client.send_command.call_args[1] @@ -193,7 +191,9 @@ def test_multiple_status_queries(send_command_response, in_progress_response, su side_effect=chain([in_progress_response, in_progress_response], repeat(success_response)) ) - command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + command_results = start_infection_monkey_agent( + aws_client, INSTANCE_ID, "windows", ISLAND_IP, TIMEOUT + ) assert command_results.status == AWSCommandStatus.SUCCESS @@ -202,7 +202,9 @@ def test_in_progress_timeout(send_command_response, in_progress_response): aws_client.send_command = MagicMock(return_value=send_command_response) aws_client.get_command_invocation = MagicMock(return_value=in_progress_response) - command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + command_results = start_infection_monkey_agent( + aws_client, INSTANCE_ID, "windows", ISLAND_IP, TIMEOUT + ) assert command_results.status == AWSCommandStatus.IN_PROGRESS @@ -211,7 +213,9 @@ def test_failed_command(send_command_response, error_response): aws_client.send_command = MagicMock(return_value=send_command_response) aws_client.get_command_invocation = MagicMock(return_value=error_response) - command_results = start_infection_monkey_agent(aws_client, INSTANCE_ID, "windows", ISLAND_IP) + command_results = start_infection_monkey_agent( + aws_client, INSTANCE_ID, "windows", ISLAND_IP, TIMEOUT + ) assert command_results.status == AWSCommandStatus.ERROR From c865bf3c157c261574944c23293bea3c50191c93 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 09:49:44 -0400 Subject: [PATCH 30/34] UT: Remove unneeded tmp_path from test_pba_file_download.flask_client() --- .../monkey_island/cc/resources/test_pba_file_download.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index e05ab3362..570d3239c 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py @@ -31,7 +31,7 @@ class MockFileStorageService(IFileStorageService): @pytest.fixture -def flask_client(build_flask_client, tmp_path): +def flask_client(build_flask_client): container = StubDIContainer() container.register(IFileStorageService, MockFileStorageService) From cb29c9ed88f6dc0b9c8fc9a5f6ea3d4946bb8168 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 09:51:49 -0400 Subject: [PATCH 31/34] Island: Add RemoteRun to resources/__init__.py --- monkey/monkey_island/cc/app.py | 2 +- monkey/monkey_island/cc/resources/__init__.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 60cc0a026..07e08ea9a 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -9,6 +9,7 @@ from werkzeug.exceptions import NotFound from common import DIContainer from monkey_island.cc.database import database, mongo +from monkey_island.cc.resources import RemoteRun from monkey_island.cc.resources.agent_controls import StopAgentCheck, StopAllAgents from monkey_island.cc.resources.attack.attack_report import AttackReport from monkey_island.cc.resources.auth.auth import Authenticate, init_jwt @@ -38,7 +39,6 @@ from monkey_island.cc.resources.pba_file_download import PBAFileDownload from monkey_island.cc.resources.pba_file_upload import FileUpload from monkey_island.cc.resources.propagation_credentials import PropagationCredentials from monkey_island.cc.resources.ransomware_report import RansomwareReport -from monkey_island.cc.resources.remote_run import RemoteRun from monkey_island.cc.resources.root import Root from monkey_island.cc.resources.security_report import SecurityReport from monkey_island.cc.resources.telemetry import Telemetry diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index e69de29bb..b5bbb1108 100644 --- a/monkey/monkey_island/cc/resources/__init__.py +++ b/monkey/monkey_island/cc/resources/__init__.py @@ -0,0 +1 @@ +from .remote_run import RemoteRun From 43d8fd87ed908e9f1011dddb4eaec8adfb5c5af8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 09:52:21 -0400 Subject: [PATCH 32/34] UT: Add unit tests for RemoteRun resource --- .../cc/resources/test_remote_run.py | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py new file mode 100644 index 000000000..17064d355 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py @@ -0,0 +1,108 @@ +import json +from unittest.mock import MagicMock + +import pytest +from tests.common import StubDIContainer + +from monkey_island.cc.services import AWSService +from monkey_island.cc.services.aws import AWSCommandResults, AWSCommandStatus + + +@pytest.fixture +def mock_aws_service(): + return MagicMock(spec=AWSService) + + +@pytest.fixture +def flask_client(build_flask_client, mock_aws_service): + container = StubDIContainer() + container.register_instance(AWSService, mock_aws_service) + + with build_flask_client(container) as flask_client: + yield flask_client + + +def test_get_invalid_action(flask_client): + response = flask_client.get("/api/remote-monkey?action=INVALID") + assert response.text.rstrip() == "{}" + + +def test_get_no_action(flask_client): + response = flask_client.get("/api/remote-monkey") + assert response.text.rstrip() == "{}" + + +def test_get_not_aws(flask_client, mock_aws_service): + mock_aws_service.island_is_running_on_aws = MagicMock(return_value=False) + response = flask_client.get("/api/remote-monkey?action=list_aws") + assert response.text.rstrip() == '{"is_aws":false}' + + +def test_get_instances(flask_client, mock_aws_service): + instances = [ + {"instance_id": "1", "name": "name1", "os": "linux", "ip_address": "1.1.1.1"}, + {"instance_id": "2", "name": "name2", "os": "windows", "ip_address": "2.2.2.2"}, + {"instance_id": "3", "name": "name3", "os": "linux", "ip_address": "3.3.3.3"}, + ] + mock_aws_service.island_is_running_on_aws = MagicMock(return_value=True) + mock_aws_service.get_managed_instances = MagicMock(return_value=instances) + + response = flask_client.get("/api/remote-monkey?action=list_aws") + + assert json.loads(response.text)["instances"] == instances + assert json.loads(response.text)["is_aws"] is True + + +# TODO: Test error cases for get() + + +def test_post_no_type(flask_client): + response = flask_client.post("/api/remote-monkey", data="{}") + assert response.status_code == 500 + + +def test_post_invalid_type(flask_client): + response = flask_client.post("/api/remote-monkey", data='{"type": "INVALID"}') + assert response.status_code == 500 + + +def test_post(flask_client, mock_aws_service): + request_body = json.dumps( + { + "type": "aws", + "instances": [ + {"instance_id": "1", "os": "linux"}, + {"instance_id": "2", "os": "linux"}, + {"instance_id": "3", "os": "windows"}, + ], + "island_ip": "127.0.0.1", + } + ) + mock_aws_service.run_agents_on_managed_instances = MagicMock( + return_value=[ + AWSCommandResults("1", 0, "", "", AWSCommandStatus.SUCCESS), + AWSCommandResults("2", 0, "some_output", "", AWSCommandStatus.IN_PROGRESS), + AWSCommandResults("3", -1, "", "some_error", AWSCommandStatus.ERROR), + ] + ) + expected_result = [ + {"instance_id": "1", "response_code": 0, "stdout": "", "stderr": "", "status": "success"}, + { + "instance_id": "2", + "response_code": 0, + "stdout": "some_output", + "stderr": "", + "status": "in_progress", + }, + { + "instance_id": "3", + "response_code": -1, + "stdout": "", + "stderr": "some_error", + "status": "error", + }, + ] + + response = flask_client.post("/api/remote-monkey", data=request_body) + + assert json.loads(response.text)["result"] == expected_result From ae0e8ddb8ef0dffb4626c4a6c79574b173d9bbaa Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 10:00:04 -0400 Subject: [PATCH 33/34] Island: Refactor RemoteRun.post() results encoding --- .../monkey_island/cc/resources/remote_run.py | 35 +++++++++---------- 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index df4b808e2..c1b47d20a 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -51,27 +51,26 @@ class RemoteRun(flask_restful.Resource): @jwt_required def post(self): body = json.loads(request.data) - resp = {} if body.get("type") == "aws": - result = self.run_aws_monkeys(body) - result = self._get_encodable_results(result) - resp["result"] = result - return jsonify(resp) + results = self.run_aws_monkeys(body) + return RemoteRun._encode_results(results) # default action return make_response({"error": "Invalid action"}, 500) @staticmethod - def _get_encodable_results(results: Sequence[AWSCommandResults]) -> str: - results_copy = [] - for result in results: - results_copy.append( - AWSCommandResults( - result.instance_id, - result.response_code, - result.stdout, - result.stderr, - result.status.name.lower(), - ) - ) - return results_copy + def _encode_results(results: Sequence[AWSCommandResults]): + result = list(map(RemoteRun._aws_command_results_to_encodable_dict, results)) + response = {"result": result} + + return jsonify(response) + + @staticmethod + def _aws_command_results_to_encodable_dict(aws_command_results: AWSCommandResults): + return { + "instance_id": aws_command_results.instance_id, + "response_code": aws_command_results.response_code, + "stdout": aws_command_results.stdout, + "stderr": aws_command_results.stderr, + "status": aws_command_results.status.name.lower(), + } From f83832dc3c64a980c288974c1ff6aee44d1e1489 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 11 May 2022 10:00:35 -0400 Subject: [PATCH 34/34] Island: Change method order in RemoteRun --- monkey/monkey_island/cc/resources/remote_run.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index c1b47d20a..de72cbc62 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -23,12 +23,6 @@ class RemoteRun(flask_restful.Resource): def __init__(self, aws_service: AWSService): self._aws_service = aws_service - def run_aws_monkeys(self, request_body) -> Sequence[AWSCommandResults]: - instances = request_body.get("instances") - island_ip = request_body.get("island_ip") - - return self._aws_service.run_agents_on_managed_instances(instances, island_ip) - @jwt_required def get(self): action = request.args.get("action") @@ -58,6 +52,12 @@ class RemoteRun(flask_restful.Resource): # default action return make_response({"error": "Invalid action"}, 500) + def run_aws_monkeys(self, request_body) -> Sequence[AWSCommandResults]: + instances = request_body.get("instances") + island_ip = request_body.get("island_ip") + + return self._aws_service.run_agents_on_managed_instances(instances, island_ip) + @staticmethod def _encode_results(results: Sequence[AWSCommandResults]): result = list(map(RemoteRun._aws_command_results_to_encodable_dict, results))