Merge pull request #2373 from guardicore/2267-tcp-scan-event

2267 tcp scan event

monkey/common/agent_events/__init__.py

@@ -1,3 +1,4 @@
 from .abstract_agent_event import AbstractAgentEvent
 from .credentials_stolen_events import CredentialsStolenEvent
 from .ping_scan_event import PingScanEvent
+from .tcp_scan_event import TCPScanEvent

monkey/common/agent_events/tcp_scan_event.py

@@ -0,0 +1,18 @@
+from ipaddress import IPv4Address
+from typing import Dict
+
+from common.types import NetworkPort, PortStatus
+
+from . import AbstractAgentEvent
+
+
+class TCPScanEvent(AbstractAgentEvent):
+    """
+    An event that occurs when the Agent performs a TCP scan on a host
+
+    Attributes:
+        :param ports: The scanned ports and their status (open/closed)
+    """
+
+    target: IPv4Address
+    ports: Dict[NetworkPort, PortStatus]

monkey/common/types.py

@@ -6,7 +6,7 @@ from ipaddress import IPv4Address
 from typing import Optional
 from uuid import UUID
 
-from pydantic import PositiveInt, conint
+from pydantic import ConstrainedInt, PositiveInt
 from typing_extensions import TypeAlias
 
 from common import OperatingSystem
@@ -18,6 +18,18 @@ HardwareID: TypeAlias = PositiveInt
 MachineID: TypeAlias = PositiveInt
 
 
+class NetworkPort(ConstrainedInt):
+    """
+    Define network port as constrainer integer.
+
+    To define a default value with this type:
+        port: NetworkPort = typing.cast(NetworkPort, 1000)
+    """
+
+    ge = 1
+    le = 65535
+
+
 @dataclass
 class PingScanData:
     response_received: bool
@@ -29,16 +41,16 @@ class PortStatus(Enum):
     An Enum representing the status of the port.
 
     This Enum represents the status of a network pork. The value of each
-    member is distincive and unique number.
+    member is the member's name in all lower-case characters.
     """
 
-    OPEN = 1
-    CLOSED = 2
+    OPEN = "open"
+    CLOSED = "closed"
 
 
 class SocketAddress(InfectionMonkeyBaseModel):
     ip: IPv4Address
-    port: conint(ge=1, le=65535)  # type: ignore[valid-type]
+    port: NetworkPort
 
     @classmethod
     def from_string(cls, address_str: str) -> SocketAddress:

monkey/tests/unit_tests/common/agent_events/test_tcp_scan_event.py

@@ -0,0 +1,95 @@
+from ipaddress import IPv4Address
+from uuid import UUID
+
+import pytest
+
+from common.agent_events import TCPScanEvent
+from common.types import PortStatus
+
+TARGET_IP_STR = "192.168.1.10"
+AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10")
+
+TCP_SCAN_EVENT = TCPScanEvent(
+    source=AGENT_ID,
+    timestamp=1664371327.4067292,
+    target=IPv4Address(TARGET_IP_STR),
+    ports={
+        22: PortStatus.OPEN,
+        80: PortStatus.CLOSED,
+        443: PortStatus.OPEN,
+        8080: PortStatus.CLOSED,
+    },
+)
+
+TCP_OBJECT_DICT = {
+    "source": AGENT_ID,
+    "timestamp": 1664371327.4067292,
+    "target": IPv4Address(TARGET_IP_STR),
+    "ports": {
+        22: PortStatus.OPEN,
+        80: PortStatus.CLOSED,
+        443: PortStatus.OPEN,
+        8080: PortStatus.CLOSED,
+    },
+}
+
+TCP_SIMPLE_DICT = {
+    "source": str(AGENT_ID),
+    "timestamp": 1664371327.4067292,
+    "target": TARGET_IP_STR,
+    "ports": {
+        22: "open",
+        80: "closed",
+        443: "open",
+        8080: "closed",
+    },
+}
+
+
+@pytest.mark.parametrize("tcp_event_dict", [TCP_OBJECT_DICT, TCP_SIMPLE_DICT])
+def test_constructor(tcp_event_dict):
+    assert TCPScanEvent(**tcp_event_dict) == TCP_SCAN_EVENT
+
+
+def test_to_dict():
+    TCP_SCAN_EVENT.dict(simplify=True) == TCP_SIMPLE_DICT
+
+
+@pytest.mark.parametrize(
+    "key, value",
+    [
+        ("target", None),
+        ("ports", "not-a-dict"),
+        ("ports", {"not-a-number": "open"}),
+        ("ports", {22: "bogus"}),
+    ],
+)
+def test_construct_invalid_field__type_error(key, value):
+    invalid_type_dict = TCP_SIMPLE_DICT.copy()
+    invalid_type_dict[key] = value
+
+    with pytest.raises(TypeError):
+        TCPScanEvent(**invalid_type_dict)
+
+
+@pytest.mark.parametrize(
+    "key, value",
+    [
+        ("target", "not-an-ip"),
+        ("ports", {99999: "closed"}),
+    ],
+)
+def test_construct_invalid_field__value_error(key, value):
+    invalid_type_dict = TCP_SIMPLE_DICT.copy()
+    invalid_type_dict[key] = value
+
+    with pytest.raises(ValueError):
+        TCPScanEvent(**invalid_type_dict)
+
+
+def test_construct__extra_fields_forbidden():
+    extra_field_dict = TCP_SIMPLE_DICT.copy()
+    extra_field_dict["extra_field"] = 99  # red balloons
+
+    with pytest.raises(ValueError):
+        TCPScanEvent(**extra_field_dict)

vulture_allowlist.py

@@ -7,8 +7,9 @@ from common.agent_configuration.agent_sub_configurations import (
     CustomPBAConfiguration,
     ScanTargetConfiguration,
 )
-from common.agent_events import PingScanEvent
+from common.agent_events import PingScanEvent, TCPScanEvent
 from common.credentials import Credentials, LMHash, NTHash
+from common.types import NetworkPort
 from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory
 from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
 from monkey_island.cc.models import Report
@@ -29,6 +30,8 @@ from monkey_island.cc.repository.zero_trust.IEventRepository import IEventReposi
 from monkey_island.cc.repository.zero_trust.IFindingRepository import IFindingRepository
 from monkey_island.cc.services import AgentSignalsService
 
+NetworkPort.ge  # unused vairable (monkey/common/types.py:28)
+NetworkPort.le  # unused variable (monkey/common/types.py:29)
 fake_monkey_dir_path  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
 set_os_linux  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
 fake_monkey_dir_path  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:57)
@@ -311,6 +314,10 @@ IAgentLogRepository.get_agent_log
 # TODO: Remove once #2268 is closed
 PingScanEvent
 
+# TODO: Remove once #2267 is closed
+TCPScanEvent
+TCPScanEvent.port_status
+
 # pydantic base models
 underscore_attrs_are_private
 extra