From a143d7206e1e7a095a0035172c744da6b9b590e5 Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 7 Oct 2022 10:50:24 +0300 Subject: [PATCH] Island: Reuse the same open port logic in scan_event_handler.py --- .../cc/agent_event_handlers/scan_event_handler.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 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 8e8d3cb25..272749807 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 copy import deepcopy from ipaddress import IPv4Interface from logging import getLogger from typing import List, Union @@ -6,7 +5,7 @@ from typing import List, Union from typing_extensions import TypeAlias from common.agent_events import PingScanEvent, TCPScanEvent -from common.types import PortStatus, SocketAddress +from common.types import NetworkService, PortStatus, SocketAddress from monkey_island.cc.models import CommunicationType, Machine, Node from monkey_island.cc.repository import ( IAgentRepository, @@ -50,7 +49,7 @@ class ScanEventHandler: logger.exception("Unable to process ping scan data") def handle_tcp_scan_event(self, event: TCPScanEvent): - num_open_ports = sum((1 for status in event.ports.values() if status == PortStatus.OPEN)) + num_open_ports = len(self._get_open_ports(event)) if num_open_ports <= 0: return @@ -99,7 +98,7 @@ class ScanEventHandler: @staticmethod def _get_open_ports(event: TCPScanEvent) -> List[int]: - return [port for port, status in event.ports.items() if status == PortStatus.OPEN] + return (port for port, status in event.ports.items() if status == PortStatus.OPEN) def _update_nodes(self, target_machine: Machine, event: ScanEvent): src_machine = self._get_source_machine(event) @@ -110,7 +109,7 @@ class ScanEventHandler: def _update_tcp_connections(self, src_node: Node, target_machine: Machine, event: TCPScanEvent): tcp_connections = set() - open_ports = (port for port, status in event.ports.items() if status == PortStatus.OPEN) + open_ports = self._get_open_ports(event) for open_port in open_ports: socket_address = SocketAddress(ip=event.target, port=open_port) tcp_connections.add(socket_address)