Island: Implement AWSService._run_agent_on_managed_instance()

This commit is contained in:
Mike Salvatore 2022-05-09 15:43:52 -04:00
parent 653bfbd24b
commit 144506c32d
2 changed files with 27 additions and 5 deletions

View File

@ -11,6 +11,7 @@ STATUS_CHECK_SLEEP_TIME = 1
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# TODO: Make sure the return type is compatible with what RemoteRun is expecting. Add typehint.
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
): ):
@ -21,6 +22,8 @@ def start_infection_monkey_agent(
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)
# TODO: Return result
def _get_run_agent_command(target_os: str, island_ip: str): def _get_run_agent_command(target_os: str, island_ip: str):
if target_os == "linux": if target_os == "linux":

View File

@ -6,6 +6,8 @@ import botocore
from common.aws.aws_instance import AWSInstance from common.aws.aws_instance import AWSInstance
from .aws_command_runner import start_infection_monkey_agent
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"
@ -66,12 +68,29 @@ class AWSService:
logger.warning("AWS client error while trying to get manage dinstances: {err}") logger.warning("AWS client error while trying to get manage dinstances: {err}")
raise err raise err
def run_agent_on_managed_instances(self, instance_ids: Iterable[str]): # TODO: Determine the return type
for id_ in instance_ids: def run_agents_on_managed_instances(
self._run_agent_on_managed_instance(id_) 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): results = []
pass # 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]]): def _filter_relevant_instance_info(raw_managed_instances_info: Sequence[Mapping[str, Any]]):