Merge pull request #2276 from guardicore/2234-refactor-reset-agent-config

2234 refactor reset agent config

monkey/monkey_island/cc/island_event_handlers/__init__.py

@@ -0,0 +1 @@
+from .reset_agent_configuration import reset_agent_configuration

monkey/monkey_island/cc/island_event_handlers/reset_agent_configuration.py

@@ -0,0 +1,33 @@
+from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
+
+
+class reset_agent_configuration:
+    """
+    Callable class that handles removal of PBA files and reset of agent configuration
+    on the Island.
+    """
+
+    def __init__(
+        self,
+        agent_configuration_repository: IAgentConfigurationRepository,
+        file_repository: IFileRepository,
+    ):
+        self._agent_configuration_repository = agent_configuration_repository
+        self._file_repository = file_repository
+
+    def __call__(self):
+        self._remove_pba_files()
+        self._agent_configuration_repository.reset_to_default()
+
+    def _remove_pba_files(self):
+        """
+        Removes PBA files from the file repository
+        """
+        agent_configuration = self._agent_configuration_repository.get_configuration()
+        custom_pbas = agent_configuration.custom_pbas
+
+        if custom_pbas.linux_filename:
+            self._file_repository.delete_file(custom_pbas.linux_filename)
+
+        if custom_pbas.windows_filename:
+            self._file_repository.delete_file(custom_pbas.windows_filename)

monkey/monkey_island/cc/models/monkey.py

