Merge pull request #522 from guardicore/feature/521-environment-check

Feature/521 environment check

monkey/common/cloud/all_instances.py

@@ -0,0 +1,12 @@
+from typing import List
+
+from common.cloud.aws.aws_instance import AwsInstance
+from common.cloud.azure.azure_instance import AzureInstance
+from common.cloud.gcp.gcp_instance import GcpInstance
+from common.cloud.instance import CloudInstance
+
+all_cloud_instances = [AwsInstance(), AzureInstance(), GcpInstance()]
+
+
+def get_all_cloud_instances() -> List[CloudInstance]:
+    return all_cloud_instances

monkey/common/cloud/aws/aws_instance.py

@@ -6,6 +6,9 @@ import logging
 
 __author__ = 'itay.mizeretz'
 
+from common.cloud.environment_names import Environment
+from common.cloud.instance import CloudInstance
+
 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"
@@ -13,10 +16,15 @@ ACCOUNT_ID_KEY = "accountId"
 logger = logging.getLogger(__name__)
 
 
-class AwsInstance(object):
+class AwsInstance(CloudInstance):
     """
     Class which gives useful information about the current instance you're on.
     """
+    def is_instance(self):
+        return self.instance_id is not None
+
+    def get_cloud_provider_name(self) -> Environment:
+        return Environment.AWS
 
     def __init__(self):
         self.instance_id = None
@@ -57,9 +65,6 @@ class AwsInstance(object):
     def get_region(self):
         return self.region
 
