forked from p34709852/monkey
Agent: Register agent on startup
This commit is contained in:
parent
80b4b05b65
commit
9c32dee739
|
@ -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:
|
||||||
"""
|
"""
|
||||||
|
|
|
@ -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.")
|
||||||
|
|
|
@ -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(
|
||||||
|
|
Loading…
Reference in New Issue