diff --git a/app/config.py b/app/config.py index 43d6641..6844903 100644 --- a/app/config.py +++ b/app/config.py @@ -120,20 +120,20 @@ class MachineConfig(): def vagrantfilepath(self) -> str: """ Vagrant specific config: The vagrant file path """ - if "vagrantfilepath" not in self.raw_config["vm_controller"]: + if not self.raw_config.vm_controller.has_key("vagrantfilepath"): raise ConfigurationError("Vagrantfilepath missing") - return self.raw_config["vm_controller"]["vagrantfilepath"] + return self.raw_config.vm_controller.vagrantfilepath def sensors(self) -> list[str]: """ Return a list of sensors configured for this machine """ - if "sensors" in self.raw_config: - return self.raw_config["sensors"] or [] + if self.raw_config.has_key("sensors"): + return self.raw_config.sensors or [] return [] def vulnerabilities(self) -> list[str]: """ Return a list of vulnerabilities configured for this machine """ - if "vulnerabilities" in self.raw_config: - return self.raw_config["vulnerabilities"] or [] + if self.raw_config.has_key("vulnerabilities"): + return self.raw_config.vulnerabilities or [] return [] def is_active(self) -> bool: @@ -307,7 +307,7 @@ class ExperimentConfig(): raise ConfigurationError("Config file is empty") try: - return int(self.raw_config["attacks"]["nap_time"]) + return int(self.raw_config.attacks.nap_time) except KeyError: return 0 diff --git a/app/config_verifier.py b/app/config_verifier.py index 46d731c..edc3b26 100644 --- a/app/config_verifier.py +++ b/app/config_verifier.py @@ -39,6 +39,11 @@ class VMController: return True return False + # def __dict__(self): + # return {"vm_type": self.vm_type, + # "vagrantfilepath": self.vagrantfilepath, + # "ip": self.ip} + @dataclass class Attacker: @@ -55,6 +60,11 @@ class Attacker: return True return False + def get(self, keyname, default=None): + if self.has_key(keyname): + return self.__dict__[keyname] + return default + @dataclass class Target: @@ -81,6 +91,11 @@ class Target: return True return False + def get(self, keyname, default=None): + if self.has_key(keyname): + return self.__dict__[keyname] + return default + @dataclass class AttackConfig: diff --git a/plugins/base/machinery.py b/plugins/base/machinery.py index 301f616..c180036 100644 --- a/plugins/base/machinery.py +++ b/plugins/base/machinery.py @@ -142,7 +142,7 @@ class MachineryPlugin(BasePlugin): # print("===========> Processing config") self.config = config - self.process_config(config.raw_config) + self.process_config(config.raw_config.__dict__) def __call_remote_run__(self, cmd: str, disown: bool = False): """ Simplifies connect and run