-    def is_aws_instance(self):
-        return self.instance_id is not None
-
     @staticmethod
     def _extract_account_id(instance_identity_document_response):
         """

monkey/common/cloud/aws/aws_service.py

@@ -4,7 +4,7 @@ import boto3
 import botocore
 from botocore.exceptions import ClientError
 
-from common.cloud.aws_instance import AwsInstance
+from common.cloud.aws.aws_instance import AwsInstance
 
 __author__ = ['itay.mizeretz', 'shay.nehmad']
 

monkey/common/cloud/azure/azure_instance.py

@@ -0,0 +1,55 @@
+import logging
+import requests
+
+from common.cloud.environment_names import Environment
+from common.cloud.instance import CloudInstance
+
+LATEST_AZURE_METADATA_API_VERSION = "2019-04-30"
+AZURE_METADATA_SERVICE_URL = "http://169.254.169.254/metadata/instance?api-version=%s" % LATEST_AZURE_METADATA_API_VERSION
+
+logger = logging.getLogger(__name__)
+
+
+class AzureInstance(CloudInstance):
+    """
+    Access to useful information about the current machine if it's an Azure VM.
+    Based on Azure metadata service: https://docs.microsoft.com/en-us/azure/virtual-machines/windows/instance-metadata-service
+    """
+    def is_instance(self):
+        return self.on_azure
+
+    def get_cloud_provider_name(self) -> Environment:
+        return Environment.AZURE
+
+    def __init__(self):
+        """
+        Determines if on Azure and if so, gets some basic metadata on this instance.
+        """
+        self.instance_name = None
+        self.instance_id = None
+        self.location = None
+        self.on_azure = False
+
+        try:
+            response = requests.get(AZURE_METADATA_SERVICE_URL, headers={"Metadata": "true"})
+            self.on_azure = True
+
+            # If not on cloud, the metadata URL is non-routable and the connection will fail.
+            # If on AWS, should get 404 since the metadata service URL is different, so bool(response) will be false.
+            if response:
+                logger.debug("On Azure. Trying to parse metadata.")
+                self.try_parse_response(response)
+            else:
+                logger.warning("On Azure, but metadata response not ok: {}".format(response.status_code))
+        except requests.RequestException:
+            logger.debug("Failed to get response from Azure metadata service: This instance is not on Azure.")
+            self.on_azure = False
+
+    def try_parse_response(self, response):
+        try:
+            response_data = response.json()
+            self.instance_name = response_data["compute"]["name"]
+            self.instance_id = response_data["compute"]["vmId"]
+            self.location = response_data["compute"]["location"]
+        except KeyError:
+            logger.exception("Error while parsing response from Azure metadata service.")

monkey/common/cloud/environment_names.py

@@ -0,0 +1,15 @@
+from enum import Enum
+
+
+class Environment(Enum):
+    UNKNOWN = "Unknown"
+    ON_PREMISE = "On Premise"
+    AZURE = "Azure"
+    AWS = "AWS"
+    GCP = "GCP"
+    ALIBABA = "Alibaba Cloud"
+    IBM = "IBM Cloud"
+    DigitalOcean = "Digital Ocean"
+
+
+ALL_ENVIRONMENTS_NAMES = [x.value for x in Environment]

monkey/common/cloud/gcp/gcp_instance.py

@@ -0,0 +1,43 @@
+import logging
+import requests
+
+from common.cloud.environment_names import Environment
+from common.cloud.instance import CloudInstance
+
+logger = logging.getLogger(__name__)
+
+
+GCP_METADATA_SERVICE_URL = "http://metadata.google.internal/"
+
+
+class GcpInstance(CloudInstance):
+    """
+    Used to determine if on GCP. See https://cloud.google.com/compute/docs/storing-retrieving-metadata#runninggce
+    """
+    def is_instance(self):
+        return self.on_gcp
+
+    def get_cloud_provider_name(self) -> Environment:
+        return Environment.GCP
+
+    def __init__(self):
+        self.on_gcp = False
+
+        try:
+            # If not on GCP, this domain shouldn't resolve.
+            response = requests.get(GCP_METADATA_SERVICE_URL)
+
+            if response:
+                logger.debug("Got ok metadata response: on GCP")
+                self.on_gcp = True
+
+                if "Metadata-Flavor" not in response.headers:
+                    logger.warning("Got unexpected GCP Metadata format")
+                else:
+                    if not response.headers["Metadata-Flavor"] == "Google":
+                        logger.warning("Got unexpected Metadata flavor: {}".format(response.headers["Metadata-Flavor"]))
+            else:
+                logger.warning("On GCP, but metadata response not ok: {}".format(response.status_code))
+        except requests.RequestException:
+            logger.debug("Failed to get response from GCP metadata service: This instance is not on GCP")
+            self.on_gcp = False

monkey/common/cloud/instance.py

@@ -0,0 +1,14 @@
+from common.cloud.environment_names import Environment
+
+
+class CloudInstance(object):
+    """
+    This is an abstract class which represents a cloud instance.
+
+    The current machine can be a cloud instance (for example EC2 instance or Azure VM).
+    """
+    def is_instance(self) -> bool:
+        raise NotImplementedError()
+
+    def get_cloud_provider_name(self) -> Environment:
+        raise NotImplementedError()

monkey/common/cmd/aws/aws_cmd_runner.py

@@ -1,6 +1,6 @@
 import logging
 
-from common.cloud.aws_service import AwsService
+from common.cloud.aws.aws_service import AwsService
 from common.cmd.aws.aws_cmd_result import AwsCmdResult
 from common.cmd.cmd_runner import CmdRunner
 from common.cmd.cmd_status import CmdStatus

monkey/common/data/system_info_collectors_names.py

@@ -0,0 +1,4 @@
+AWS_COLLECTOR = "AwsCollector"
+HOSTNAME_COLLECTOR = "HostnameCollector"
+ENVIRONMENT_COLLECTOR = "EnvironmentCollector"
+PROCESS_LIST_COLLECTOR = "ProcessListCollector"

monkey/infection_monkey/config.py

@@ -1,6 +1,6 @@
 import hashlib
-import os
 import json
+import os
 import sys
 import uuid
 from abc import ABCMeta
@@ -125,6 +125,7 @@ class Configuration(object):
 
     finger_classes = []
     exploiter_classes = []
+    system_info_collectors_classes = []
 
     # how many victims to look for in a single scan iteration
     victims_max_find = 100

monkey/infection_monkey/monkey.py

@@ -32,11 +32,17 @@ from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from infection_monkey.exploit.HostExploiter import HostExploiter
 
+MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, shutting down"
+
 __author__ = 'itamar'
 
 LOG = logging.getLogger(__name__)
 
 
+class PlannedShutdownException(Exception):
+    pass
+
+
 class InfectionMonkey(object):
     def __init__(self, args):
         self._keep_running = False
@@ -87,22 +93,18 @@ class InfectionMonkey(object):
                 LOG.debug("Default server: %s is already in command servers list" % self._default_server)
 
     def start(self):
-        LOG.info("Monkey is running...")
+        try:
+            LOG.info("Monkey is starting...")
 
+            LOG.debug("Starting the setup phase.")
             # Sets island's IP and port for monkey to communicate to
-        if not self.set_default_server():
+            self.set_default_server()
-            return
             self.set_default_port()
 
             # Create a dir for monkey files if there isn't one
             create_monkey_dir()
 
-        if WindowsUpgrader.should_upgrade():
+            self.upgrade_to_64_if_needed()
-            self._upgrading_to_64 = True
-            self._singleton.unlock()
-            LOG.info("32bit monkey running on 64bit Windows. Upgrading.")
-            WindowsUpgrader.upgrade(self._opts)
-            return
 
             ControlClient.wakeup(parent=self._parent)
             ControlClient.load_control_config()
@@ -110,9 +112,7 @@ class InfectionMonkey(object):
             if is_windows_os():
                 T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
-        if not WormConfiguration.alive:
+            self.shutdown_by_not_alive_config()
-            LOG.info("Marked not alive from configuration")
-            return
 
             if firewall.is_enabled():
                 firewall.add_firewall_rule()
@@ -124,20 +124,12 @@ class InfectionMonkey(object):
             StateTelem(is_done=False).send()
             TunnelTelem().send()
 
-        if WormConfiguration.collect_system_info:
+            LOG.debug("Starting the post-breach phase.")
-            LOG.debug("Calling system info collection")
+            self.collect_system_info_if_configured()
-            system_info_collector = SystemInfoCollector()
+            PostBreach().execute_all_configured()
-            system_info = system_info_collector.get_info()
-            SystemInfoTelem(system_info).send()
 
-        # Executes post breach actions
+            LOG.debug("Starting the propagation phase.")
-        PostBreach().execute()
+            self.shutdown_by_max_depth_reached()
-
-        if 0 == WormConfiguration.depth:
-            TraceTelem("Reached max depth, shutting down").send()
-            return
-        else:
-            LOG.debug("Running with depth: %d" % WormConfiguration.depth)
 
             for iteration_index in range(WormConfiguration.max_iterations):
                 ControlClient.keepalive()
@@ -224,6 +216,35 @@ class InfectionMonkey(object):
             if monkey_tunnel:
                 monkey_tunnel.stop()
                 monkey_tunnel.join()
+        except PlannedShutdownException:
+            LOG.info("A planned shutdown of the Monkey occurred. Logging the reason and finishing execution.")
+            LOG.exception("Planned shutdown, reason:")
+
+    def shutdown_by_max_depth_reached(self):
+        if 0 == WormConfiguration.depth:
+            TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send()
+            raise PlannedShutdownException(MAX_DEPTH_REACHED_MESSAGE)
+        else:
+            LOG.debug("Running with depth: %d" % WormConfiguration.depth)
+
+    def collect_system_info_if_configured(self):
+        if WormConfiguration.collect_system_info:
+            LOG.debug("Calling system info collection")
+            system_info_collector = SystemInfoCollector()
+            system_info = system_info_collector.get_info()
+            SystemInfoTelem(system_info).send()
+
+    def shutdown_by_not_alive_config(self):
+        if not WormConfiguration.alive:
+            raise PlannedShutdownException("Marked 'not alive' from configuration.")
+
+    def upgrade_to_64_if_needed(self):
+        if WindowsUpgrader.should_upgrade():
+            self._upgrading_to_64 = True
+            self._singleton.unlock()
+            LOG.info("32bit monkey running on 64bit Windows. Upgrading.")
+            WindowsUpgrader.upgrade(self._opts)
+            raise PlannedShutdownException("Finished upgrading from 32bit to 64bit.")
 
     def cleanup(self):
         LOG.info("Monkey cleanup started")
@@ -346,9 +367,11 @@ class InfectionMonkey(object):
             self._default_server_port = ''
 
     def set_default_server(self):
+        """
+        Sets the default server for the Monkey to communicate back to.
+        :raises PlannedShutdownException if couldn't find the server.
+        """
         if not ControlClient.find_server(default_tunnel=self._default_tunnel):
-            LOG.info("Monkey couldn't find server. Going down.")
+            raise PlannedShutdownException("Monkey couldn't find server with {} default tunnel.".format(self._default_tunnel))
-            return False
         self._default_server = WormConfiguration.current_server
         LOG.debug("default server set to: %s" % self._default_server)
-        return True

monkey/infection_monkey/post_breach/post_breach_handler.py

@@ -20,7 +20,7 @@ class PostBreach(object):
         self.os_is_linux = not is_windows_os()
         self.pba_list = self.config_to_pba_list()
 
-    def execute(self):
+    def execute_all_configured(self):
         """
         Executes all post breach actions.
         """

monkey/infection_monkey/system_info/__init__.py

@@ -6,9 +6,9 @@ import psutil
 from enum import IntEnum
 
 from infection_monkey.network.info import get_host_subnets
-from infection_monkey.system_info.aws_collector import AwsCollector
 from infection_monkey.system_info.azure_cred_collector import AzureCollector
 from infection_monkey.system_info.netstat_collector import NetstatCollector
+from system_info.system_info_collectors_handler import SystemInfoCollectorsHandler
 
 LOG = logging.getLogger(__name__)
 
@@ -61,50 +61,12 @@ class InfoCollector(object):
         self.info = {}
 
     def get_info(self):
-        self.get_hostname()
+        # Collect all hardcoded
-        self.get_process_list()
         self.get_network_info()
         self.get_azure_info()
-        self.get_aws_info()
 
-    def get_hostname(self):
+        # Collect all plugins
-        """
+        SystemInfoCollectorsHandler().execute_all_configured()
-        Adds the fully qualified computer hostname to the system information.
-        :return: None. Updates class information
-        """
-        LOG.debug("Reading hostname")
-        self.info['hostname'] = socket.getfqdn()
-
-    def get_process_list(self):
-        """
-        Adds process information from the host to the system information.
-        Currently lists process name, ID, parent ID, command line
-        and the full image path of each process.
-        :return: None. Updates class information
-        """
-        LOG.debug("Reading process list")
-        processes = {}
-        for process in psutil.process_iter():
-            try:
-                processes[process.pid] = {"name": process.name(),
-                                          "pid": process.pid,
-                                          "ppid": process.ppid(),
-                                          "cmdline": " ".join(process.cmdline()),
-                                          "full_image_path": process.exe(),
-                                          }
-            except (psutil.AccessDenied, WindowsError):
-                # we may be running as non root
-                # and some processes are impossible to acquire in Windows/Linux
-                # in this case we'll just add what we can
-                processes[process.pid] = {"name": "null",
-                                          "pid": process.pid,
-                                          "ppid": process.ppid(),
-                                          "cmdline": "ACCESS DENIED",
-                                          "full_image_path": "null",
-                                          }
-                continue
-
-        self.info['process_list'] = processes
 
     def get_network_info(self):
         """
