p17860254
/
monkey
forked from p15670423/monkey
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Island: Extract methods in handle_scan_data

This commit is contained in:
Kekoa Kaaikala 2022-09-27 18:58:13 +00:00
parent eacd426969
commit 89c6e2b7bc
1 changed files with 25 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
monkey/monkey_island/cc/agent_event_handlers/handle_scan_data.py
View File

@ -28,30 +28,34 @@ class handle_scan_data:
def __call__(self, event: PingScanEvent):
try:
# Get or create the destination machine
# NOTE: Assuming IP's are unique for now
if not isinstance(event.target, IPv4Address):
logger.error("Unable to process scan data: Unknown target")
return
dest_machines = self._machine_repository.get_machines_by_ip(event.target)
if not dest_machines:
machine = Machine(id=self._machine_repository.get_new_id())
dest_machines = [machine]
self._machine_repository.upsert_machine(machine)
# Update the destination machine
dest_machine = dest_machines[0]
if event.scan_data.os is not None:
dest_machine.operating_system = event.scan_data.os
self._machine_repository.upsert_machine(dest_machine)
# Get the source machine
agent = self._agent_repository.get_agent_by_id(event.source)
src_machine = self._machine_repository.get_machine_by_id(agent.machine_id)
dest_machine = self._get_destination_machine(event)
self._update_destination_machine(dest_machine, event)
src_machine = self._get_source_machine(event)
# Update or create the node
self._node_repository.upsert_communication(
src_machine.id, dest_machine.id, CommunicationType.SCANNED
)
except (RetrievalError, StorageError, UnknownRecordError) as err:
except (RetrievalError, StorageError, TypeError, UnknownRecordError) as err:
logger.error(f"Unable to process scan data: {err}")
def _get_destination_machine(self, event: PingScanEvent) -> Machine:
# NOTE: Assuming IP's are unique for now
if not isinstance(event.target, IPv4Address):
raise TypeError("Unknown target")
dest_machines = self._machine_repository.get_machines_by_ip(event.target)
if not dest_machines:
machine = Machine(id=self._machine_repository.get_new_id())
dest_machines = [machine]
self._machine_repository.upsert_machine(machine)
return dest_machines[0]
def _get_source_machine(self, event: PingScanEvent) -> Machine:
agent = self._agent_repository.get_agent_by_id(event.source)
return self._machine_repository.get_machine_by_id(agent.machine_id)
def _update_destination_machine(self, machine: Machine, event: PingScanEvent):
if event.scan_data.os is not None:
machine.operating_system = event.scan_data.os
self._machine_repository.upsert_machine(machine)