From 144506c32d586125c5258e76121b9a37fb86c6a1 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Mon, 9 May 2022 15:43:52 -0400
Subject: [PATCH] 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]]):