@@ -150,11 +112,3 @@ class InfoCollector(object):
         except Exception:
             # If we failed to collect azure info, no reason to fail all the collection. Log and continue.
             LOG.error("Failed collecting Azure info.", exc_info=True)
-
-    def get_aws_info(self):
-        # noinspection PyBroadException
-        try:
-            self.info['aws'] = AwsCollector().get_aws_info()
-        except Exception:
-            # If we failed to collect aws info, no reason to fail all the collection. Log and continue.
-            LOG.error("Failed collecting AWS info.", exc_info=True)

monkey/infection_monkey/system_info/aws_collector.py

@@ -1,29 +0,0 @@
-import logging
-
-from common.cloud.aws_instance import AwsInstance
-
-__author__ = 'itay.mizeretz'
-
-LOG = logging.getLogger(__name__)
-
-
-class AwsCollector(object):
-    """
-    Extract info from AWS machines
-    """
-
-    @staticmethod
-    def get_aws_info():
-        LOG.info("Collecting AWS info")
-        aws = AwsInstance()
-        info = {}
-        if aws.is_aws_instance():
-            LOG.info("Machine is an AWS instance")
-            info = \
-                {
-                    'instance_id': aws.get_instance_id()
-                }
-        else:
-            LOG.info("Machine is NOT an AWS instance")
-
-        return info

