From b10327af5c3d0c936606ce821f58725aed80b146 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Tue, 30 Aug 2022 15:36:30 +0000
Subject: [PATCH 1/2] Agent: Register agent interfaces

---
 monkey/infection_monkey/master/automated_master.py |  6 +++---
 monkey/infection_monkey/master/control_channel.py  | 14 ++++++++------
 monkey/infection_monkey/master/ip_scanner.py       |  2 +-
 monkey/infection_monkey/master/propagator.py       |  2 +-
 monkey/infection_monkey/network/info.py            |  4 ++++
 5 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/monkey/infection_monkey/master/automated_master.py b/monkey/infection_monkey/master/automated_master.py
index e491e4537..45b329ad1 100644
--- a/monkey/infection_monkey/master/automated_master.py
+++ b/monkey/infection_monkey/master/automated_master.py
@@ -2,7 +2,7 @@ import logging
 import threading
 import time
 from ipaddress import IPv4Interface
-from typing import Any, Callable, Iterable, List, Optional
+from typing import Any, Callable, Collection, List, Optional
 
 from common.agent_configuration import CustomPBAConfiguration, PluginConfiguration
 from common.utils import Timer
@@ -206,7 +206,7 @@ class AutomatedMaster(IMaster):
 
     def _run_pbas(
         self,
-        plugins: Iterable[PluginConfiguration],
+        plugins: Collection[PluginConfiguration],
         callback: Callable[[Any], None],
         custom_pba_options: CustomPBAConfiguration,
     ):
@@ -221,7 +221,7 @@ class AutomatedMaster(IMaster):
 
     def _run_plugins(
         self,
-        plugins: Iterable[PluginConfiguration],
+        plugins: Collection[PluginConfiguration],
         plugin_type: str,
         callback: Callable[[Any], None],
     ):
diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py
index 8fa64d461..93fa67fa6 100644
--- a/monkey/infection_monkey/master/control_channel.py
+++ b/monkey/infection_monkey/master/control_channel.py
@@ -1,26 +1,28 @@
 import json
 import logging
 from pprint import pformat
-from typing import Mapping, Optional, Sequence
+from typing import MutableMapping, Optional, Sequence
 from uuid import UUID
 
 import requests
+from urllib3 import disable_warnings
 
 from common import AgentRegistrationData
 from common.agent_configuration import AgentConfiguration
 from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
 from common.credentials import Credentials
 from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError
+from infection_monkey.network.info import get_network_interfaces
 from infection_monkey.utils import agent_process
 from infection_monkey.utils.ids import get_agent_id, get_machine_id
 
-requests.packages.urllib3.disable_warnings()
+disable_warnings()  # noqa: DUO131
 
 logger = logging.getLogger(__name__)
 
 
 class ControlChannel(IControlChannel):
-    def __init__(self, server: str, agent_id: str, proxies: Mapping[str, str]):
+    def __init__(self, server: str, agent_id: str, proxies: MutableMapping[str, str]):
         self._agent_id = agent_id
         self._control_channel_server = server
         self._proxies = proxies
@@ -32,7 +34,7 @@ class ControlChannel(IControlChannel):
             start_time=agent_process.get_start_time(),
             parent_id=parent,
             cc_server=self._control_channel_server,
-            network_interfaces=[],  # TODO: Populate this
+            network_interfaces=get_network_interfaces(),
         )
 
         try:
@@ -70,8 +72,8 @@ class ControlChannel(IControlChannel):
             )
             response.raise_for_status()
 
-            response = json.loads(response.content.decode())
-            return response["stop_agent"]
+            json_response = json.loads(response.content.decode())
+            return json_response["stop_agent"]
         except (
             json.JSONDecodeError,
             requests.exceptions.ConnectionError,
diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py
index f3e733839..bbabd7ae7 100644
--- a/monkey/infection_monkey/master/ip_scanner.py
+++ b/monkey/infection_monkey/master/ip_scanner.py
@@ -41,7 +41,7 @@ class IPScanner:
     ):
         # Pre-fill a Queue with all IPs to scan so that threads know they can safely exit when the
         # queue is empty.
-        addresses = Queue()
+        addresses: Queue = Queue()
         for address in addresses_to_scan:
             addresses.put(address)
 
diff --git a/monkey/infection_monkey/master/propagator.py b/monkey/infection_monkey/master/propagator.py
index 62f489657..c11e38d46 100644
--- a/monkey/infection_monkey/master/propagator.py
+++ b/monkey/infection_monkey/master/propagator.py
@@ -45,7 +45,7 @@ class Propagator:
         self._exploiter = exploiter
         self._victim_host_factory = victim_host_factory
         self._local_network_interfaces = local_network_interfaces
-        self._hosts_to_exploit = None
+        self._hosts_to_exploit: Queue = Queue()
 
     def propagate(
         self, propagation_config: PropagationConfiguration, current_depth: int, stop: Event
diff --git a/monkey/infection_monkey/network/info.py b/monkey/infection_monkey/network/info.py
index 2ae49b056..304d37ff6 100644
--- a/monkey/infection_monkey/network/info.py
+++ b/monkey/infection_monkey/network/info.py
@@ -25,6 +25,10 @@ class NetworkAddress:
     domain: str
 
 
+def get_network_interfaces() -> List[IPv4Interface]:
+    return get_local_network_interfaces()
+
+
 def get_local_network_interfaces() -> List[IPv4Interface]:
     return [IPv4Interface(f"{i['addr']}/{i['netmask']}") for i in get_host_subnets()]
 

From a954df3ed89ad359201d0fb6e1e967f8054a775d Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 31 Aug 2022 15:34:48 +0000
Subject: [PATCH 2/2] Agent: Replace get_local_network_interfaces

Replaced get_local_network_interfaces() with get_network_interfaces()
---
 monkey/infection_monkey/monkey.py       | 4 ++--
 monkey/infection_monkey/network/info.py | 4 ----
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index 220eba1f8..115b67e24 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/monkey/infection_monkey/monkey.py
@@ -41,7 +41,7 @@ from infection_monkey.master import AutomatedMaster
 from infection_monkey.master.control_channel import ControlChannel
 from infection_monkey.model import VictimHostFactory
 from infection_monkey.network.firewall import app as firewall
-from infection_monkey.network.info import get_local_network_interfaces
+from infection_monkey.network.info import get_network_interfaces
 from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter
 from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter
 from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter
@@ -240,7 +240,7 @@ class InfectionMonkey:
 
     @staticmethod
     def _get_local_network_interfaces() -> List[IPv4Interface]:
-        local_network_interfaces = get_local_network_interfaces()
+        local_network_interfaces = get_network_interfaces()
         for interface in local_network_interfaces:
             logger.debug(f"Found local interface {str(interface)}")
 
diff --git a/monkey/infection_monkey/network/info.py b/monkey/infection_monkey/network/info.py
index 304d37ff6..12748b8a0 100644
--- a/monkey/infection_monkey/network/info.py
+++ b/monkey/infection_monkey/network/info.py
@@ -26,10 +26,6 @@ class NetworkAddress:
 
 
 def get_network_interfaces() -> List[IPv4Interface]:
-    return get_local_network_interfaces()
-
-
-def get_local_network_interfaces() -> List[IPv4Interface]:
     return [IPv4Interface(f"{i['addr']}/{i['netmask']}") for i in get_host_subnets()]