Merge pull request #2389 from guardicore/2269-notify-relay-on-propagation

2269 notify relay on propagation
This commit is contained in:
Mike Salvatore 2022-10-03 14:46:14 -04:00 committed by GitHub
commit a3ce870b64
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 90 additions and 112 deletions

View File

@ -0,0 +1 @@
from .notify_relay_on_propagation import notify_relay_on_propagation

View File

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

View File

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

View File

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

View File

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

View File

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