Island: Reuse the same open port logic in scan_event_handler.py

This commit is contained in:
vakarisz 2022-10-07 10:50:24 +03:00
parent d0d37ce595
commit a143d7206e
1 changed files with 4 additions and 5 deletions

View File

@ -1,4 +1,3 @@
from copy import deepcopy
from ipaddress import IPv4Interface from ipaddress import IPv4Interface
from logging import getLogger from logging import getLogger
from typing import List, Union from typing import List, Union
@ -6,7 +5,7 @@ from typing import List, Union
from typing_extensions import TypeAlias from typing_extensions import TypeAlias
from common.agent_events import PingScanEvent, TCPScanEvent 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.models import CommunicationType, Machine, Node
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
IAgentRepository, IAgentRepository,
@ -50,7 +49,7 @@ class ScanEventHandler:
logger.exception("Unable to process ping scan data") logger.exception("Unable to process ping scan data")
def handle_tcp_scan_event(self, event: TCPScanEvent): 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: if num_open_ports <= 0:
return return
@ -99,7 +98,7 @@ class ScanEventHandler:
@staticmethod @staticmethod
def _get_open_ports(event: TCPScanEvent) -> List[int]: 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): def _update_nodes(self, target_machine: Machine, event: ScanEvent):
src_machine = self._get_source_machine(event) 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): def _update_tcp_connections(self, src_node: Node, target_machine: Machine, event: TCPScanEvent):
tcp_connections = set() 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: for open_port in open_ports:
socket_address = SocketAddress(ip=event.target, port=open_port) socket_address = SocketAddress(ip=event.target, port=open_port)
tcp_connections.add(socket_address) tcp_connections.add(socket_address)