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 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],
):

View File

@ -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,

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
# queue is empty.
addresses = Queue()
addresses: Queue = Queue()
for address in addresses_to_scan:
addresses.put(address)

View File

@ -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

View File

@ -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()]