monkey/infection_monkey/system_info/collectors/__init__.py

@@ -0,0 +1,3 @@
+"""
+This package holds all the dynamic (plugin) collectors
+"""

monkey/infection_monkey/system_info/collectors/aws_collector.py

@@ -0,0 +1,31 @@
+import logging
+
+from common.cloud.aws.aws_instance import AwsInstance
+from common.data.system_info_collectors_names import AWS_COLLECTOR
+from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+
+
+logger = logging.getLogger(__name__)
+
+
+class AwsCollector(SystemInfoCollector):
+    """
+    Extract info from AWS machines.
+    """
+    def __init__(self):
+        super().__init__(name=AWS_COLLECTOR)
+
+    def collect(self) -> dict:
+        logger.info("Collecting AWS info")
+        aws = AwsInstance()
+        info = {}
+        if aws.is_instance():
+            logger.info("Machine is an AWS instance")
+            info = \
+                {
+                    'instance_id': aws.get_instance_id()
+                }
+        else:
+            logger.info("Machine is NOT an AWS instance")
+
+        return info

monkey/infection_monkey/system_info/collectors/environment_collector.py

@@ -0,0 +1,24 @@
+from common.cloud.all_instances import get_all_cloud_instances
+from common.cloud.environment_names import Environment
+from common.data.system_info_collectors_names import ENVIRONMENT_COLLECTOR
+from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+
+
+def get_monkey_environment() -> Environment:
+    """
+    Get the Monkey's running environment.
+    :return: One of the cloud providers if on cloud; otherwise, assumes "on premise".
+    """
+    for instance in get_all_cloud_instances():
+        if instance.is_instance():
+            return instance.get_cloud_provider_name()
+
+    return Environment.ON_PREMISE
+
+
+class EnvironmentCollector(SystemInfoCollector):
+    def __init__(self):
+        super().__init__(name=ENVIRONMENT_COLLECTOR)
+
+    def collect(self) -> dict:
+        return {"environment": get_monkey_environment()}

monkey/infection_monkey/system_info/collectors/hostname_collector.py

@@ -0,0 +1,16 @@
+import logging
+import socket
+
+from common.data.system_info_collectors_names import HOSTNAME_COLLECTOR
+from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+
+
+logger = logging.getLogger(__name__)
+
+
+class HostnameCollector(SystemInfoCollector):
+    def __init__(self):
+        super().__init__(name=HOSTNAME_COLLECTOR)
+
+    def collect(self) -> dict:
+        return {"hostname": socket.getfqdn()}

monkey/infection_monkey/system_info/collectors/process_list_collector.py

@@ -0,0 +1,50 @@
+import logging
+import psutil
+
+from common.data.system_info_collectors_names import PROCESS_LIST_COLLECTOR
+from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+
+logger = logging.getLogger(__name__)
+
+# Linux doesn't have WindowsError
+try:
+    WindowsError
+except NameError:
+    # noinspection PyShadowingBuiltins
+    WindowsError = psutil.AccessDenied
+
+
+class ProcessListCollector(SystemInfoCollector):
+    def __init__(self):
+        super().__init__(name=PROCESS_LIST_COLLECTOR)
+
+    def collect(self) -> dict:
+        """
+        Adds process information from the host to the system information.
+        Currently lists process name, ID, parent ID, command line
+        and the full image path of each process.
+        """
+        logger.debug("Reading process list")
+        processes = {}
+        for process in psutil.process_iter():
+            try:
+                processes[process.pid] = {
+                    "name": process.name(),
+                    "pid": process.pid,
+                    "ppid": process.ppid(),
+                    "cmdline": " ".join(process.cmdline()),
+                    "full_image_path": process.exe(),
+                }
+            except (psutil.AccessDenied, WindowsError):
+                # we may be running as non root and some processes are impossible to acquire in Windows/Linux.
+                # In this case we'll just add what we know.
+                processes[process.pid] = {
+                    "name": "null",
+                    "pid": process.pid,
+                    "ppid": process.ppid(),
+                    "cmdline": "ACCESS DENIED",
+                    "full_image_path": "null",
+                }
+                continue
+
+        return {'process_list': processes}

monkey/infection_monkey/system_info/system_info_collector.py

