Island: Fix ping scan handler to add machines IP

This commit is contained in:
vakarisz 2022-09-29 17:50:25 +03:00 committed by Kekoa Kaaikala
parent ff2b04c703
commit 9b30770777
2 changed files with 12 additions and 14 deletions

View File

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

View File

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