diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py index 25135231f..39075750a 100644 --- a/monkey/infection_monkey/i_control_channel.py +++ b/monkey/infection_monkey/i_control_channel.py @@ -1,11 +1,23 @@ import abc -from typing import Sequence +from typing import Optional, Sequence +from uuid import UUID from common.agent_configuration import AgentConfiguration from common.credentials import Credentials class IControlChannel(metaclass=abc.ABCMeta): + @abc.abstractmethod + def register_agent(self, parent_id: Optional[UUID] = None): + """ + Registers this agent with the Island when this agent starts + + :param parent: The ID of the parent that spawned this agent, or None if this agent has no + parent + :raises IslandCommunicationError: If the agent cannot be successfully registered + """ + pass + @abc.abstractmethod def should_agent_stop(self) -> bool: """ diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index a4709b124..8fa64d461 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -1,14 +1,18 @@ import json import logging from pprint import pformat -from typing import Mapping, Sequence +from typing import Mapping, Optional, Sequence +from uuid import UUID import requests +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.utils import agent_process +from infection_monkey.utils.ids import get_agent_id, get_machine_id requests.packages.urllib3.disable_warnings() @@ -21,6 +25,34 @@ class ControlChannel(IControlChannel): self._control_channel_server = server self._proxies = proxies + def register_agent(self, parent: Optional[UUID] = None): + agent_registration_data = AgentRegistrationData( + id=get_agent_id(), + machine_hardware_id=get_machine_id(), + start_time=agent_process.get_start_time(), + parent_id=parent, + cc_server=self._control_channel_server, + network_interfaces=[], # TODO: Populate this + ) + + try: + url = f"https://{self._control_channel_server}/api/agents" + response = requests.post( # noqa: DUO123 + url, + json=agent_registration_data.dict(simplify=True), + verify=False, + proxies=self._proxies, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + except ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + requests.exceptions.TooManyRedirects, + requests.exceptions.HTTPError, + ) as e: + raise IslandCommunicationError(e) + def should_agent_stop(self) -> bool: if not self._control_channel_server: logger.error("Agent should stop because it can't connect to the C&C server.") diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 98b2cab78..7d6962e88 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -178,6 +178,7 @@ class InfectionMonkey: control_channel = ControlChannel( self._control_client.server_address, GUID, self._control_client.proxies ) + control_channel.register_agent(self._opts.parent) config = control_channel.get_config() self._monkey_inbound_tunnel = self._control_client.create_control_tunnel(