forked from p15670423/monkey
Merge pull request #2389 from guardicore/2269-notify-relay-on-propagation
2269 notify relay on propagation
This commit is contained in:
commit
a3ce870b64
|
@ -0,0 +1 @@
|
|||
from .notify_relay_on_propagation import notify_relay_on_propagation
|
|
@ -0,0 +1,31 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from common.agent_events import PropagationEvent
|
||||
from infection_monkey.network.relay import TCPRelay
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class notify_relay_on_propagation:
|
||||
"""
|
||||
Notifies a TCPRelay of potential relay users if propagation is successful
|
||||
"""
|
||||
|
||||
def __init__(self, tcp_relay: Optional[TCPRelay]):
|
||||
"""
|
||||
:param tcp_relay: A TCPRelay to notify on successful propagation
|
||||
"""
|
||||
self._tcp_relay = tcp_relay
|
||||
|
||||
def __call__(self, event: PropagationEvent):
|
||||
"""
|
||||
Notify a TCPRelay of potential relay users if propagation is successful
|
||||
|
||||
:param event: A `PropagationEvent`
|
||||
"""
|
||||
if self._tcp_relay is None:
|
||||
return
|
||||
|
||||
if event.success:
|
||||
self._tcp_relay.add_potential_user(event.target)
|
|
@ -13,7 +13,7 @@ from common.agent_event_serializers import (
|
|||
AgentEventSerializerRegistry,
|
||||
register_common_agent_event_serializers,
|
||||
)
|
||||
from common.agent_events import CredentialsStolenEvent
|
||||
from common.agent_events import CredentialsStolenEvent, PropagationEvent
|
||||
from common.agent_registration_data import AgentRegistrationData
|
||||
from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue
|
||||
from common.network.network_utils import get_my_ip_addresses, get_network_interfaces
|
||||
|
@ -22,6 +22,7 @@ from common.utils.argparse_types import positive_int
|
|||
from common.utils.attack_utils import ScanStatus, UsageEnum
|
||||
from common.version import get_version
|
||||
from infection_monkey.agent_event_forwarder import AgentEventForwarder
|
||||
from infection_monkey.agent_event_handlers import notify_relay_on_propagation
|
||||
from infection_monkey.config import GUID
|
||||
from infection_monkey.control import ControlClient
|
||||
from infection_monkey.credential_collectors import (
|
||||
|
@ -80,9 +81,6 @@ from infection_monkey.puppet.puppet import Puppet
|
|||
from infection_monkey.system_singleton import SystemSingleton
|
||||
from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
|
||||
from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
|
||||
from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
|
||||
ExploitInterceptingTelemetryMessenger,
|
||||
)
|
||||
from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
|
||||
LegacyTelemetryMessengerAdapter,
|
||||
)
|
||||
|
@ -278,15 +276,11 @@ class InfectionMonkey:
|
|||
|
||||
victim_host_factory = self._build_victim_host_factory(local_network_interfaces)
|
||||
|
||||
telemetry_messenger = ExploitInterceptingTelemetryMessenger(
|
||||
self._telemetry_messenger, self._relay
|
||||
)
|
||||
|
||||
self._master = AutomatedMaster(
|
||||
self._current_depth,
|
||||
servers,
|
||||
puppet,
|
||||
telemetry_messenger,
|
||||
self._telemetry_messenger,
|
||||
victim_host_factory,
|
||||
self._control_channel,
|
||||
local_network_interfaces,
|
||||
|
@ -313,6 +307,7 @@ class InfectionMonkey:
|
|||
agent_event_queue.subscribe_all_events(
|
||||
AgentEventForwarder(self._island_api_client, agent_event_serializer_registry).send_event
|
||||
)
|
||||
agent_event_queue.subscribe_type(PropagationEvent, notify_relay_on_propagation(self._relay))
|
||||
|
||||
def _build_puppet(
|
||||
self,
|
||||
|
|
|
@ -1,41 +0,0 @@
|
|||
from functools import singledispatch
|
||||
from ipaddress import IPv4Address
|
||||
|
||||
from infection_monkey.network.relay import TCPRelay
|
||||
from infection_monkey.telemetry.exploit_telem import ExploitTelem
|
||||
from infection_monkey.telemetry.i_telem import ITelem
|
||||
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
|
||||
|
||||
|
||||
class ExploitInterceptingTelemetryMessenger(ITelemetryMessenger):
|
||||
def __init__(self, telemetry_messenger: ITelemetryMessenger, relay: TCPRelay):
|
||||
self._telemetry_messenger = telemetry_messenger
|
||||
self._relay = relay
|
||||
|
||||
def send_telemetry(self, telemetry: ITelem):
|
||||
_send_telemetry(telemetry, self._telemetry_messenger, self._relay)
|
||||
|
||||
|
||||
# Note: We can use @singledispatchmethod instead of @singledispatch if we migrate to Python 3.8 or
|
||||
# later.
|
||||
@singledispatch
|
||||
def _send_telemetry(
|
||||
telemetry: ITelem,
|
||||
telemetry_messenger: ITelemetryMessenger,
|
||||
relay: TCPRelay,
|
||||
):
|
||||
telemetry_messenger.send_telemetry(telemetry)
|
||||
|
||||
|
||||
@_send_telemetry.register
|
||||
def _(
|
||||
telemetry: ExploitTelem,
|
||||
telemetry_messenger: ITelemetryMessenger,
|
||||
relay: TCPRelay,
|
||||
):
|
||||
if telemetry.propagation_result is True:
|
||||
if relay:
|
||||
address = IPv4Address(str(telemetry.host["ip_addr"]))
|
||||
relay.add_potential_user(address)
|
||||
|
||||
telemetry_messenger.send_telemetry(telemetry)
|
|
@ -0,0 +1,54 @@
|
|||
from ipaddress import IPv4Address
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
|
||||
from common.agent_events import PropagationEvent
|
||||
from infection_monkey.agent_event_handlers import notify_relay_on_propagation
|
||||
from infection_monkey.network.relay import TCPRelay
|
||||
|
||||
TARGET_ADDRESS = IPv4Address("192.168.1.10")
|
||||
|
||||
SUCCESSFUL_PROPAGATION_EVENT = PropagationEvent(
|
||||
source=UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5"),
|
||||
target=TARGET_ADDRESS,
|
||||
tags=frozenset({"test"}),
|
||||
success=True,
|
||||
exploiter_name="test_exploiter",
|
||||
)
|
||||
|
||||
FAILED_PROPAGATION_EVENT = PropagationEvent(
|
||||
source=UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5"),
|
||||
target=TARGET_ADDRESS,
|
||||
tags=frozenset({"test"}),
|
||||
success=False,
|
||||
exploiter_name="test_exploiter",
|
||||
error_message="everything is broken",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tcp_relay():
|
||||
return MagicMock(spec=TCPRelay)
|
||||
|
||||
|
||||
def test_relay_notified_on_successful_propation(mock_tcp_relay):
|
||||
handler = notify_relay_on_propagation(mock_tcp_relay)
|
||||
handler(SUCCESSFUL_PROPAGATION_EVENT)
|
||||
|
||||
mock_tcp_relay.add_potential_user.assert_called_once_with(TARGET_ADDRESS)
|
||||
|
||||
|
||||
def test_relay_not_notified_on_successful_propation(mock_tcp_relay):
|
||||
handler = notify_relay_on_propagation(mock_tcp_relay)
|
||||
handler(FAILED_PROPAGATION_EVENT)
|
||||
|
||||
mock_tcp_relay.add_potential_user.assert_not_called()
|
||||
|
||||
|
||||
def test_handler_doesnt_raise_if_relay_is_none():
|
||||
handler = notify_relay_on_propagation(None)
|
||||
|
||||
# Raises AttributeError on failure
|
||||
handler(SUCCESSFUL_PROPAGATION_EVENT)
|
|
@ -1,62 +0,0 @@
|
|||
from unittest.mock import MagicMock
|
||||
|
||||
from infection_monkey.i_puppet.i_puppet import ExploiterResultData
|
||||
from infection_monkey.model.host import VictimHost
|
||||
from infection_monkey.telemetry.exploit_telem import ExploitTelem
|
||||
from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
|
||||
ExploitInterceptingTelemetryMessenger,
|
||||
)
|
||||
|
||||
|
||||
class MockExploitTelem(ExploitTelem):
|
||||
def __init__(self, propagation_success):
|
||||
erd = ExploiterResultData()
|
||||
erd.propagation_success = propagation_success
|
||||
super().__init__("TestExploiter", VictimHost("127.0.0.1"), erd)
|
||||
|
||||
def get_data(self):
|
||||
return {}
|
||||
|
||||
|
||||
def test_generic_telemetry(TestTelem):
|
||||
mock_telemetry_messenger = MagicMock()
|
||||
mock_relay = MagicMock()
|
||||
|
||||
telemetry_messenger = ExploitInterceptingTelemetryMessenger(
|
||||
mock_telemetry_messenger, mock_relay
|
||||
)
|
||||
|
||||
telemetry_messenger.send_telemetry(TestTelem())
|
||||
|
||||
assert mock_telemetry_messenger.send_telemetry.called
|
||||
assert not mock_relay.add_potential_user.called
|
||||
|
||||
|
||||
def test_propagation_successful_exploit_telemetry():
|
||||
mock_telemetry_messenger = MagicMock()
|
||||
mock_relay = MagicMock()
|
||||
mock_exploit_telem = MockExploitTelem(True)
|
||||
|
||||
telemetry_messenger = ExploitInterceptingTelemetryMessenger(
|
||||
mock_telemetry_messenger, mock_relay
|
||||
)
|
||||
|
||||
telemetry_messenger.send_telemetry(mock_exploit_telem)
|
||||
|
||||
assert mock_telemetry_messenger.send_telemetry.called
|
||||
assert mock_relay.add_potential_user.called
|
||||
|
||||
|
||||
def test_propagation_failed_exploit_telemetry():
|
||||
mock_telemetry_messenger = MagicMock()
|
||||
mock_relay = MagicMock()
|
||||
mock_exploit_telem = MockExploitTelem(False)
|
||||
|
||||
telemetry_messenger = ExploitInterceptingTelemetryMessenger(
|
||||
mock_telemetry_messenger, mock_relay
|
||||
)
|
||||
|
||||
telemetry_messenger.send_telemetry(mock_exploit_telem)
|
||||
|
||||
assert mock_telemetry_messenger.send_telemetry.called
|
||||
assert not mock_relay.add_potential_user.called
|
Loading…
Reference in New Issue