WIP - adding the environment collector as a plugin, still some work to do

monkey/common/cloud/environment_names.py

@@ -0,0 +1,3 @@
+ON_PREMISE = "On Premise"
+AZURE = "Azure"
+AWS = "AWS"

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}

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

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