p41275903
/
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

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