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

@@ -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()
 

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

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)

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:

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)

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)

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)

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):
         """

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]:

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)