diff --git a/monkey/common/cloud/aws/aws_instance.py b/monkey/common/cloud/aws/aws_instance.py index c77220d17..03c5482ba 100644 --- a/monkey/common/cloud/aws/aws_instance.py +++ b/monkey/common/cloud/aws/aws_instance.py @@ -6,7 +6,7 @@ import logging __author__ = 'itay.mizeretz' -from common.cloud.environment_names import AWS +from common.cloud.environment_names import Environment from common.cloud.instance import CloudInstance AWS_INSTANCE_METADATA_LOCAL_IP_ADDRESS = "169.254.169.254" @@ -23,8 +23,8 @@ class AwsInstance(CloudInstance): def is_instance(self): return self.instance_id is not None - def get_cloud_provider_name(self) -> str: - return AWS + def get_cloud_provider_name(self) -> Environment: + return Environment.AWS def __init__(self): self.instance_id = None diff --git a/monkey/common/cloud/azure/azure_instance.py b/monkey/common/cloud/azure/azure_instance.py index f0d5a8044..ec910fb98 100644 --- a/monkey/common/cloud/azure/azure_instance.py +++ b/monkey/common/cloud/azure/azure_instance.py @@ -1,7 +1,7 @@ import logging import requests -from common.cloud.environment_names import AZURE +from common.cloud.environment_names import Environment from common.cloud.instance import CloudInstance LATEST_AZURE_METADATA_API_VERSION = "2019-04-30" @@ -18,8 +18,8 @@ class AzureInstance(CloudInstance): def is_instance(self): return self.on_azure - def get_cloud_provider_name(self) -> str: - return AZURE + def get_cloud_provider_name(self) -> Environment: + return Environment.AZURE def __init__(self): """ diff --git a/monkey/common/cloud/environment_names.py b/monkey/common/cloud/environment_names.py index 0c8655753..945d438ce 100644 --- a/monkey/common/cloud/environment_names.py +++ b/monkey/common/cloud/environment_names.py @@ -1,12 +1,15 @@ -# When adding a new environment to this file, make sure to add it to ALL_ENV_NAMES as well! +from enum import Enum -UNKNOWN = "Unknown" -ON_PREMISE = "On Premise" -AZURE = "Azure" -AWS = "AWS" -GCP = "GCP" -ALIBABA = "Alibaba Cloud" -IBM = "IBM Cloud" -DigitalOcean = "Digital Ocean" -ALL_ENV_NAMES = [UNKNOWN, ON_PREMISE, AZURE, AWS, GCP, ALIBABA, IBM, DigitalOcean] +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] diff --git a/monkey/common/cloud/gcp/gcp_instance.py b/monkey/common/cloud/gcp/gcp_instance.py index 26738db43..184465bf5 100644 --- a/monkey/common/cloud/gcp/gcp_instance.py +++ b/monkey/common/cloud/gcp/gcp_instance.py @@ -1,7 +1,7 @@ import logging import requests -from common.cloud.environment_names import GCP +from common.cloud.environment_names import Environment from common.cloud.instance import CloudInstance logger = logging.getLogger(__name__) @@ -17,8 +17,8 @@ class GcpInstance(CloudInstance): def is_instance(self): return self.on_gcp - def get_cloud_provider_name(self) -> str: - return GCP + def get_cloud_provider_name(self) -> Environment: + return Environment.GCP def __init__(self): self.on_gcp = False diff --git a/monkey/common/cloud/instance.py b/monkey/common/cloud/instance.py index 61ab4c734..abe0c7910 100644 --- a/monkey/common/cloud/instance.py +++ b/monkey/common/cloud/instance.py @@ -1,3 +1,6 @@ +from common.cloud.environment_names import Environment + + class CloudInstance(object): """ This is an abstract class which represents a cloud instance. @@ -7,5 +10,5 @@ class CloudInstance(object): def is_instance(self) -> bool: raise NotImplementedError() - def get_cloud_provider_name(self) -> str: + def get_cloud_provider_name(self) -> Environment: raise NotImplementedError() diff --git a/monkey/infection_monkey/system_info/collectors/environment_collector.py b/monkey/infection_monkey/system_info/collectors/environment_collector.py index 12a58b6d5..7a953fce9 100644 --- a/monkey/infection_monkey/system_info/collectors/environment_collector.py +++ b/monkey/infection_monkey/system_info/collectors/environment_collector.py @@ -1,10 +1,10 @@ from common.cloud.all_instances import get_all_cloud_instances -from common.cloud.environment_names import ON_PREMISE +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() -> str: +def get_monkey_environment() -> Environment: """ Get the Monkey's running environment. :return: One of the cloud providers if on cloud; otherwise, assumes "on premise". @@ -13,7 +13,7 @@ def get_monkey_environment() -> str: if instance.is_instance(): return instance.get_cloud_provider_name() - return ON_PREMISE + return Environment.ON_PREMISE class EnvironmentCollector(SystemInfoCollector): diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index 65dda0850..da6d880b4 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -45,7 +45,7 @@ class Monkey(Document): command_control_channel = EmbeddedDocumentField(CommandControlChannel) # Environment related fields - environment = StringField(default=environment_names.UNKNOWN, choices=environment_names.ALL_ENV_NAMES) + 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.