@@ -0,0 +1,38 @@
+from infection_monkey.config import WormConfiguration
+from infection_monkey.utils.plugins.plugin import Plugin
+from abc import ABCMeta, abstractmethod
+
+import infection_monkey.system_info.collectors
+
+
+class SystemInfoCollector(Plugin, metaclass=ABCMeta):
+    """
+    ABC for system info collection. See system_info_collector_handler for more info. Basically, to implement a new system info
+    collector, inherit from this class in an implementation in the infection_monkey.system_info.collectors class, and override
+    the 'collect' method. Don't forget to parse your results in the Monkey Island and to add the collector to the configuration
+    as well - see monkey_island.cc.services.processing.system_info_collectors for examples.
+
+    See the Wiki page "How to add a new System Info Collector to the Monkey?" for a detailed guide.
+    """
+    def __init__(self, name="unknown"):
+        self.name = name
+
+    @staticmethod
+    def should_run(class_name) -> bool:
+        return class_name in WormConfiguration.system_info_collectors_classes
+
+    @staticmethod
+    def base_package_file():
+        return infection_monkey.system_info.collectors.__file__
+
+    @staticmethod
+    def base_package_name():
+        return infection_monkey.system_info.collectors.__package__
+
+    @abstractmethod
+    def collect(self) -> dict:
+        """
+        Collect the relevant information and return it in a dictionary.
+        To be implemented by each collector.
+        """
+        raise NotImplementedError()

monkey/infection_monkey/system_info/system_info_collectors_handler.py

@@ -0,0 +1,33 @@
+import logging
+from typing import Sequence
+
+from infection_monkey.system_info.system_info_collector import SystemInfoCollector
+from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
+
+LOG = logging.getLogger(__name__)
+
+
+class SystemInfoCollectorsHandler(object):
+    def __init__(self):
+        self.collectors_list = self.config_to_collectors_list()
+
+    def execute_all_configured(self):
+        successful_collections = 0
+        system_info_telemetry = {}
+        for collector in self.collectors_list:
+            try:
+                LOG.debug("Executing system info collector: '{}'".format(collector.name))
+                collected_info = collector.collect()
+                system_info_telemetry[collector.name] = collected_info
+                successful_collections += 1
+            except Exception as e:
+                # If we failed one collector, no need to stop execution. Log and continue.
+                LOG.error("Collector {} failed. Error info: {}".format(collector.name, e))
+        LOG.info("All system info collectors executed. Total {} executed, out of which {} collected successfully.".
+                 format(len(self.collectors_list), successful_collections))
+
+        SystemInfoTelem({"collectors": system_info_telemetry}).send()
+
+    @staticmethod
+    def config_to_collectors_list() -> Sequence[SystemInfoCollector]:
+        return SystemInfoCollector.get_instances()

monkey/monkey_island/cc/environment/aws.py

@@ -1,6 +1,6 @@
 import monkey_island.cc.auth
 from monkey_island.cc.environment import Environment
-from common.cloud.aws_instance import AwsInstance
+from common.cloud.aws.aws_instance import AwsInstance
 
 __author__ = 'itay.mizeretz'
 

monkey/monkey_island/cc/models/monkey.py

@@ -9,6 +9,7 @@ from monkey_island.cc.models.monkey_ttl import MonkeyTtl, create_monkey_ttl_docu
 from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
 from monkey_island.cc.models.command_control_channel import CommandControlChannel
 from monkey_island.cc.utils import local_ip_addresses
+from common.cloud import environment_names
 
 MAX_MONKEYS_AMOUNT_TO_CACHE = 100
 
@@ -42,6 +43,9 @@ class Monkey(Document):
     ttl_ref = ReferenceField(MonkeyTtl)
     tunnel = ReferenceField("self")
     command_control_channel = EmbeddedDocumentField(CommandControlChannel)
+
+    # Environment related fields
+    environment = StringField(default=environment_names.Environment.UNKNOWN, choices=environment_names.ALL_ENVIRONMENTS_NAMES)
     aws_instance_id = StringField(required=False)  # This field only exists when the monkey is running on an AWS
 
     # instance. See https://github.com/guardicore/monkey/issues/426.
@@ -55,7 +59,8 @@ class Monkey(Document):
             raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex, str(db_id)))
 
     @staticmethod
-    def get_single_monkey_by_guid(monkey_guid):
+    # See https://www.python.org/dev/peps/pep-0484/#forward-references
+    def get_single_monkey_by_guid(monkey_guid) -> 'Monkey':
         try:
             return Monkey.objects.get(guid=monkey_guid)
         except DoesNotExist as ex:

monkey/monkey_island/cc/resources/remote_run.py

@@ -6,7 +6,7 @@ import flask_restful
 
 from monkey_island.cc.auth import jwt_required
 from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService
-from common.cloud.aws_service import AwsService
+from common.cloud.aws.aws_service import AwsService
 
 CLIENT_ERROR_FORMAT = "ClientError, error message: '{}'. Probably, the IAM role that has been associated with the " \
                       "instance doesn't permit SSM calls. "

monkey/monkey_island/cc/services/config_schema.py

@@ -1,3 +1,5 @@
+from common.data.system_info_collectors_names import *
+
 WARNING_SIGN = " \u26A0"
 
 SCHEMA = {
@@ -99,6 +101,44 @@ SCHEMA = {
                 }
             ]
         },
