diff --git a/monkey/common/cloud/environment_names.py b/monkey/common/cloud/environment_names.py new file mode 100644 index 000000000..f9e881a5a --- /dev/null +++ b/monkey/common/cloud/environment_names.py @@ -0,0 +1,3 @@ +ON_PREMISE = "On Premise" +AZURE = "Azure" +AWS = "AWS" diff --git a/monkey/infection_monkey/system_info/collectors/__init__.py b/monkey/infection_monkey/system_info/collectors/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/infection_monkey/system_info/collectors/environment_collector.py b/monkey/infection_monkey/system_info/collectors/environment_collector.py new file mode 100644 index 000000000..1459b5633 --- /dev/null +++ b/monkey/infection_monkey/system_info/collectors/environment_collector.py @@ -0,0 +1,21 @@ +from common.cloud.aws_instance import AwsInstance +from common.cloud.azure.azure_instance import AzureInstance +from common.cloud.environment_names import ON_PREMISE, AZURE, AWS +from infection_monkey.system_info.system_info_collector import SystemInfoCollector + + +class EnvironmentCollector(SystemInfoCollector): + def __init__(self): + super(EnvironmentCollector, self).__init__(name="EnvironmentCollector") + + def collect(self) -> dict: + # Check if on any cloud env. Default is on prem. + if AwsInstance().is_aws_instance(): + env = AWS + elif AzureInstance().is_azure_instance(): + env = AZURE + # TODO: elif GcpInstance().is_gcp_instance(): + else: + env = ON_PREMISE + + return {"environment": env} diff --git a/monkey/infection_monkey/system_info/system_info_collector.py b/monkey/infection_monkey/system_info/system_info_collector.py new file mode 100644 index 000000000..863e0e0ab --- /dev/null +++ b/monkey/infection_monkey/system_info/system_info_collector.py @@ -0,0 +1,30 @@ +from config import WormConfiguration +from infection_monkey.utils.plugins.plugin import Plugin + + +import infection_monkey.system_info.collectors + + +class SystemInfoCollector(Plugin): + def __init__(self, name="unknown"): + self.name = name + + @staticmethod + def should_run(class_name) -> bool: + return class_name in WormConfiguration.system_info_collectors + + @staticmethod + def base_package_file(): + return infection_monkey.system_info.collectors.__file__ + + @staticmethod + def base_package_name(): + return infection_monkey.system_info.collectors.__package__ + + def collect(self) -> dict: + """ + Collect the relevant information and return it in a dictionary. + To be implemented by each collector. + TODO should this be an abstractmethod, or will that ruin the plugin system somehow? if can be abstract should add UT + """ + raise NotImplementedError() diff --git a/monkey/infection_monkey/system_info/system_info_handler.py b/monkey/infection_monkey/system_info/system_info_handler.py new file mode 100644 index 000000000..dd085300f --- /dev/null +++ b/monkey/infection_monkey/system_info/system_info_handler.py @@ -0,0 +1,34 @@ +import logging +from typing import Sequence + +from infection_monkey.system_info.system_info_collector import SystemInfoCollector +from telemetry.system_info_telem import SystemInfoTelem + +LOG = logging.getLogger(__name__) + +PATH_TO_COLLECTORS = "infection_monkey.system_info.collectors." + + +# TODO Add new collectors to config and config schema +class SystemInfo(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: + 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)) + # TODO Send SystemInfoTelem() + + @staticmethod + def config_to_collectors_list() -> Sequence[SystemInfoCollector]: + return SystemInfoCollector.get_instances()