From f3a5a7090b60b20df3e78b3f7421b3ada666f537 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Mon, 2 May 2022 14:51:06 +0300 Subject: [PATCH] Agent, Island, Common: Refactor AwsService from class to package This also changes AwsInstance from singleton and instead the aws_service package is used as one --- monkey/common/aws/aws_instance.py | 4 - monkey/common/aws/aws_service.py | 104 +++++++++++------- monkey/common/cmd/aws/aws_cmd_runner.py | 4 +- .../utils/aws_environment_check.py | 10 +- monkey/monkey_island/cc/app.py | 6 - .../monkey_island/cc/resources/remote_run.py | 11 +- .../monkey_island/cc/services/initialize.py | 5 + .../cc/services/remote_run_aws.py | 33 ------ .../cc/services/reporting/aws_exporter.py | 4 +- .../cc/services/reporting/exporter_init.py | 5 +- 10 files changed, 79 insertions(+), 107 deletions(-) diff --git a/monkey/common/aws/aws_instance.py b/monkey/common/aws/aws_instance.py index 9f7b81180..d99c87117 100644 --- a/monkey/common/aws/aws_instance.py +++ b/monkey/common/aws/aws_instance.py @@ -6,8 +6,6 @@ from typing import Optional, Tuple import requests -from common.utils.code_utils import Singleton - AWS_INSTANCE_METADATA_LOCAL_IP_ADDRESS = "169.254.169.254" AWS_LATEST_METADATA_URI_PREFIX = "http://{0}/latest/".format(AWS_INSTANCE_METADATA_LOCAL_IP_ADDRESS) ACCOUNT_ID_KEY = "accountId" @@ -29,8 +27,6 @@ class AwsInstance: Class which gives useful information about the current instance you're on. """ - __metaclass__ = Singleton - def __init__(self): self._is_instance, self._instance_info = AwsInstance._fetch_instance_info() diff --git a/monkey/common/aws/aws_service.py b/monkey/common/aws/aws_service.py index 84710f839..3198d6e0f 100644 --- a/monkey/common/aws/aws_service.py +++ b/monkey/common/aws/aws_service.py @@ -1,4 +1,7 @@ import logging +from functools import wraps +from threading import Event +from typing import Callable, Optional import boto3 import botocore @@ -26,49 +29,66 @@ def filter_instance_data_from_aws_response(response): ] -class AwsService(object): +aws_instance: Optional[AwsInstance] = None +AWS_INFO_FETCH_TIMEOUT = 10.0 # Seconds +init_done = Event() + + +def initialize(): + global aws_instance + aws_instance = AwsInstance() + init_done.set() + + +def wait_init_done(fnc: Callable): + @wraps(fnc) + def inner(): + awaited = init_done.wait(AWS_INFO_FETCH_TIMEOUT) + if not awaited: + logger.error( + f"AWS service couldn't initialize in time! " + f"Current timeout is {AWS_INFO_FETCH_TIMEOUT}, " + f"but AWS info took longer to fetch from metadata server." + ) + return + fnc() + + return inner + + +@wait_init_done +def is_on_aws(): + return aws_instance.is_instance + + +@wait_init_done +def get_region(): + return aws_instance.region + + +@wait_init_done +def get_client(client_type): + return boto3.client(client_type, region_name=aws_instance.region) + + +@wait_init_done +def get_instances(): """ - A wrapper class around the boto3 client and session modules, which supplies various AWS - services. + Get the information for all instances with the relevant roles. - This class will assume: - 1. That it's running on an EC2 instance - 2. That the instance is associated with the correct IAM role. See - https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#iam-role - for details. + This function will assume that it's running on an EC2 instance with the correct IAM role. + See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#iam + -role for details. + + :raises: botocore.exceptions.ClientError if can't describe local instance information. + :return: All visible instances from this instance """ + local_ssm_client = boto3.client("ssm", aws_instance.region) + try: + response = local_ssm_client.describe_instance_information() - region = None - - @staticmethod - def set_region(region): - AwsService.region = region - - @staticmethod - def get_client(client_type, region=None): - return boto3.client( - client_type, region_name=region if region is not None else AwsService.region - ) - - @staticmethod - def get_instances(): - """ - Get the information for all instances with the relevant roles. - - This function will assume that it's running on an EC2 instance with the correct IAM role. - See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#iam - -role for details. - - :raises: botocore.exceptions.ClientError if can't describe local instance information. - :return: All visible instances from this instance - """ - current_instance = AwsInstance() - local_ssm_client = boto3.client("ssm", current_instance.region) - try: - response = local_ssm_client.describe_instance_information() - - filtered_instances_data = filter_instance_data_from_aws_response(response) - return filtered_instances_data - except botocore.exceptions.ClientError as e: - logger.warning("AWS client error while trying to get instances: " + e) - raise e + filtered_instances_data = filter_instance_data_from_aws_response(response) + return filtered_instances_data + except botocore.exceptions.ClientError as e: + logger.warning("AWS client error while trying to get instances: " + e) + raise e diff --git a/monkey/common/cmd/aws/aws_cmd_runner.py b/monkey/common/cmd/aws/aws_cmd_runner.py index 03d4e9e4f..1e00c6b35 100644 --- a/monkey/common/cmd/aws/aws_cmd_runner.py +++ b/monkey/common/cmd/aws/aws_cmd_runner.py @@ -1,7 +1,7 @@ import logging import time -from common.aws.aws_service import AwsService +from common.aws import aws_service from common.cmd.aws.aws_cmd_result import AwsCmdResult from common.cmd.cmd_runner import CmdRunner from common.cmd.cmd_status import CmdStatus @@ -18,7 +18,7 @@ class AwsCmdRunner(CmdRunner): super(AwsCmdRunner, self).__init__(is_linux) self.instance_id = instance_id self.region = region - self.ssm = AwsService.get_client("ssm", region) + self.ssm = aws_service.get_client("ssm", region) def query_command(self, command_id): time.sleep(2) diff --git a/monkey/infection_monkey/utils/aws_environment_check.py b/monkey/infection_monkey/utils/aws_environment_check.py index 6bdbb5c85..fe7316e57 100644 --- a/monkey/infection_monkey/utils/aws_environment_check.py +++ b/monkey/infection_monkey/utils/aws_environment_check.py @@ -1,6 +1,6 @@ import logging -from common.aws.aws_instance import AwsInstance +from common.aws import aws_service from infection_monkey.telemetry.aws_instance_telem import AWSInstanceTelemetry from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import ( LegacyTelemetryMessengerAdapter, @@ -10,16 +10,12 @@ from infection_monkey.utils.threading import create_daemon_thread logger = logging.getLogger(__name__) -def _running_on_aws(aws_instance: AwsInstance) -> bool: - return aws_instance.is_instance - - def _report_aws_environment(telemetry_messenger: LegacyTelemetryMessengerAdapter): logger.info("Collecting AWS info") - aws_instance = AwsInstance() + aws_instance = aws_service.initialize() - if _running_on_aws(aws_instance): + if aws_service.is_on_aws(): logger.info("Machine is an AWS instance") telemetry_messenger.send_telemetry(AWSInstanceTelemetry(aws_instance.instance_id)) else: diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 97f7dfca1..60cc0a026 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,7 +1,6 @@ import os import uuid from datetime import timedelta -from threading import Thread from typing import Type import flask_restful @@ -49,7 +48,6 @@ from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFinding from monkey_island.cc.resources.zero_trust.zero_trust_report import ZeroTrustReport from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.server_utils.custom_json_encoder import CustomJSONEncoder -from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService from monkey_island.cc.services.representations import output_json HOME_FILE = "index.html" @@ -104,10 +102,6 @@ def init_app_services(app): with app.app_context(): database.init() - # If running on AWS, this will initialize the instance data, which is used "later" in the - # execution of the island. Run on a daemon thread since it's slow. - Thread(target=RemoteRunAwsService.init, name="AWS check thread", daemon=True).start() - def init_app_url_rules(app): app.add_url_rule("/", "serve_home", serve_home) diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index 864fb4848..dd9a4eaf6 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -4,7 +4,7 @@ import flask_restful from botocore.exceptions import ClientError, NoCredentialsError from flask import jsonify, make_response, request -from common.aws.aws_service import AwsService +from common.aws import aws_service from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService @@ -19,10 +19,6 @@ NO_CREDS_ERROR_FORMAT = ( class RemoteRun(flask_restful.Resource): - def __init__(self): - super(RemoteRun, self).__init__() - RemoteRunAwsService.init() - def run_aws_monkeys(self, request_body): instances = request_body.get("instances") island_ip = request_body.get("island_ip") @@ -32,11 +28,11 @@ class RemoteRun(flask_restful.Resource): def get(self): action = request.args.get("action") if action == "list_aws": - is_aws = RemoteRunAwsService.is_running_on_aws() + is_aws = aws_service.is_on_aws() resp = {"is_aws": is_aws} if is_aws: try: - resp["instances"] = AwsService.get_instances() + resp["instances"] = aws_service.get_instances() except NoCredentialsError as e: resp["error"] = NO_CREDS_ERROR_FORMAT.format(e) return jsonify(resp) @@ -52,7 +48,6 @@ class RemoteRun(flask_restful.Resource): body = json.loads(request.data) resp = {} if body.get("type") == "aws": - RemoteRunAwsService.update_aws_region_authless() result = self.run_aws_monkeys(body) resp["result"] = result return jsonify(resp) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 4a4b2e4af..06b2473f8 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -1,6 +1,8 @@ from pathlib import Path +from threading import Thread from common import DIContainer +from common.aws import aws_service from monkey_island.cc.services import DirectoryFileStorageService, IFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -14,6 +16,9 @@ def initialize_services(data_dir: Path) -> DIContainer: IFileStorageService, DirectoryFileStorageService(data_dir / "custom_pbas") ) + # Takes a while so it's best to start it in the background + Thread(target=aws_service.initialize, name="AwsService initialization", daemon=True).start() + # This is temporary until we get DI all worked out. PostBreachFilesService.initialize(container.resolve(IFileStorageService)) LocalMonkeyRunService.initialize(data_dir) diff --git a/monkey/monkey_island/cc/services/remote_run_aws.py b/monkey/monkey_island/cc/services/remote_run_aws.py index dd15bff07..3dbb15477 100644 --- a/monkey/monkey_island/cc/services/remote_run_aws.py +++ b/monkey/monkey_island/cc/services/remote_run_aws.py @@ -1,34 +1,13 @@ import logging -from threading import Event -from common.aws.aws_instance import AwsInstance -from common.aws.aws_service import AwsService from common.cmd.aws.aws_cmd_runner import AwsCmdRunner from common.cmd.cmd import Cmd from common.cmd.cmd_runner import CmdRunner logger = logging.getLogger(__name__) -AWS_INFO_FETCH_TIMEOUT = 10 # Seconds -aws_info_fetch_done = Event() class RemoteRunAwsService: - aws_instance = None - - def __init__(self): - pass - - @staticmethod - def init(): - """ - Initializes service. Subsequent calls to this function have no effect. - Must be called at least once (in entire monkey lifetime) before usage of functions - :return: None - """ - if RemoteRunAwsService.aws_instance is None: - RemoteRunAwsService.aws_instance = AwsInstance() - aws_info_fetch_done.set() - @staticmethod def run_aws_monkeys(instances, island_ip): """ @@ -47,18 +26,6 @@ class RemoteRunAwsService: lambda _, result: result.is_success, ) - @staticmethod - def is_running_on_aws(): - aws_info_fetch_done.wait(AWS_INFO_FETCH_TIMEOUT) - return RemoteRunAwsService.aws_instance.is_instance - - @staticmethod - def update_aws_region_authless(): - """ - Updates the AWS region without auth params (via IAM role) - """ - AwsService.set_region(RemoteRunAwsService.aws_instance.region) - @staticmethod def _run_aws_monkey_cmd_async(instance_id, is_linux, island_ip): """ diff --git a/monkey/monkey_island/cc/services/reporting/aws_exporter.py b/monkey/monkey_island/cc/services/reporting/aws_exporter.py index ed8925985..41d133b19 100644 --- a/monkey/monkey_island/cc/services/reporting/aws_exporter.py +++ b/monkey/monkey_island/cc/services/reporting/aws_exporter.py @@ -5,6 +5,7 @@ from datetime import datetime import boto3 from botocore.exceptions import UnknownServiceError +from common.aws import aws_service from common.aws.aws_instance import AwsInstance from monkey_island.cc.services.reporting.exporter import Exporter @@ -35,8 +36,7 @@ class AWSExporter(Exporter): logger.info("No issues were found by the monkey, no need to send anything") return True - # Not suppressing error here on purpose. - current_aws_region = AwsInstance().region + current_aws_region = aws_service.get_region() for machine in issues_list: for issue in issues_list[machine]: diff --git a/monkey/monkey_island/cc/services/reporting/exporter_init.py b/monkey/monkey_island/cc/services/reporting/exporter_init.py index c19f3d5e3..8a84cdadd 100644 --- a/monkey/monkey_island/cc/services/reporting/exporter_init.py +++ b/monkey/monkey_island/cc/services/reporting/exporter_init.py @@ -1,6 +1,6 @@ import logging -from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService +from common.aws import aws_service from monkey_island.cc.services.reporting.aws_exporter import AWSExporter from monkey_island.cc.services.reporting.report_exporter_manager import ReportExporterManager @@ -22,8 +22,7 @@ def populate_exporter_list(): def try_add_aws_exporter_to_manager(manager): # noinspection PyBroadException try: - RemoteRunAwsService.init() - if RemoteRunAwsService.is_running_on_aws(): + if aws_service.is_on_aws(): manager.add_exporter_to_list(AWSExporter) except Exception: logger.error("Failed adding aws exporter to manager. Exception info:", exc_info=True)