p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Remove handle_ping_scan_event callable class

This commit is contained in:
Ilija Lazoroski 2022-09-30 16:20:23 +02:00
parent 8002080c8b
commit f23093dc78
3 changed files with 0 additions and 345 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
monkey/monkey_island/cc/agent_event_handlers/__init__.py
View File

@ -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_stolen_credentials_to_repository import save_stolen_credentials_to_repository
from .scan_event_handler import ScanEventHandler

76
monkey/monkey_island/cc/agent_event_handlers/handle_ping_scan_event.py
View File

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

268
monkey/tests/unit_tests/monkey_island/cc/agent_event_handlers/test_handle_ping_scan_event.py
View File

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