p34709852
/
monkey
forked from p15670423/monkey
1
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Add timeout parameter to start_infection_monkey_agent()

This commit is contained in:
Mike Salvatore 2022-05-11 08:18:40 -04:00
parent 3702ecbc8c
commit 7c0616eecd
3 changed files with 34 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
monkey/monkey_island/cc/services/aws/aws_command_runner.py
View File

@ -7,7 +7,6 @@ import botocore
from common.utils import Timer from common.utils import Timer
REMOTE_COMMAND_TIMEOUT = 5
STATUS_CHECK_SLEEP_TIME = 1 STATUS_CHECK_SLEEP_TIME = 1
LINUX_DOCUMENT_NAME = "AWS-RunShellScript" LINUX_DOCUMENT_NAME = "AWS-RunShellScript"
WINDOWS_DOCUMENT_NAME = "AWS-RunPowerShellScript" WINDOWS_DOCUMENT_NAME = "AWS-RunPowerShellScript"
@ -35,7 +34,11 @@ class AWSCommandResults:
def start_infection_monkey_agent( 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: ) -> AWSCommandResults:
""" """
Run a command on a remote AWS instance 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 = _get_run_agent_command(target_os, island_ip)
command_id = _run_command_async(aws_client, target_instance_id, target_os, command) 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) 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( 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 = Timer()
timer.set(REMOTE_COMMAND_TIMEOUT) timer.set(timeout)
while not timer.is_expired(): while not timer.is_expired():
time.sleep(STATUS_CHECK_SLEEP_TIME) time.sleep(STATUS_CHECK_SLEEP_TIME)

15
monkey/monkey_island/cc/services/aws/aws_service.py
View File

@ -11,6 +11,7 @@ from common.utils.code_utils import queue_to_list
from .aws_command_runner import AWSCommandResults, start_infection_monkey_agent from .aws_command_runner import AWSCommandResults, start_infection_monkey_agent
DEFAULT_REMOTE_COMMAND_TIMEOUT = 5
INSTANCE_INFORMATION_LIST_KEY = "InstanceInformationList" INSTANCE_INFORMATION_LIST_KEY = "InstanceInformationList"
INSTANCE_ID_KEY = "InstanceId" INSTANCE_ID_KEY = "InstanceId"
COMPUTER_NAME_KEY = "ComputerName" COMPUTER_NAME_KEY = "ComputerName"
@ -72,12 +73,16 @@ class AWSService:
raise err raise err
def run_agents_on_managed_instances( 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]: ) -> Sequence[AWSCommandResults]:
""" """
Run an agent on one or more managed AWS instances. Run an agent on one or more managed AWS instances.
:param instances: An iterable of instances that the agent will be run on :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 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 :return: A sequence of AWSCommandResults
""" """
@ -86,7 +91,7 @@ class AWSService:
for i in instances: for i in instances:
t = Thread( t = Thread(
target=self._run_agent_on_managed_instance, 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, daemon=True,
) )
t.start() t.start()
@ -98,10 +103,12 @@ class AWSService:
return queue_to_list(results_queue) return queue_to_list(results_queue)
def _run_agent_on_managed_instance( 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) 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) results_queue.put(command_results)

26
monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py
View File

@ -11,6 +11,7 @@ from monkey_island.cc.services.aws.aws_command_runner import (
start_infection_monkey_agent, start_infection_monkey_agent,
) )
TIMEOUT = 0.03
INSTANCE_ID = "BEEFFACE" INSTANCE_ID = "BEEFFACE"
ISLAND_IP = "127.0.0.1" ISLAND_IP = "127.0.0.1"
""" """
@ -129,9 +130,6 @@ def error_response():
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
def patch_timeouts(monkeypatch): def patch_timeouts(monkeypatch):
monkeypatch.setattr(
"monkey_island.cc.services.aws.aws_command_runner.REMOTE_COMMAND_TIMEOUT", 0.03
)
monkeypatch.setattr( monkeypatch.setattr(
"monkey_island.cc.services.aws.aws_command_runner.STATUS_CHECK_SLEEP_TIME", 0.01 "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): 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() successful_mock_client.send_command.assert_called_once()
call_args_kwargs = successful_mock_client.send_command.call_args[1] 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): 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() successful_mock_client.send_command.assert_called_once()
call_args_kwargs = successful_mock_client.send_command.call_args[1] 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): 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() successful_mock_client.send_command.assert_called_once()
call_args_kwargs = successful_mock_client.send_command.call_args[1] 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): 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() successful_mock_client.send_command.assert_called_once()
call_args_kwargs = successful_mock_client.send_command.call_args[1] 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): 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() successful_mock_client.send_command.assert_called_once()
call_args_kwargs = successful_mock_client.send_command.call_args[1] 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)) 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 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.send_command = MagicMock(return_value=send_command_response)
aws_client.get_command_invocation = MagicMock(return_value=in_progress_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 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.send_command = MagicMock(return_value=send_command_response)
aws_client.get_command_invocation = MagicMock(return_value=error_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 assert command_results.status == AWSCommandStatus.ERROR