Island: Call get_or_create_target_machine() from ScanEventHandler

This commit is contained in:
Mike Salvatore 2022-10-05 15:01:40 -04:00
parent c4052bc5ad
commit e1f32177e9
1 changed files with 3 additions and 11 deletions

View File

@ -1,4 +1,3 @@
from ipaddress import IPv4Interface
from logging import getLogger from logging import getLogger
from typing import List, Union from typing import List, Union
@ -16,6 +15,8 @@ from monkey_island.cc.repository import (
UnknownRecordError, UnknownRecordError,
) )
from .utils import get_or_create_target_machine
ScanEvent: TypeAlias = Union[PingScanEvent, TCPScanEvent] ScanEvent: TypeAlias = Union[PingScanEvent, TCPScanEvent]
logger = getLogger(__name__) logger = getLogger(__name__)
@ -65,16 +66,7 @@ class ScanEventHandler:
logger.exception("Unable to process tcp scan data") logger.exception("Unable to process tcp scan data")
def _get_target_machine(self, event: ScanEvent) -> Machine: def _get_target_machine(self, event: ScanEvent) -> Machine:
try: return get_or_create_target_machine(self._machine_repository, event.target)
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 _get_source_node(self, event: ScanEvent) -> Node: def _get_source_node(self, event: ScanEvent) -> Node:
machine = self._get_source_machine(event) machine = self._get_source_machine(event)