From fdb54f6b8d9b3269a361b65cf1bac61e4e9897e8 Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Sun, 5 Jan 2020 16:23:22 +0200 Subject: [PATCH] Extracted function in EnvCollector for reuse in other parts of the Monkey --- .../collectors/environment_collector.py | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) 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()}