From 676d46307bf4c761bf59f56226e53dc314022155 Mon Sep 17 00:00:00 2001
From: Shay Nehmad <shay.nehmad@guardicore.com>
Date: Wed, 8 Jan 2020 11:20:49 +0200
Subject: [PATCH] Using the generic CloudInstance class to determine
 environment

---
 monkey/common/cloud/aws/aws_instance.py           |  9 ++++++---
 monkey/common/cloud/azure/azure_instance.py       |  9 ++++++---
 monkey/common/cloud/instance.py                   | 15 +++++++++++++++
 .../collectors/environment_collector.py           | 15 ++++++---------
 4 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/monkey/common/cloud/aws/aws_instance.py b/monkey/common/cloud/aws/aws_instance.py
index 301881894..c77220d17 100644
--- a/monkey/common/cloud/aws/aws_instance.py
+++ b/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):
         """
diff --git a/monkey/common/cloud/azure/azure_instance.py b/monkey/common/cloud/azure/azure_instance.py
index 5222c7620..f0d5a8044 100644
--- a/monkey/common/cloud/azure/azure_instance.py
+++ b/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
diff --git a/monkey/common/cloud/instance.py b/monkey/common/cloud/instance.py
index 52dd56b02..2e702b867 100644
--- a/monkey/common/cloud/instance.py
+++ b/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
diff --git a/monkey/infection_monkey/system_info/collectors/environment_collector.py b/monkey/infection_monkey/system_info/collectors/environment_collector.py
index 208bbfa42..56df5906b 100644
--- a/monkey/infection_monkey/system_info/collectors/environment_collector.py
+++ b/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():
-    # Check if on any cloud env. Default is on prem.
-    if AwsInstance().is_instance():
-        env = AWS
-    elif AzureInstance().is_instance():
-        env = AZURE
-    # TODO: elif GcpInstance().is_gcp_instance():
-    else:
-        env = ON_PREMISE
+def get_monkey_environment() -> str:
+    env = ON_PREMISE
+    for instance in CloudInstance.get_all_cloud_instances():
+        if instance.is_instance():
+            env = instance.get_cloud_provider_name()
     return env