@@ -19,7 +19,7 @@ from mongoengine import (
 from monkey_island.cc.models.command_control_channel import CommandControlChannel
 from monkey_island.cc.models.monkey_ttl import MonkeyTtl, create_monkey_ttl_document
 from monkey_island.cc.server_utils.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
-from monkey_island.cc.services.utils.network_utils import get_ip_addresses
+from monkey_island.cc.server_utils.network_utils import get_ip_addresses
 
 
 class ParentNotFoundError(Exception):

monkey/monkey_island/cc/resources/reset_agent_configuration.py

@@ -2,22 +2,22 @@ from http import HTTPStatus
 
 from flask import make_response
 
+from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
-from monkey_island.cc.services import RepositoryService
 
 
 class ResetAgentConfiguration(AbstractResource):
     urls = ["/api/reset-agent-configuration"]
 
-    def __init__(self, repository_service: RepositoryService):
+    def __init__(self, island_event_queue: IIslandEventQueue):
-        self._repository_service = repository_service
+        self._island_event_queue = island_event_queue
 
     @jwt_required
     def post(self):
         """
         Reset the agent configuration to its default values
         """
-        self._repository_service.reset_agent_configuration()
+        self._island_event_queue.publish(IslandEventTopic.RESET_AGENT_CONFIGURATION)
 
         return make_response({}, HTTPStatus.OK)

monkey/monkey_island/cc/server_setup.py

@@ -33,10 +33,10 @@ from monkey_island.cc.server_utils.consts import (  # noqa: E402
     MONKEY_ISLAND_ABS_PATH,
 )
 from monkey_island.cc.server_utils.island_logger import reset_logger, setup_logging  # noqa: E402
+from monkey_island.cc.server_utils.network_utils import get_ip_addresses  # noqa: E402
 from monkey_island.cc.services.initialize import initialize_services  # noqa: E402
-from monkey_island.cc.services.utils.network_utils import get_ip_addresses  # noqa: E402
-from monkey_island.cc.setup import PyWSGILoggingFilter  # noqa: E402
 from monkey_island.cc.setup import island_config_options_validator  # noqa: E402
+from monkey_island.cc.setup import PyWSGILoggingFilter, setup_island_event_handlers  # noqa: E402
 from monkey_island.cc.setup.data_dir import IncompatibleDataDirectory, setup_data_dir  # noqa: E402
 from monkey_island.cc.setup.gevent_hub_error_handler import GeventHubErrorHandler  # noqa: E402
 from monkey_island.cc.setup.island_config_options import IslandConfigOptions  # noqa: E402
@@ -58,6 +58,7 @@ def run_monkey_island():
 
     _send_analytics(deployment, version)
     container = _initialize_di_container(ip_addresses, version, config_options.data_dir)
+    setup_island_event_handlers(container)
 
     _initialize_mongodb_connection(config_options.start_mongodb, config_options.data_dir)
     _start_island_server(ip_addresses, island_args.setup_only, config_options, container)

monkey/monkey_island/cc/server_utils/network_utils.py

@@ -0,0 +1,39 @@
+import ipaddress
+from typing import Sequence
+
+from netifaces import AF_INET, ifaddresses, interfaces
+from ring import lru
+
+# TODO: This functionality is duplicated in the agent. Unify them after 2216-tcp-relay is merged
+
+
+# The local IP addresses list should not change often. Therefore, we can cache the result and never
+# call this function more than once. This stopgap measure is here since this function is called a
+# lot of times during the report generation. This means that if the interfaces of the Island machine
+# change, the Island process needs to be restarted.
+@lru(maxsize=1)
+def get_ip_addresses() -> Sequence[str]:
+    ip_list = []
+    for interface in interfaces():
+        addresses = ifaddresses(interface).get(AF_INET, [])
+        ip_list.extend([link["addr"] for link in addresses if link["addr"] != "127.0.0.1"])
+    return ip_list
+
+
+# The subnets list should not change often. Therefore, we can cache the result and never call this
+# function more than once. This stopgap measure is here since this function is called a lot of times
+# during the report generation. This means that if the interfaces or subnets of the Island machine
+# change, the Island process needs to be restarted.
+@lru(maxsize=1)
+def get_subnets():
+    subnets = []
+    for interface in interfaces():
+        addresses = ifaddresses(interface).get(AF_INET, [])
+        subnets.extend(
+            [
+                ipaddress.ip_interface(link["addr"] + "/" + link["netmask"]).network
+                for link in addresses
+                if link["addr"] != "127.0.0.1"
+            ]
+        )
+    return subnets

monkey/monkey_island/cc/services/node.py

@@ -7,9 +7,9 @@ import monkey_island.cc.services.log
 from monkey_island.cc import models
 from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
+from monkey_island.cc.server_utils.network_utils import get_ip_addresses
 from monkey_island.cc.services.edge.displayed_edge import DisplayedEdgeService
 from monkey_island.cc.services.edge.edge import EdgeService
-from monkey_island.cc.services.utils.network_utils import get_ip_addresses
 from monkey_island.cc.services.utils.node_states import NodeStates
 
 

monkey/monkey_island/cc/services/reporting/report.py

@@ -10,6 +10,7 @@ from monkey_island.cc.database import mongo
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.models.report import get_report, save_report
 from monkey_island.cc.repository import IAgentConfigurationRepository, ICredentialsRepository
+from monkey_island.cc.server_utils.network_utils import get_ip_addresses, get_subnets
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.services.reporting.exploitations.manual_exploitation import get_manual_monkeys
 from monkey_island.cc.services.reporting.exploitations.monkey_exploitation import (
@@ -19,7 +20,6 @@ from monkey_island.cc.services.reporting.pth_report import PTHReportService
 from monkey_island.cc.services.reporting.report_generation_synchronisation import (
     safe_generate_regular_report,
 )
-from monkey_island.cc.services.utils.network_utils import get_ip_addresses, get_subnets
 
 from .. import AWSService
 from . import aws_exporter

monkey/monkey_island/cc/services/utils/network_utils.py

@@ -1,85 +0,0 @@
-import array
-import ipaddress
-import socket
-import struct
-import sys
-from typing import Sequence
-
-from netifaces import AF_INET, ifaddresses, interfaces
-from ring import lru
-
-# Local ips function
-# TODO: I can't find anywhere these are actually used. Confirm this is the case, remove these
-#       functions, and test.
-if sys.platform == "win32":
-
-    def local_ips():
-        local_hostname = socket.gethostname()
-        return socket.gethostbyname_ex(local_hostname)[2]
-
-else:
-    import fcntl
-
-    def local_ips():
-        result = []
-        try:
-            is_64bits = sys.maxsize > 2**32
-            struct_size = 40 if is_64bits else 32
-            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-            max_possible = 8  # initial value
-            while True:
-                struct_bytes = max_possible * struct_size
-                names = array.array("B", "\0" * struct_bytes)
-                outbytes = struct.unpack(
-                    "iL",
-                    fcntl.ioctl(
-                        s.fileno(),
-                        0x8912,  # SIOCGIFCONF
-                        struct.pack("iL", struct_bytes, names.buffer_info()[0]),
-                    ),
-                )[0]
-                if outbytes == struct_bytes:
-                    max_possible *= 2
-                else:
-                    break
-            namestr = names.tostring()
-
-            for i in range(0, outbytes, struct_size):
-                addr = socket.inet_ntoa(namestr[i + 20 : i + 24])
-                if not addr.startswith("127"):
-                    result.append(addr)
-                    # name of interface is (namestr[i:i+16].split('\0', 1)[0]
-        finally:
-            return result
-
-
-# The local IP addresses list should not change often. Therefore, we can cache the result and never
-# call this function more than once. This stopgap measure is here since this function is called a
-# lot of times during the report generation. This means that if the interfaces of the Island machine
-# change, the Island process needs to be restarted.
-@lru(maxsize=1)
-def get_ip_addresses() -> Sequence[str]:
-    ip_list = []
-    for interface in interfaces():
-        addresses = ifaddresses(interface).get(AF_INET, [])
-        ip_list.extend([link["addr"] for link in addresses if link["addr"] != "127.0.0.1"])
-    return ip_list
-
-
-# The subnets list should not change often. Therefore, we can cache the result and never call this
-# function more than once. This stopgap measure is here since this function is called a lot of times
-# during the report generation. This means that if the interfaces or subnets of the Island machine
-# change, the Island process needs to be restarted.
-@lru(maxsize=1)
-def get_subnets():
-    subnets = []
-    for interface in interfaces():
-        addresses = ifaddresses(interface).get(AF_INET, [])
-        subnets.extend(
-            [
-                ipaddress.ip_interface(link["addr"] + "/" + link["netmask"]).network
-                for link in addresses
-                if link["addr"] != "127.0.0.1"
-            ]
-        )
-    return subnets

monkey/monkey_island/cc/setup/__init__.py

@@ -1 +1,2 @@
 from .pywsgi_logging_filter import PyWSGILoggingFilter
+from .island_event_handlers import setup_island_event_handlers

monkey/monkey_island/cc/setup/island_event_handlers.py

@@ -0,0 +1,11 @@
+from common import DIContainer
+from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
+from monkey_island.cc.island_event_handlers import reset_agent_configuration
+
+
+def setup_island_event_handlers(container: DIContainer):
+    event_queue = container.resolve(PyPubSubIslandEventQueue)
+
+    event_queue.subscribe(
+        IslandEventTopic.RESET_AGENT_CONFIGURATION, container.resolve(reset_agent_configuration)
+    )

monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_reset_agent_configuration.py

@@ -0,0 +1,59 @@
+from unittest.mock import MagicMock
+
+import pytest
+from tests.monkey_island import InMemoryAgentConfigurationRepository
+
+from common.agent_configuration import AgentConfiguration
+from monkey_island.cc.island_event_handlers import reset_agent_configuration
+from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
+
+LINUX_FILENAME = "linux_pba_file.sh"
+WINDOWS_FILENAME = "windows_pba_file.ps1"
+
+
+@pytest.fixture
+def agent_configuration(default_agent_configuration: AgentConfiguration) -> AgentConfiguration:
+    custom_pbas = default_agent_configuration.custom_pbas.copy(
+        update={"linux_filename": LINUX_FILENAME, "windows_filename": WINDOWS_FILENAME},
+    )
+    return default_agent_configuration.copy(update={"custom_pbas": custom_pbas})
+
+
+@pytest.fixture
+def agent_configuration_repository(
+    agent_configuration: AgentConfiguration,
+) -> IAgentConfigurationRepository:
+    agent_configuration_repository = InMemoryAgentConfigurationRepository()
+    agent_configuration_repository.store_configuration(agent_configuration)
+
+    return agent_configuration_repository
+
+
+@pytest.fixture
+def mock_file_repository() -> IFileRepository:
+    return MagicMock(spec=IFileRepository)
+
+
+@pytest.fixture
+def callable_reset_agent_configuration(
+    agent_configuration_repository: IAgentConfigurationRepository,
+    mock_file_repository: IFileRepository,
+) -> reset_agent_configuration:
+    return reset_agent_configuration(agent_configuration_repository, mock_file_repository)
+
+
+def test_reset_configuration__remove_pba_files(
+    callable_reset_agent_configuration, mock_file_repository
+):
+    callable_reset_agent_configuration()
+
+    assert mock_file_repository.delete_file.called_with(LINUX_FILENAME)
+    assert mock_file_repository.delete_file.called_with(WINDOWS_FILENAME)
+
+
+def test_reset_configuration__agent_configuration_changed(
+    callable_reset_agent_configuration, agent_configuration_repository, agent_configuration
+):
+    callable_reset_agent_configuration()
+
+    assert agent_configuration_repository.get_configuration() != agent_configuration