forked from p34709852/monkey
Island: Fix ping scan handler to add machines IP
This commit is contained in:
parent
ff2b04c703
commit
9b30770777
|
@ -1,5 +1,7 @@
|
|||
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 (
|
||||
|
@ -15,6 +17,11 @@ logger = getLogger(__name__)
|
|||
|
||||
|
||||
class handle_ping_scan_event:
|
||||
"""
|
||||
Handles ping scan event and makes changes to Machine and Node states based on it
|
||||
:param event: Ping scan event
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent_repository: IAgentRepository,
|
||||
|
@ -46,7 +53,10 @@ class handle_ping_scan_event:
|
|||
dest_machines = self._machine_repository.get_machines_by_ip(event.target)
|
||||
return dest_machines[0]
|
||||
except UnknownRecordError:
|
||||
machine = Machine(id=self._machine_repository.get_new_id())
|
||||
machine = Machine(
|
||||
id=self._machine_repository.get_new_id(),
|
||||
network_interfaces=[IPv4Interface(event.target)],
|
||||
)
|
||||
self._machine_repository.upsert_machine(machine)
|
||||
return machine
|
||||
|
||||
|
|
|
@ -130,7 +130,7 @@ def test_handle_ping_scan_event__upserts_machine(
|
|||
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)
|
||||
machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError)
|
||||
handler(EVENT)
|
||||
|
||||
expected_machine = TARGET_MACHINE.copy()
|
||||
|
@ -193,18 +193,6 @@ def test_handle_ping_scan_event__node_not_upserted_if_no_matching_machine(
|
|||
assert not node_repository.upsert_communication.called
|
||||
|
||||
|
||||
def test_handle_ping_scan_event__upserts_machine_if_not_existed(
|
||||
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, operating_system=OperatingSystem.LINUX)
|
||||
|
||||
machine_repository.upsert_machine.assert_called_with(expected_machine)
|
||||
|
||||
|
||||
@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,
|
||||
|
|
Loading…
Reference in New Issue