Extracted function in EnvCollector for reuse in other parts of the Monkey

This commit is contained in:
Shay Nehmad 2020-01-05 16:23:22 +02:00
parent c0331f84ff
commit fdb54f6b8d
1 changed files with 13 additions and 10 deletions

View File

@ -4,11 +4,7 @@ from common.cloud.environment_names import ON_PREMISE, AZURE, AWS
from infection_monkey.system_info.system_info_collector import SystemInfoCollector from infection_monkey.system_info.system_info_collector import SystemInfoCollector
class EnvironmentCollector(SystemInfoCollector): def get_monkey_environment():
def __init__(self):
super(EnvironmentCollector, self).__init__(name="EnvironmentCollector")
def collect(self) -> dict:
# Check if on any cloud env. Default is on prem. # Check if on any cloud env. Default is on prem.
if AwsInstance().is_aws_instance(): if AwsInstance().is_aws_instance():
env = AWS env = AWS
@ -17,5 +13,12 @@ class EnvironmentCollector(SystemInfoCollector):
# TODO: elif GcpInstance().is_gcp_instance(): # TODO: elif GcpInstance().is_gcp_instance():
else: else:
env = ON_PREMISE env = ON_PREMISE
return env
return {"environment": env}
class EnvironmentCollector(SystemInfoCollector):
def __init__(self):
super(EnvironmentCollector, self).__init__(name="EnvironmentCollector")
def collect(self) -> dict:
return {"environment": get_monkey_environment()}