+        "system_info_collectors_classes": {
+            "title": "System Information Collectors",
+            "type": "string",
+            "anyOf": [
+                {
+                    "type": "string",
+                    "enum": [
+                        ENVIRONMENT_COLLECTOR
+                    ],
+                    "title": "Collect which environment this machine is on (on prem/cloud)",
+                    "attack_techniques": []
+                },
+                {
+                    "type": "string",
+                    "enum": [
+                        AWS_COLLECTOR
+                    ],
+                    "title": "If on AWS, collect more information about the instance",
+                    "attack_techniques": []
+                },
+                {
+                    "type": "string",
+                    "enum": [
+                        HOSTNAME_COLLECTOR
+                    ],
+                    "title": "Collect the machine's hostname",
+                    "attack_techniques": []
+                },
+{
+                    "type": "string",
+                    "enum": [
+                        PROCESS_LIST_COLLECTOR
+                    ],
+                    "title": "Collect running processes on the machine",
+                    "attack_techniques": []
+                },
+            ],
+        },
         "post_breach_acts": {
             "title": "Post breach actions",
             "type": "string",
@@ -433,6 +473,21 @@ SCHEMA = {
                             "attack_techniques": ["T1003"],
                             "description": "Determines whether to use Mimikatz"
                         },
+                        "system_info_collectors_classes": {
+                            "title": "System info collectors",
+                            "type": "array",
+                            "uniqueItems": True,
+                            "items": {
+                                "$ref": "#/definitions/system_info_collectors_classes"
+                            },
+                            "default": [
+                                ENVIRONMENT_COLLECTOR,
+                                AWS_COLLECTOR,
+                                HOSTNAME_COLLECTOR,
+                                PROCESS_LIST_COLLECTOR
+                            ],
+                            "description": "Determines which system information collectors will collect information."
+                        },
                     }
                 },
                 "life_cycle": {

monkey/monkey_island/cc/services/remote_run_aws.py

@@ -1,7 +1,7 @@
 import logging
 
-from common.cloud.aws_instance import AwsInstance
+from common.cloud.aws.aws_instance import AwsInstance
-from common.cloud.aws_service import AwsService
+from common.cloud.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
@@ -54,7 +54,7 @@ class RemoteRunAwsService:
 
     @staticmethod
     def is_running_on_aws():
-        return RemoteRunAwsService.aws_instance.is_aws_instance()
+        return RemoteRunAwsService.aws_instance.is_instance()
 
     @staticmethod
     def update_aws_region_authless():

monkey/monkey_island/cc/services/reporting/aws_exporter.py

@@ -5,7 +5,7 @@ from datetime import datetime
 import boto3
 from botocore.exceptions import UnknownServiceError
 
-from common.cloud.aws_instance import AwsInstance
+from common.cloud.aws.aws_instance import AwsInstance
 from monkey_island.cc.environment.environment import load_server_configuration_from_file
 from monkey_island.cc.services.reporting.exporter import Exporter
 

monkey/monkey_island/cc/services/reporting/report.py

@@ -386,10 +386,11 @@ class ReportService:
     @staticmethod
     def get_monkey_subnets(monkey_guid):
         network_info = mongo.db.telemetry.find_one(
-            {'telem_category': 'system_info', 'monkey_guid': monkey_guid},
+            {'telem_category': 'system_info',
+             'monkey_guid': monkey_guid},
             {'data.network_info.networks': 1}
         )
-        if network_info is None:
+        if network_info is None or not network_info["data"]:
             return []
 
         return \

monkey/monkey_island/cc/services/reporting/report_generation_synchronisation.py

@@ -15,8 +15,10 @@ __regular_report_generating_lock = threading.Semaphore()
 def safe_generate_reports():
     # Entering the critical section; Wait until report generation is available.
     __report_generating_lock.acquire()
+    try:
         report = safe_generate_regular_report()
         attack_report = safe_generate_attack_report()
+    finally:
         # Leaving the critical section.
         __report_generating_lock.release()
     return report, attack_report
@@ -25,8 +27,10 @@ def safe_generate_reports():
 def safe_generate_regular_report():
     # Local import to avoid circular imports
     from monkey_island.cc.services.reporting.report import ReportService
+    try:
         __regular_report_generating_lock.acquire()
         report = ReportService.generate_report()
+    finally:
         __regular_report_generating_lock.release()
     return report
 
@@ -34,8 +38,10 @@ def safe_generate_regular_report():
 def safe_generate_attack_report():
     # Local import to avoid circular imports
     from monkey_island.cc.services.attack.attack_report import AttackReportService
+    try:
         __attack_report_generating_lock.acquire()
         attack_report = AttackReportService.generate_new_report()
+    finally:
         __attack_report_generating_lock.release()
     return attack_report
 

monkey/monkey_island/cc/services/telemetry/processing/system_info.py

@@ -1,25 +1,23 @@
 import logging
 
-from monkey_island.cc.database import mongo
-from monkey_island.cc.models import Monkey
-from monkey_island.cc.services import mimikatz_utils
-from monkey_island.cc.services.node import NodeService
-from monkey_island.cc.services.config import ConfigService
-from monkey_island.cc.services.telemetry.zero_trust_tests.antivirus_existence import test_antivirus_existence
-from monkey_island.cc.services.wmi_handler import WMIHandler
 from monkey_island.cc.encryptor import encryptor
+from monkey_island.cc.services import mimikatz_utils
+from monkey_island.cc.services.config import ConfigService
+from monkey_island.cc.services.node import NodeService
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \
+    SystemInfoTelemetryDispatcher
+from monkey_island.cc.services.wmi_handler import WMIHandler
 
 logger = logging.getLogger(__name__)
 
 
 def process_system_info_telemetry(telemetry_json):
+    dispatcher = SystemInfoTelemetryDispatcher()
     telemetry_processing_stages = [
         process_ssh_info,
         process_credential_info,
         process_mimikatz_and_wmi_info,
-        process_aws_data,
+        dispatcher.dispatch_collector_results_to_relevant_processors
-        update_db_with_new_hostname,
-        test_antivirus_existence,
     ]
 
     # Calling safe_process_telemetry so if one of the stages fail, we log and move on instead of failing the rest of
@@ -34,7 +32,7 @@ def safe_process_telemetry(processing_function, telemetry_json):
         processing_function(telemetry_json)
     except Exception as err:
         logger.error(
-            "Error {} while in {} stage of processing telemetry.".format(str(err), processing_function.func_name),
+            "Error {} while in {} stage of processing telemetry.".format(str(err), processing_function.__name__),
             exc_info=True)
 
 
@@ -104,14 +102,3 @@ def process_mimikatz_and_wmi_info(telemetry_json):
         wmi_handler = WMIHandler(monkey_id, telemetry_json['data']['wmi'], users_secrets)
         wmi_handler.process_and_handle_wmi_info()
 
-
-def process_aws_data(telemetry_json):
-    if 'aws' in telemetry_json['data']:
-        if 'instance_id' in telemetry_json['data']['aws']:
-            monkey_id = NodeService.get_monkey_by_guid(telemetry_json['monkey_guid']).get('_id')
-            mongo.db.monkey.update_one({'_id': monkey_id},
-                                       {'$set': {'aws_instance_id': telemetry_json['data']['aws']['instance_id']}})
-
-
-def update_db_with_new_hostname(telemetry_json):
-    Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid']).set_hostname(telemetry_json['data']['hostname'])

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/aws.py

@@ -0,0 +1,15 @@
+import logging
+
+from monkey_island.cc.models.monkey import Monkey
+
+logger = logging.getLogger(__name__)
+
+
+def process_aws_telemetry(collector_results, monkey_guid):
+    relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
+
+    if "instance_id" in collector_results:
+        instance_id = collector_results["instance_id"]
+        relevant_monkey.aws_instance_id = instance_id
+        relevant_monkey.save()
+        logger.debug("Updated Monkey {} with aws instance id {}".format(str(relevant_monkey), instance_id))

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/environment.py

@@ -0,0 +1,12 @@
+import logging
+
+from monkey_island.cc.models.monkey import Monkey
+
+logger = logging.getLogger(__name__)
+
+
+def process_environment_telemetry(collector_results, monkey_guid):
+    relevant_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
+    relevant_monkey.environment = collector_results["environment"]
+    relevant_monkey.save()
+    logger.debug("Updated Monkey {} with env {}".format(str(relevant_monkey), collector_results))

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/hostname.py

@@ -0,0 +1,9 @@
+import logging
+
+from monkey_island.cc.models.monkey import Monkey
+
+logger = logging.getLogger(__name__)
+
+
+def process_hostname_telemetry(collector_results, monkey_guid):
+    Monkey.get_single_monkey_by_guid(monkey_guid).set_hostname(collector_results["hostname"])

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/system_info_telemetry_dispatcher.py

@@ -0,0 +1,63 @@
+import logging
+import typing
+
+from common.data.system_info_collectors_names import *
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.aws import process_aws_telemetry
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.environment import process_environment_telemetry
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.hostname import process_hostname_telemetry
+from monkey_island.cc.services.telemetry.zero_trust_tests.antivirus_existence import test_antivirus_existence
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSORS = {
+    AWS_COLLECTOR: [process_aws_telemetry],
+    ENVIRONMENT_COLLECTOR: [process_environment_telemetry],
+    HOSTNAME_COLLECTOR: [process_hostname_telemetry],
+    PROCESS_LIST_COLLECTOR: [test_antivirus_existence]
+}
+
+
+class SystemInfoTelemetryDispatcher(object):
+    def __init__(self, collector_to_parsing_functions: typing.Mapping[str, typing.List[typing.Callable]] = None):
+        """
+        :param collector_to_parsing_functions: Map between collector names and a list of functions that process the output of
+        that collector. If `None` is supplied, uses the default one; This should be the normal flow, overriding the
+        collector->functions mapping is useful mostly for testing.
+        """
+        if collector_to_parsing_functions is None:
+            collector_to_parsing_functions = SYSTEM_INFO_COLLECTOR_TO_TELEMETRY_PROCESSORS
+        self.collector_to_processing_functions = collector_to_parsing_functions
+
+    def dispatch_collector_results_to_relevant_processors(self, telemetry_json):
+        """
+        If the telemetry has collectors' results, dispatches the results to the relevant processing functions.
+        :param telemetry_json: Telemetry sent from the Monkey
+        """
+        if "collectors" in telemetry_json["data"]:
+            self.dispatch_single_result_to_relevant_processor(telemetry_json)
+
+    def dispatch_single_result_to_relevant_processor(self, telemetry_json):
+        relevant_monkey_guid = telemetry_json['monkey_guid']
+
+        for collector_name, collector_results in telemetry_json["data"]["collectors"].items():
+            self.dispatch_result_of_single_collector_to_processing_functions(
+                collector_name,
+                collector_results,
+                relevant_monkey_guid)
+
+    def dispatch_result_of_single_collector_to_processing_functions(
+            self,
+            collector_name,
+            collector_results,
+            relevant_monkey_guid):
+        if collector_name in self.collector_to_processing_functions:
+            for processing_function in self.collector_to_processing_functions[collector_name]:
+                # noinspection PyBroadException
+                try:
+                    processing_function(collector_results, relevant_monkey_guid)
+                except Exception as e:
+                    logger.error(
+                        "Error {} while processing {} system info telemetry".format(str(e), collector_name),
+                        exc_info=True)
+        else:
+            logger.warning("Unknown system info collector name: {}".format(collector_name))

monkey/monkey_island/cc/services/telemetry/processing/system_info_collectors/test_system_info_telemetry_dispatcher.py

@@ -0,0 +1,57 @@
+import uuid
+
+from monkey_island.cc.models import Monkey
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \
+    SystemInfoTelemetryDispatcher
+from monkey_island.cc.services.telemetry.processing.system_info_collectors.system_info_telemetry_dispatcher import \
+    process_aws_telemetry
+from monkey_island.cc.testing.IslandTestCase import IslandTestCase
+
+TEST_SYS_INFO_TO_PROCESSING = {
+    "AwsCollector": [process_aws_telemetry],
+}
+
+
+class SystemInfoTelemetryDispatcherTest(IslandTestCase):
+    def test_dispatch_to_relevant_collector_bad_inputs(self):
+        self.fail_if_not_testing_env()
+
+        dispatcher = SystemInfoTelemetryDispatcher(TEST_SYS_INFO_TO_PROCESSING)
+
+        # Bad format telem JSONs - throws
+        bad_empty_telem_json = {}
+        self.assertRaises(KeyError, dispatcher.dispatch_collector_results_to_relevant_processors, bad_empty_telem_json)
+        bad_no_data_telem_json = {"monkey_guid": "bla"}
+        self.assertRaises(KeyError, dispatcher.dispatch_collector_results_to_relevant_processors, bad_no_data_telem_json)
+        bad_no_monkey_telem_json = {"data": {"collectors": {"AwsCollector": "Bla"}}}
+        self.assertRaises(KeyError, dispatcher.dispatch_collector_results_to_relevant_processors, bad_no_monkey_telem_json)
+
+        # Telem JSON with no collectors - nothing gets dispatched
+        good_telem_no_collectors = {"monkey_guid": "bla", "data": {"bla": "bla"}}
+        good_telem_empty_collectors = {"monkey_guid": "bla", "data": {"bla": "bla", "collectors": {}}}
+
+        dispatcher.dispatch_collector_results_to_relevant_processors(good_telem_no_collectors)
+        dispatcher.dispatch_collector_results_to_relevant_processors(good_telem_empty_collectors)
+
+    def test_dispatch_to_relevant_collector(self):
+        self.fail_if_not_testing_env()
+        self.clean_monkey_db()
+
+        a_monkey = Monkey(guid=str(uuid.uuid4()))
+        a_monkey.save()
+
+        dispatcher = SystemInfoTelemetryDispatcher()
+
+        # JSON with results - make sure functions are called
+        instance_id = "i-0bd2c14bd4c7d703f"
+        telem_json = {
+            "data": {
+                "collectors": {
+                    "AwsCollector": {"instance_id": instance_id},
+                }
+            },
+            "monkey_guid": a_monkey.guid
+        }
+        dispatcher.dispatch_collector_results_to_relevant_processors(telem_json)
+
+        self.assertEquals(Monkey.get_single_monkey_by_guid(a_monkey.guid).aws_instance_id, instance_id)

monkey/monkey_island/cc/services/telemetry/zero_trust_tests/antivirus_existence.py

@@ -7,16 +7,16 @@ from monkey_island.cc.models.zero_trust.event import Event
 from monkey_island.cc.services.telemetry.zero_trust_tests.known_anti_viruses import ANTI_VIRUS_KNOWN_PROCESS_NAMES
 
 
-def test_antivirus_existence(telemetry_json):
+def test_antivirus_existence(process_list_json, monkey_guid):
-    current_monkey = Monkey.get_single_monkey_by_guid(telemetry_json['monkey_guid'])
+    current_monkey = Monkey.get_single_monkey_by_guid(monkey_guid)
-    if 'process_list' in telemetry_json['data']:
+
     process_list_event = Event.create_event(
         title="Process list",
         message="Monkey on {} scanned the process list".format(current_monkey.hostname),
         event_type=zero_trust_consts.EVENT_TYPE_MONKEY_LOCAL)
     events = [process_list_event]
 
-        av_processes = filter_av_processes(telemetry_json)
+    av_processes = filter_av_processes(process_list_json["process_list"])
 
     for process in av_processes:
         events.append(Event.create_event(
@@ -35,8 +35,8 @@ def test_antivirus_existence(telemetry_json):
     )
 
 
-def filter_av_processes(telemetry_json):
+def filter_av_processes(process_list):
-    all_processes = list(telemetry_json['data']['process_list'].items())
+    all_processes = list(process_list.items())
     av_processes = []
     for process in all_processes:
         process_name = process[1]['name']