diff --git a/monkey/infection_monkey/system_info/collectors/environment_collector.py b/monkey/infection_monkey/system_info/collectors/environment_collector.py index 1459b5633..523989f61 100644 --- a/monkey/infection_monkey/system_info/collectors/environment_collector.py +++ b/monkey/infection_monkey/system_info/collectors/environment_collector.py @@ -4,18 +4,21 @@ from common.cloud.environment_names import ON_PREMISE, AZURE, AWS 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_aws_instance(): + env = AWS + elif AzureInstance().is_azure_instance(): + env = AZURE + # TODO: elif GcpInstance().is_gcp_instance(): + else: + env = ON_PREMISE + return env + + 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} + return {"environment": get_monkey_environment()}