1960 deserialize config

monkey/common/configuration/__init__.py

@@ -1,6 +1,15 @@
-from .agent_configuration import (
-    AgentConfiguration,
-    AgentConfigurationSchema,
+from .agent_configuration import AgentConfiguration, InvalidConfigurationError
+from .agent_sub_configurations import (
+    CustomPBAConfiguration,
+    PluginConfiguration,
+    ScanTargetConfiguration,
+    ICMPScanConfiguration,
+    TCPScanConfiguration,
+    NetworkScanConfiguration,
+    ExploitationOptionsConfiguration,
+    ExploiterConfiguration,
+    ExploitationConfiguration,
+    PropagationConfiguration,
 )
 from .default_agent_configuration import (
     DEFAULT_AGENT_CONFIGURATION_JSON,

monkey/common/configuration/agent_configuration.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import List
+from typing import Any, List, Mapping
 
-from marshmallow import Schema, fields, post_load
+from marshmallow import Schema, fields
+from marshmallow.exceptions import MarshmallowError
 
 from .agent_sub_configuration_schemas import (
     CustomPBAConfigurationSchema,
@@ -15,6 +18,15 @@
 )
 
 
+class InvalidConfigurationError(Exception):
+    pass
+
+
+INVALID_CONFIGURATION_ERROR_MESSAGE = (
+    "Cannot construct an AgentConfiguration object with the supplied, invalid data:"
+)
+
+
 @dataclass(frozen=True)
 class AgentConfiguration:
     keep_tunnel_open_time: float
@@ -24,6 +36,57 @@ class AgentConfiguration:
     payloads: List[PluginConfiguration]
     propagation: PropagationConfiguration
 
+    def __post_init__(self):
+        # This will raise an exception if the object is invalid. Calling this in __post__init()
+        # makes it impossible to construct an invalid object
+        try:
+            AgentConfigurationSchema().dump(self)
+        except Exception as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
+
+    @staticmethod
+    def from_mapping(config_mapping: Mapping[str, Any]) -> AgentConfiguration:
+        """
+        Construct an AgentConfiguration from a Mapping
+
+        :param config_mapping: A Mapping that represents an AgentConfiguration
+        :return: An AgentConfiguration
+        :raises: InvalidConfigurationError if the provided Mapping does not represent a valid
+                 AgentConfiguration
+        """
+
+        try:
+            config_dict = AgentConfigurationSchema().load(config_mapping)
+            return AgentConfiguration(**config_dict)
+        except MarshmallowError as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
+
+    @staticmethod
+    def from_json(config_json: str) -> AgentConfiguration:
+        """
+        Construct an AgentConfiguration from a JSON string
+
+        :param config_json: A JSON string that represents an AgentConfiguration
+        :return: An AgentConfiguration
+        :raises: InvalidConfigurationError if the provided JSON does not represent a valid
+                 AgentConfiguration
+        """
+        try:
+            config_dict = AgentConfigurationSchema().loads(config_json)
+            return AgentConfiguration(**config_dict)
+        except MarshmallowError as err:
+            raise InvalidConfigurationError(f"{INVALID_CONFIGURATION_ERROR_MESSAGE}: {err}")
+
+    @staticmethod
+    def to_json(config: AgentConfiguration) -> str:
+        """
+        Serialize an AgentConfiguration to JSON
+
+        :param config: An AgentConfiguration
+        :return: A JSON string representing the AgentConfiguration
+        """
+        return AgentConfigurationSchema().dumps(config)
+
 
 class AgentConfigurationSchema(Schema):
     keep_tunnel_open_time = fields.Float()
@@ -32,7 +95,3 @@ class AgentConfigurationSchema(Schema):
     credential_collectors = fields.List(fields.Nested(PluginConfigurationSchema))
     payloads = fields.List(fields.Nested(PluginConfigurationSchema))
     propagation = fields.Nested(PropagationConfigurationSchema)
-
-    @post_load
-    def _make_agent_configuration(self, data, **kwargs):
-        return AgentConfiguration(**data)

monkey/common/configuration/default_agent_configuration.py

@@ -1,4 +1,4 @@
-from . import AgentConfiguration, AgentConfigurationSchema
+from . import AgentConfiguration
 
 DEFAULT_AGENT_CONFIGURATION_JSON = """{
         "keep_tunnel_open_time": 30,
@@ -204,5 +204,4 @@
 
 
 def build_default_agent_configuration() -> AgentConfiguration:
-    schema = AgentConfigurationSchema()
-    return schema.loads(DEFAULT_AGENT_CONFIGURATION_JSON)
+    return AgentConfiguration.from_json(DEFAULT_AGENT_CONFIGURATION_JSON)

monkey/common/operating_systems.py

@@ -2,5 +2,12 @@
 
 
 class OperatingSystems(Enum):
+    """
+    An Enum representing all supported operating systems
+
+    This Enum represents all operating systems that Infection Monkey supports. The value of each
+    member is the member's name in all lower-case characters.
+    """
+
     LINUX = "linux"
     WINDOWS = "windows"

monkey/common/utils/exceptions.py

@@ -38,5 +38,7 @@ class DomainControllerNameFetchError(FailedExploitationError):
     """Raise on failed attempt to extract domain controller's name"""
 
 
+# TODO: This has been replaced by common.configuration.InvalidConfigurationError. Use that error
+#       instead and remove this one.
 class InvalidConfigurationError(Exception):
     """Raise when configuration is invalid"""

monkey/infection_monkey/exploit/caching_agent_repository.py

@@ -5,6 +5,7 @@
 
 import requests
 
+from common import OperatingSystems
 from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT
 
 from . import IAgentRepository
@@ -22,18 +23,22 @@ def __init__(self, island_url: str, proxies: Mapping[str, str]):
         self._proxies = proxies
         self._lock = threading.Lock()
 
-    def get_agent_binary(self, os: str, architecture: str = None) -> io.BytesIO:
+    def get_agent_binary(
+        self, operating_system: OperatingSystems, architecture: str = None
+    ) -> io.BytesIO:
         # If multiple calls to get_agent_binary() are made simultaneously before the result of
         # _download_binary_from_island() is cached, then multiple requests will be sent to the
         # island. Add a mutex in front of the call to _download_agent_binary_from_island() so
         # that only one request per OS will be sent to the island.
         with self._lock:
-            return io.BytesIO(self._download_binary_from_island(os))
+            return io.BytesIO(self._download_binary_from_island(operating_system))
 
     @lru_cache(maxsize=None)
-    def _download_binary_from_island(self, os: str) -> bytes:
+    def _download_binary_from_island(self, operating_system: OperatingSystems) -> bytes:
+        os_name = operating_system.value
+
         response = requests.get(  # noqa: DUO123
-            f"{self._island_url}/api/agent-binaries/{os}",
+            f"{self._island_url}/api/agent-binaries/{os_name}",
             verify=False,
             proxies=self._proxies,
             timeout=MEDIUM_REQUEST_TIMEOUT,

monkey/infection_monkey/exploit/i_agent_repository.py

@@ -1,6 +1,8 @@
 import abc
 import io
 
+from common import OperatingSystems
+
 # TODO: The Island also has an IAgentRepository with a totally different interface. At the moment,
 #       the Island and Agent have different needs, but at some point we should unify these.
 
@@ -13,12 +15,13 @@ class IAgentRepository(metaclass=abc.ABCMeta):
     """
 
     @abc.abstractmethod
-    def get_agent_binary(self, os: str, architecture: str = None) -> io.BytesIO:
+    def get_agent_binary(
+        self, operating_system: OperatingSystems, architecture: str = None
+    ) -> io.BytesIO:
         """
         Retrieve the appropriate agent binary from the repository.
-        :param str os: The name of the operating system on which the agent binary will run
-        :param str architecture: Reserved
+        :param operating_system: The name of the operating system on which the agent binary will run
+        :param architecture: Reserved
         :return: A file-like object for the requested agent binary
-        :rtype: io.BytesIO
         """
         pass

monkey/infection_monkey/exploit/log4shell.py

@@ -129,7 +129,7 @@ def _build_command(self, path: PurePath, http_path) -> str:
         }
 
     def _build_java_class(self, exploit_command: str) -> bytes:
-        if OperatingSystems.LINUX in self.host.os["type"]:
+        if OperatingSystems.LINUX == self.host.os["type"]:
             return build_exploit_bytecode(exploit_command, LINUX_EXPLOIT_TEMPLATE_PATH)
         else:
             return build_exploit_bytecode(exploit_command, WINDOWS_EXPLOIT_TEMPLATE_PATH)

monkey/infection_monkey/exploit/powershell.py

@@ -2,6 +2,7 @@
 from pathlib import Path, PurePath
 from typing import List, Optional
 
+from common import OperatingSystems
 from infection_monkey.exploit.HostExploiter import HostExploiter
 from infection_monkey.exploit.powershell_utils.auth_options import AuthOptions, get_auth_options
 from infection_monkey.exploit.powershell_utils.credentials import (
@@ -162,7 +163,7 @@ def _copy_monkey_binary_to_victim(self, monkey_path_on_victim: PurePath):
                 temp_monkey_binary_filepath.unlink()
 
     def _create_local_agent_file(self, binary_path):
-        agent_binary_bytes = self.agent_repository.get_agent_binary("windows")
+        agent_binary_bytes = self.agent_repository.get_agent_binary(OperatingSystems.WINDOWS)
         with open(binary_path, "wb") as f:
             f.write(agent_binary_bytes.getvalue())
 

monkey/infection_monkey/exploit/tools/http_tools.py

@@ -57,6 +57,6 @@ def create_locked_transfer(
         httpd.start()
         lock.acquire()
         return (
-            "http://%s:%s/%s" % (local_ip, local_port, urllib.parse.quote(host.os["type"])),
+            "http://%s:%s/%s" % (local_ip, local_port, urllib.parse.quote(host.os["type"].value)),
             httpd,
         )

monkey/infection_monkey/i_control_channel.py

@@ -1,5 +1,7 @@
 import abc
 
+from common.configuration import AgentConfiguration
+
 
 class IControlChannel(metaclass=abc.ABCMeta):
     @abc.abstractmethod
@@ -11,10 +13,10 @@ def should_agent_stop(self) -> bool:
         """
 
     @abc.abstractmethod
-    def get_config(self) -> dict:
+    def get_config(self) -> AgentConfiguration:
         """
-        :return: A dictionary containing Agent Configuration
-        :rtype: dict
+        :return: An AgentConfiguration object
+        :rtype: AgentConfiguration
         """
         pass
 

monkey/infection_monkey/master/automated_master.py

@@ -1,8 +1,9 @@
 import logging
 import threading
 import time
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Optional
 
+from common.configuration import PluginConfiguration
 from common.utils import Timer
 from infection_monkey.credential_store import ICredentialsStore
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
@@ -13,7 +14,7 @@
 from infection_monkey.telemetry.credentials_telem import CredentialsTelem
 from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
 from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
-from infection_monkey.utils.propagation import should_propagate
+from infection_monkey.utils.propagation import maximum_depth_reached
 from infection_monkey.utils.threading import create_daemon_thread, interruptible_iter
 
 from . import Exploiter, IPScanner, Propagator
@@ -111,7 +112,7 @@ def _wait_for_master_stop_condition(self):
             time.sleep(CHECK_FOR_TERMINATE_INTERVAL_SEC)
 
     @staticmethod
-    def _try_communicate_with_island(fn: Callable[[], Any], max_tries: int):
+    def _try_communicate_with_island(fn: Callable[[], Any], max_tries: int) -> Any:
         tries = 0
         while tries < max_tries:
             try:
@@ -141,7 +142,7 @@ def _run_simulation(self):
         try:
             config = AutomatedMaster._try_communicate_with_island(
                 self._control_channel.get_config, CHECK_FOR_CONFIG_COUNT
-            )["config"]
+            )
         except IslandCommunicationError as e:
             logger.error(f"An error occurred while fetching configuration: {e}")
             return
@@ -150,15 +151,15 @@ def _run_simulation(self):
             target=self._run_plugins,
             name="CredentialCollectorThread",
             args=(
-                config["credential_collectors"],
+                config.credential_collectors,
                 "credential collector",
                 self._collect_credentials,
             ),
         )
         pba_thread = create_daemon_thread(
             target=self._run_pbas,
             name="PBAThread",
-            args=(config["post_breach_actions"].items(), self._run_pba, config["custom_pbas"]),
+            args=(config.post_breach_actions, self._run_pba, config.custom_pbas),
         )
 
         credential_collector_thread.start()
@@ -173,52 +174,56 @@ def _run_simulation(self):
         current_depth = self._current_depth if self._current_depth is not None else 0
         logger.info(f"Current depth is {current_depth}")
 
-        if should_propagate(self._control_channel.get_config(), self._current_depth):
-            self._propagator.propagate(config["propagation"], current_depth, self._stop)
+        if maximum_depth_reached(config.propagation.maximum_depth, self._current_depth):
+            self._propagator.propagate(config.propagation, current_depth, self._stop)
         else:
             logger.info("Skipping propagation: maximum depth reached")
 
         payload_thread = create_daemon_thread(
             target=self._run_plugins,
             name="PayloadThread",
-            args=(config["payloads"].items(), "payload", self._run_payload),
+            args=(config.payloads, "payload", self._run_payload),
         )
         payload_thread.start()
         payload_thread.join()
 
         pba_thread.join()
 
-    def _collect_credentials(self, collector: str):
-        credentials = self._puppet.run_credential_collector(collector, {})
+    def _collect_credentials(self, collector: PluginConfiguration):
+        credentials = self._puppet.run_credential_collector(collector.name, collector.options)
 
         if credentials:
             self._telemetry_messenger.send_telemetry(CredentialsTelem(credentials))
         else:
             logger.debug(f"No credentials were collected by {collector}")
 
-    def _run_pba(self, pba: Tuple[str, Dict]):
-        name = pba[0]
-        options = pba[1]
-
-        for pba_data in self._puppet.run_pba(name, options):
+    def _run_pba(self, pba: PluginConfiguration):
+        for pba_data in self._puppet.run_pba(pba.name, pba.options):
             self._telemetry_messenger.send_telemetry(PostBreachTelem(pba_data))
 
-    def _run_payload(self, payload: Tuple[str, Dict]):
-        name = payload[0]
-        options = payload[1]
-
-        self._puppet.run_payload(name, options, self._stop)
+    def _run_payload(self, payload: PluginConfiguration):
+        self._puppet.run_payload(payload.name, payload.options, self._stop)
 
     def _run_pbas(
-        self, plugins: Iterable[Any], callback: Callable[[Any], None], custom_pba_options: Mapping
+        self,
+        plugins: Iterable[PluginConfiguration],
+        callback: Callable[[Any], None],
+        custom_pba_options: Dict,
     ):
         self._run_plugins(plugins, "post-breach action", callback)
 
         if custom_pba_is_enabled(custom_pba_options):
-            self._run_plugins([("CustomPBA", custom_pba_options)], "post-breach action", callback)
+            self._run_plugins(
+                [PluginConfiguration(name="CustomPBA", options=custom_pba_options)],
+                "post-breach action",
+                callback,
+            )
 
     def _run_plugins(
-        self, plugins: Iterable[Any], plugin_type: str, callback: Callable[[Any], None]
+        self,
+        plugins: Iterable[PluginConfiguration],
+        plugin_type: str,
+        callback: Callable[[Any], None],
     ):
         logger.info(f"Running {plugin_type}s")
         logger.debug(f"Found {len(plugins)} {plugin_type}(s) to run")