forked from p15670423/monkey
Island: Remove handle_ping_scan_event callable class
This commit is contained in:
parent
8002080c8b
commit
f23093dc78
|
@ -1,4 +1,3 @@
|
||||||
from .handle_ping_scan_event import handle_ping_scan_event
|
|
||||||
from .save_event_to_event_repository import save_event_to_event_repository
|
from .save_event_to_event_repository import save_event_to_event_repository
|
||||||
from .save_stolen_credentials_to_repository import save_stolen_credentials_to_repository
|
from .save_stolen_credentials_to_repository import save_stolen_credentials_to_repository
|
||||||
from .scan_event_handler import ScanEventHandler
|
from .scan_event_handler import ScanEventHandler
|
||||||
|
|
|
@ -1,76 +0,0 @@
|
||||||
from logging import getLogger
|
|
||||||
|
|
||||||
from pydantic.json import IPv4Interface
|
|
||||||
|
|
||||||
from common.agent_events import PingScanEvent
|
|
||||||
from monkey_island.cc.models import CommunicationType, Machine
|
|
||||||
from monkey_island.cc.repository import (
|
|
||||||
IAgentRepository,
|
|
||||||
IMachineRepository,
|
|
||||||
INodeRepository,
|
|
||||||
RetrievalError,
|
|
||||||
StorageError,
|
|
||||||
UnknownRecordError,
|
|
||||||
)
|
|
||||||
|
|
||||||
logger = getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
class handle_ping_scan_event:
|
|
||||||
"""
|
|
||||||
Handles ping scan event and makes changes to Machine and Node states based on it
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
agent_repository: IAgentRepository,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
):
|
|
||||||
self._agent_repository = agent_repository
|
|
||||||
self._machine_repository = machine_repository
|
|
||||||
self._node_repository = node_repository
|
|
||||||
|
|
||||||
def __call__(self, event: PingScanEvent):
|
|
||||||
"""
|
|
||||||
:param event: Ping scan event to process
|
|
||||||
"""
|
|
||||||
if not event.response_received:
|
|
||||||
return
|
|
||||||
|
|
||||||
try:
|
|
||||||
target_machine = self._get_target_machine(event)
|
|
||||||
|
|
||||||
self._update_target_machine_os(target_machine, event)
|
|
||||||
self._update_nodes(target_machine, event)
|
|
||||||
except (RetrievalError, StorageError, UnknownRecordError):
|
|
||||||
logger.exception("Unable to process ping scan data")
|
|
||||||
|
|
||||||
def _get_target_machine(self, event: PingScanEvent) -> Machine:
|
|
||||||
try:
|
|
||||||
target_machines = self._machine_repository.get_machines_by_ip(event.target)
|
|
||||||
return target_machines[0]
|
|
||||||
except UnknownRecordError:
|
|
||||||
machine = Machine(
|
|
||||||
id=self._machine_repository.get_new_id(),
|
|
||||||
network_interfaces=[IPv4Interface(event.target)],
|
|
||||||
)
|
|
||||||
self._machine_repository.upsert_machine(machine)
|
|
||||||
return machine
|
|
||||||
|
|
||||||
def _update_target_machine_os(self, machine: Machine, event: PingScanEvent):
|
|
||||||
if event.os is not None and machine.operating_system is None:
|
|
||||||
machine.operating_system = event.os
|
|
||||||
self._machine_repository.upsert_machine(machine)
|
|
||||||
|
|
||||||
def _update_nodes(self, target_machine, event):
|
|
||||||
src_machine = self._get_source_machine(event)
|
|
||||||
|
|
||||||
# Update or create the node
|
|
||||||
self._node_repository.upsert_communication(
|
|
||||||
src_machine.id, target_machine.id, CommunicationType.SCANNED
|
|
||||||
)
|
|
||||||
|
|
||||||
def _get_source_machine(self, event: PingScanEvent) -> Machine:
|
|
||||||
agent = self._agent_repository.get_agent_by_id(event.source)
|
|
||||||
return self._machine_repository.get_machine_by_id(agent.machine_id)
|
|
|
@ -1,268 +0,0 @@
|
||||||
from ipaddress import IPv4Address, IPv4Interface
|
|
||||||
from itertools import count
|
|
||||||
from unittest.mock import MagicMock
|
|
||||||
from uuid import UUID
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
from common import OperatingSystem
|
|
||||||
from common.agent_events import PingScanEvent
|
|
||||||
from common.types import SocketAddress
|
|
||||||
from monkey_island.cc.agent_event_handlers import handle_ping_scan_event
|
|
||||||
from monkey_island.cc.models import Agent, CommunicationType, Machine
|
|
||||||
from monkey_island.cc.repository import (
|
|
||||||
IAgentRepository,
|
|
||||||
IMachineRepository,
|
|
||||||
INodeRepository,
|
|
||||||
RetrievalError,
|
|
||||||
StorageError,
|
|
||||||
UnknownRecordError,
|
|
||||||
)
|
|
||||||
|
|
||||||
SEED_ID = 99
|
|
||||||
AGENT_ID = UUID("1d8ce743-a0f4-45c5-96af-91106529d3e2")
|
|
||||||
MACHINE_ID = 11
|
|
||||||
CC_SERVER = SocketAddress(ip="10.10.10.100", port="5000")
|
|
||||||
AGENT = Agent(id=AGENT_ID, machine_id=MACHINE_ID, start_time=0, parent_id=None, cc_server=CC_SERVER)
|
|
||||||
PINGER_MACHINE = Machine(
|
|
||||||
id=MACHINE_ID,
|
|
||||||
hardware_id=5,
|
|
||||||
network_interfaces=[IPv4Interface("10.10.10.99/24")],
|
|
||||||
)
|
|
||||||
TARGET_MACHINE = Machine(
|
|
||||||
id=33,
|
|
||||||
hardware_id=9,
|
|
||||||
network_interfaces=[IPv4Interface("10.10.10.1/24")],
|
|
||||||
)
|
|
||||||
EVENT = PingScanEvent(
|
|
||||||
source=AGENT_ID,
|
|
||||||
target=IPv4Address("10.10.10.1"),
|
|
||||||
response_received=True,
|
|
||||||
os=OperatingSystem.LINUX,
|
|
||||||
)
|
|
||||||
|
|
||||||
EVENT_NO_RESPONSE = PingScanEvent(
|
|
||||||
source=AGENT_ID,
|
|
||||||
target=IPv4Address("10.10.10.1"),
|
|
||||||
response_received=False,
|
|
||||||
os=OperatingSystem.LINUX,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def agent_repository() -> IAgentRepository:
|
|
||||||
agent_repository = MagicMock(spec=IAgentRepository)
|
|
||||||
agent_repository.upsert_agent = MagicMock()
|
|
||||||
agent_repository.get_agent_by_id = MagicMock(return_value=AGENT)
|
|
||||||
return agent_repository
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def machine_repository() -> IMachineRepository:
|
|
||||||
machine_repository = MagicMock(spec=IMachineRepository)
|
|
||||||
machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID))
|
|
||||||
machine_repository.upsert_machine = MagicMock()
|
|
||||||
return machine_repository
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def node_repository() -> INodeRepository:
|
|
||||||
node_repository = MagicMock(spec=INodeRepository)
|
|
||||||
node_repository.upsert_communication = MagicMock()
|
|
||||||
return node_repository
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
|
||||||
def handler(agent_repository, machine_repository, node_repository) -> handle_ping_scan_event:
|
|
||||||
return handle_ping_scan_event(agent_repository, machine_repository, node_repository)
|
|
||||||
|
|
||||||
|
|
||||||
machines = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE}
|
|
||||||
machines_by_id = {MACHINE_ID: PINGER_MACHINE, TARGET_MACHINE.id: TARGET_MACHINE}
|
|
||||||
machines_by_ip = {
|
|
||||||
IPv4Address("10.10.10.99"): [PINGER_MACHINE],
|
|
||||||
IPv4Address("10.10.10.1"): [TARGET_MACHINE],
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def machine_from_id(id: int):
|
|
||||||
return machines_by_id[id]
|
|
||||||
|
|
||||||
|
|
||||||
def machines_from_ip(ip: IPv4Address):
|
|
||||||
return machines_by_ip[ip]
|
|
||||||
|
|
||||||
|
|
||||||
class error_machine_by_id:
|
|
||||||
"""Raise an error if the machine with the called ID matches the stored ID"""
|
|
||||||
|
|
||||||
def __init__(self, id: int, error):
|
|
||||||
self.id = id
|
|
||||||
self.error = error
|
|
||||||
|
|
||||||
def __call__(self, id: int):
|
|
||||||
if id == self.id:
|
|
||||||
raise self.error
|
|
||||||
else:
|
|
||||||
return machine_from_id(id)
|
|
||||||
|
|
||||||
|
|
||||||
class error_machine_by_ip:
|
|
||||||
"""Raise an error if the machine with the called IP matches the stored ID"""
|
|
||||||
|
|
||||||
def __init__(self, id: int, error):
|
|
||||||
self.id = id
|
|
||||||
self.error = error
|
|
||||||
|
|
||||||
def __call__(self, ip: IPv4Address):
|
|
||||||
print(f"IP is: {ip}")
|
|
||||||
machines = machines_from_ip(ip)
|
|
||||||
if machines[0].id == self.id:
|
|
||||||
print(f"Raise error: {self.error}")
|
|
||||||
raise self.error
|
|
||||||
else:
|
|
||||||
print(f"Return machine: {machines}")
|
|
||||||
return machines
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_ping_scan_event__target_machine_not_exists(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
):
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
expected_machine = Machine(id=SEED_ID, network_interfaces=[IPv4Interface(EVENT.target)])
|
|
||||||
expected_machine.operating_system = EVENT.os
|
|
||||||
machine_repository.upsert_machine.assert_called_with(expected_machine)
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_ping_scan_event__target_machine_already_exists(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
):
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
expected_machine = TARGET_MACHINE.copy()
|
|
||||||
expected_machine.operating_system = OperatingSystem.LINUX
|
|
||||||
machine_repository.upsert_machine.assert_called_with(expected_machine)
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_ping_scan_event__upserts_node(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
):
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(return_value=[TARGET_MACHINE])
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
node_repository.upsert_communication.assert_called_with(
|
|
||||||
PINGER_MACHINE.id, TARGET_MACHINE.id, CommunicationType.SCANNED
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_ping_scan_event__node_not_upserted_if_no_matching_agent(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
agent_repository: IAgentRepository,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
):
|
|
||||||
agent_repository.get_agent_by_id = MagicMock(side_effect=UnknownRecordError)
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(return_value=TARGET_MACHINE)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
assert not node_repository.upsert_communication.called
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_ping_scan_event__node_not_upserted_if_no_matching_machine(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
):
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=UnknownRecordError)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
assert not node_repository.upsert_communication.called
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("id", [PINGER_MACHINE.id, TARGET_MACHINE.id])
|
|
||||||
def test_handle_scan_data__node_not_upserted_if_machine_retrievalerror(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
id,
|
|
||||||
):
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(
|
|
||||||
side_effect=error_machine_by_id(id, RetrievalError)
|
|
||||||
)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(
|
|
||||||
side_effect=error_machine_by_ip(id, RetrievalError)
|
|
||||||
)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
assert not node_repository.upsert_communication.called
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_scan_data__machine_not_upserted_if_os_is_none(
|
|
||||||
handler: handle_ping_scan_event, machine_repository: IMachineRepository
|
|
||||||
):
|
|
||||||
event = PingScanEvent(source=EVENT.source, target=EVENT.target, response_received=True, os=None)
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
|
|
||||||
|
|
||||||
handler(event)
|
|
||||||
|
|
||||||
assert not machine_repository.upsert_machine.called
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_scan_data__machine_not_upserted_if_existing_machine_has_os(
|
|
||||||
handler: handle_ping_scan_event, machine_repository: IMachineRepository
|
|
||||||
):
|
|
||||||
machine_with_os = TARGET_MACHINE
|
|
||||||
machine_with_os.operating_system = OperatingSystem.WINDOWS
|
|
||||||
machine_repository.get_machine_by_ip = MagicMock(return_value=machine_with_os)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
assert not machine_repository.upsert_machine.called
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_scan_data__node_not_upserted_if_machine_storageerror(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
):
|
|
||||||
target_machine = TARGET_MACHINE
|
|
||||||
target_machine.operating_system = None
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
|
|
||||||
machine_repository.upsert_machine = MagicMock(side_effect=StorageError)
|
|
||||||
|
|
||||||
handler(EVENT)
|
|
||||||
|
|
||||||
assert not node_repository.upsert_communication.called
|
|
||||||
|
|
||||||
|
|
||||||
def test_handle_scan_data__failed_ping(
|
|
||||||
handler: handle_ping_scan_event,
|
|
||||||
machine_repository: IMachineRepository,
|
|
||||||
node_repository: INodeRepository,
|
|
||||||
):
|
|
||||||
machine_repository.upsert_machine = MagicMock(side_effect=StorageError)
|
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
|
||||||
machine_repository.get_machines_by_ip = MagicMock(side_effect=machines_from_ip)
|
|
||||||
|
|
||||||
handler(EVENT_NO_RESPONSE)
|
|
||||||
|
|
||||||
assert not node_repository.upsert_communication.called
|
|
||||||
assert not machine_repository.upsert_machine.called
|
|
Loading…
Reference in New Issue