Agent: Register agent interfaces

This commit is contained in:
Kekoa Kaaikala 2022-08-30 15:36:30 +00:00
parent 70a9251c5b
commit b10327af5c
5 changed files with 17 additions and 11 deletions

View File

@ -2,7 +2,7 @@ import logging
import threading import threading
import time import time
from ipaddress import IPv4Interface 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.agent_configuration import CustomPBAConfiguration, PluginConfiguration
from common.utils import Timer from common.utils import Timer
@ -206,7 +206,7 @@ class AutomatedMaster(IMaster):
def _run_pbas( def _run_pbas(
self, self,
plugins: Iterable[PluginConfiguration], plugins: Collection[PluginConfiguration],
callback: Callable[[Any], None], callback: Callable[[Any], None],
custom_pba_options: CustomPBAConfiguration, custom_pba_options: CustomPBAConfiguration,
): ):
@ -221,7 +221,7 @@ class AutomatedMaster(IMaster):
def _run_plugins( def _run_plugins(
self, self,
plugins: Iterable[PluginConfiguration], plugins: Collection[PluginConfiguration],
plugin_type: str, plugin_type: str,
callback: Callable[[Any], None], callback: Callable[[Any], None],
): ):

View File

@ -1,26 +1,28 @@
import json import json
import logging import logging
from pprint import pformat from pprint import pformat
from typing import Mapping, Optional, Sequence from typing import MutableMapping, Optional, Sequence
from uuid import UUID from uuid import UUID
import requests import requests
from urllib3 import disable_warnings
from common import AgentRegistrationData from common import AgentRegistrationData
from common.agent_configuration import AgentConfiguration from common.agent_configuration import AgentConfiguration
from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
from common.credentials import Credentials from common.credentials import Credentials
from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError 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 import agent_process
from infection_monkey.utils.ids import get_agent_id, get_machine_id 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__) logger = logging.getLogger(__name__)
class ControlChannel(IControlChannel): 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._agent_id = agent_id
self._control_channel_server = server self._control_channel_server = server
self._proxies = proxies self._proxies = proxies
@ -32,7 +34,7 @@ class ControlChannel(IControlChannel):
start_time=agent_process.get_start_time(), start_time=agent_process.get_start_time(),
parent_id=parent, parent_id=parent,
cc_server=self._control_channel_server, cc_server=self._control_channel_server,
network_interfaces=[], # TODO: Populate this network_interfaces=get_network_interfaces(),
) )
try: try:
@ -70,8 +72,8 @@ class ControlChannel(IControlChannel):
) )
response.raise_for_status() response.raise_for_status()
response = json.loads(response.content.decode()) json_response = json.loads(response.content.decode())
return response["stop_agent"] return json_response["stop_agent"]
except ( except (
json.JSONDecodeError, json.JSONDecodeError,
requests.exceptions.ConnectionError, requests.exceptions.ConnectionError,

View File

@ -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 # Pre-fill a Queue with all IPs to scan so that threads know they can safely exit when the
# queue is empty. # queue is empty.
addresses = Queue() addresses: Queue = Queue()
for address in addresses_to_scan: for address in addresses_to_scan:
addresses.put(address) addresses.put(address)

View File

@ -45,7 +45,7 @@ class Propagator:
self._exploiter = exploiter self._exploiter = exploiter
self._victim_host_factory = victim_host_factory self._victim_host_factory = victim_host_factory
self._local_network_interfaces = local_network_interfaces self._local_network_interfaces = local_network_interfaces
self._hosts_to_exploit = None self._hosts_to_exploit: Queue = Queue()
def propagate( def propagate(
self, propagation_config: PropagationConfiguration, current_depth: int, stop: Event self, propagation_config: PropagationConfiguration, current_depth: int, stop: Event

View File

@ -25,6 +25,10 @@ class NetworkAddress:
domain: str domain: str
def get_network_interfaces() -> List[IPv4Interface]:
return get_local_network_interfaces()
def get_local_network_interfaces() -> List[IPv4Interface]: def get_local_network_interfaces() -> List[IPv4Interface]:
return [IPv4Interface(f"{i['addr']}/{i['netmask']}") for i in get_host_subnets()] return [IPv4Interface(f"{i['addr']}/{i['netmask']}") for i in get_host_subnets()]