Agent: Register agent on startup

This commit is contained in:
Mike Salvatore 2022-08-22 14:37:34 -04:00
parent 80b4b05b65
commit 9c32dee739
3 changed files with 47 additions and 2 deletions

View File

@ -1,11 +1,23 @@
import abc import abc
from typing import Sequence from typing import Optional, Sequence
from uuid import UUID
from common.agent_configuration import AgentConfiguration from common.agent_configuration import AgentConfiguration
from common.credentials import Credentials from common.credentials import Credentials
class IControlChannel(metaclass=abc.ABCMeta): 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 @abc.abstractmethod
def should_agent_stop(self) -> bool: def should_agent_stop(self) -> bool:
""" """

View File

@ -1,14 +1,18 @@
import json import json
import logging import logging
from pprint import pformat from pprint import pformat
from typing import Mapping, Sequence from typing import Mapping, Optional, Sequence
from uuid import UUID
import requests import requests
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.utils import agent_process
from infection_monkey.utils.ids import get_agent_id, get_machine_id
requests.packages.urllib3.disable_warnings() requests.packages.urllib3.disable_warnings()
@ -21,6 +25,34 @@ class ControlChannel(IControlChannel):
self._control_channel_server = server self._control_channel_server = server
self._proxies = proxies 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: def should_agent_stop(self) -> bool:
if not self._control_channel_server: if not self._control_channel_server:
logger.error("Agent should stop because it can't connect to the C&C server.") logger.error("Agent should stop because it can't connect to the C&C server.")

View File

@ -178,6 +178,7 @@ class InfectionMonkey:
control_channel = ControlChannel( control_channel = ControlChannel(
self._control_client.server_address, GUID, self._control_client.proxies self._control_client.server_address, GUID, self._control_client.proxies
) )
control_channel.register_agent(self._opts.parent)
config = control_channel.get_config() config = control_channel.get_config()
self._monkey_inbound_tunnel = self._control_client.create_control_tunnel( self._monkey_inbound_tunnel = self._control_client.create_control_tunnel(