From b10327af5c3d0c936606ce821f58725aed80b146 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 30 Aug 2022 15:36:30 +0000 Subject: [PATCH] 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()]