Using the generic CloudInstance class to determine environment

monkey/common/cloud/aws/aws_instance.py

@@ -6,6 +6,7 @@ import logging
 
 __author__ = 'itay.mizeretz'
 
+from common.cloud.environment_names import AWS
 from common.cloud.instance import CloudInstance
 
 AWS_INSTANCE_METADATA_LOCAL_IP_ADDRESS = "169.254.169.254"
@@ -19,6 +20,11 @@ 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) -> str:
+        return AWS
 
     def __init__(self):
         self.instance_id = None
@@ -59,9 +65,6 @@ class AwsInstance(CloudInstance):
     def get_region(self):
         return self.region
 
-    def is_instance(self):
-        return self.instance_id is not None
-
     @staticmethod
     def _extract_account_id(instance_identity_document_response):
         """

monkey/common/cloud/azure/azure_instance.py

@@ -1,6 +1,7 @@
 import logging
 import requests
 
+from common.cloud.environment_names import AZURE
 from common.cloud.instance import CloudInstance
 
 LATEST_AZURE_METADATA_API_VERSION = "2019-04-30"
@@ -14,6 +15,11 @@ 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) -> str:
+        return AZURE
 
     def __init__(self):
         """
@@ -47,6 +53,3 @@ class AzureInstance(CloudInstance):
             self.location = response_data["compute"]["location"]
         except KeyError:
             logger.exception("Error while parsing response from Azure metadata service.")
-
-    def is_instance(self):
-        return self.on_azure

monkey/common/cloud/instance.py

@@ -1,3 +1,18 @@
+from typing import List
+
+from common.cloud.aws.aws_instance import AwsInstance
+from common.cloud.azure.azure_instance import AzureInstance
+
+
 class CloudInstance(object):
     def is_instance(self) -> bool:
         raise NotImplementedError()
+
+    def get_cloud_provider_name(self) -> str:
+        raise NotImplementedError()
+
+    all_cloud_instances = [AwsInstance(), AzureInstance()]
+
+    @staticmethod
+    def get_all_cloud_instances() -> List['CloudInstance']:
+        return CloudInstance.all_cloud_instances

monkey/infection_monkey/system_info/collectors/environment_collector.py

@@ -1,18 +1,15 @@
 from common.cloud.aws.aws_instance import AwsInstance
 from common.cloud.azure.azure_instance import AzureInstance
 from common.cloud.environment_names import ON_PREMISE, AZURE, AWS
+from common.cloud.instance import CloudInstance
 from infection_monkey.system_info.system_info_collector import SystemInfoCollector
 
 
-def get_monkey_environment():
+def get_monkey_environment() -> str:
-    # Check if on any cloud env. Default is on prem.
+    env = ON_PREMISE
-    if AwsInstance().is_instance():
+    for instance in CloudInstance.get_all_cloud_instances():
-        env = AWS
+        if instance.is_instance():
-    elif AzureInstance().is_instance():
+            env = instance.get_cloud_provider_name()
-        env = AZURE
-    # TODO: elif GcpInstance().is_gcp_instance():
-    else:
-        env = ON_PREMISE
     return env