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 logging import getLogger
|
||||||
|
|
||||||
|
from pydantic.json import IPv4Interface
|
||||||
|
|
||||||
from common.agent_events import PingScanEvent
|
from common.agent_events import PingScanEvent
|
||||||
from monkey_island.cc.models import CommunicationType, Machine
|
from monkey_island.cc.models import CommunicationType, Machine
|
||||||
from monkey_island.cc.repository import (
|
from monkey_island.cc.repository import (
|
||||||
|
@ -15,6 +17,11 @@ logger = getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
class handle_ping_scan_event:
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
agent_repository: IAgentRepository,
|
agent_repository: IAgentRepository,
|
||||||
|
@ -46,7 +53,10 @@ class handle_ping_scan_event:
|
||||||
dest_machines = self._machine_repository.get_machines_by_ip(event.target)
|
dest_machines = self._machine_repository.get_machines_by_ip(event.target)
|
||||||
return dest_machines[0]
|
return dest_machines[0]
|
||||||
except UnknownRecordError:
|
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)
|
self._machine_repository.upsert_machine(machine)
|
||||||
return machine
|
return machine
|
||||||
|
|
||||||
|
|
|
@ -130,7 +130,7 @@ def test_handle_ping_scan_event__upserts_machine(
|
||||||
machine_repository: IMachineRepository,
|
machine_repository: IMachineRepository,
|
||||||
):
|
):
|
||||||
machine_repository.get_machine_by_id = MagicMock(side_effect=machine_from_id)
|
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)
|
handler(EVENT)
|
||||||
|
|
||||||
expected_machine = TARGET_MACHINE.copy()
|
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
|
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])
|
@pytest.mark.parametrize("id", [PINGER_MACHINE.id, TARGET_MACHINE.id])
|
||||||
def test_handle_scan_data__node_not_upserted_if_machine_retrievalerror(
|
def test_handle_scan_data__node_not_upserted_if_machine_retrievalerror(
|
||||||
handler: handle_ping_scan_event,
|
handler: handle_ping_scan_event,
|
||||||
|
|
Loading…
Reference in New Issue