From e1f32177e9dc3b145ffbb131caf1be78ef61eaea Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 5 Oct 2022 15:01:40 -0400 Subject: [PATCH] Island: Call get_or_create_target_machine() from ScanEventHandler --- .../cc/agent_event_handlers/scan_event_handler.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py b/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py index b79fc2558..b7c222830 100644 --- a/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py +++ b/monkey/monkey_island/cc/agent_event_handlers/scan_event_handler.py @@ -1,4 +1,3 @@ -from ipaddress import IPv4Interface from logging import getLogger from typing import List, Union @@ -16,6 +15,8 @@ from monkey_island.cc.repository import ( UnknownRecordError, ) +from .utils import get_or_create_target_machine + ScanEvent: TypeAlias = Union[PingScanEvent, TCPScanEvent] logger = getLogger(__name__) @@ -65,16 +66,7 @@ class ScanEventHandler: logger.exception("Unable to process tcp scan data") def _get_target_machine(self, event: ScanEvent) -> 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 + return get_or_create_target_machine(self._machine_repository, event.target) def _get_source_node(self, event: ScanEvent) -> Node: machine = self._get_source_machine(event)