Agent: Move PingScanData to common

monkey/common/types.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from ipaddress import IPv4Address
+from typing import Optional
 from uuid import UUID
 
 from pydantic import PositiveInt, conint
 from typing_extensions import TypeAlias
 
+from common import OperatingSystem
 from common.base_models import InfectionMonkeyBaseModel
 from common.network.network_utils import address_to_ip_port
 
@@ -14,6 +17,12 @@ HardwareID: TypeAlias = PositiveInt
 MachineID: TypeAlias = PositiveInt
 
 
+@dataclass
+class PingScanData:
+    response_received: bool
+    os: Optional[OperatingSystem]
+
+
 class SocketAddress(InfectionMonkeyBaseModel):
     ip: IPv4Address
     port: conint(ge=1, le=65535)  # type: ignore[valid-type]

monkey/infection_monkey/i_puppet/__init__.py

@@ -2,7 +2,6 @@ from .plugin_type import PluginType
 from .i_puppet import (
     IPuppet,
     ExploiterResultData,
-    PingScanData,
     PortScanData,
     FingerprintData,
     PortStatus,

monkey/infection_monkey/i_puppet/i_fingerprinter.py

@@ -1,7 +1,9 @@
 from abc import abstractmethod
 from typing import Dict
 
-from . import FingerprintData, PingScanData, PortScanData
+from common.types import PingScanData
+
+from . import FingerprintData, PortScanData
 
 
 class IFingerprinter:

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -5,8 +5,8 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, Iterable, Mapping, Optional, Sequence
 
-from common import OperatingSystem
 from common.credentials import Credentials
+from common.types import PingScanData
 from infection_monkey.model import VictimHost
 
 from . import PluginType
@@ -32,12 +32,6 @@ class ExploiterResultData:
     error_message: str = ""
 
 
-@dataclass
-class PingScanData:
-    response_received: bool
-    os: Optional[OperatingSystem]
-
-
 PortScanData = namedtuple("PortScanData", ["port", "status", "banner", "service"])
 FingerprintData = namedtuple("FingerprintData", ["os_type", "os_version", "services"])
 PostBreachData = namedtuple("PostBreachData", ["display_name", "command", "result"])

monkey/infection_monkey/master/ip_scan_results.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
 from typing import Dict
 
-from infection_monkey.i_puppet import FingerprintData, PingScanData, PortScanData
+from common.types import PingScanData
+from infection_monkey.i_puppet import FingerprintData, PortScanData
 
 Port = int
 FingerprinterName = str

monkey/infection_monkey/master/ip_scanner.py

@@ -9,13 +9,8 @@ from common.agent_configuration.agent_sub_configurations import (
     NetworkScanConfiguration,
     PluginConfiguration,
 )
-from infection_monkey.i_puppet import (
+from common.types import PingScanData
-    FingerprintData,
+from infection_monkey.i_puppet import FingerprintData, IPuppet, PortScanData, PortStatus
-    IPuppet,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
 from infection_monkey.network import NetworkAddress
 from infection_monkey.utils.threading import interruptible_iter, run_worker_threads
 

monkey/infection_monkey/master/propagator.py

@@ -10,13 +10,8 @@ from common.agent_configuration import (
     PropagationConfiguration,
     ScanTargetConfiguration,
 )
-from infection_monkey.i_puppet import (
+from common.types import PingScanData
-    ExploiterResultData,
+from infection_monkey.i_puppet import ExploiterResultData, FingerprintData, PortScanData, PortStatus
-    FingerprintData,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
 from infection_monkey.model import VictimHost, VictimHostFactory
 from infection_monkey.network import NetworkAddress
 from infection_monkey.network_scanning.scan_target_generator import compile_scan_target_list

monkey/infection_monkey/network_scanning/elasticsearch_fingerprinter.py

@@ -5,13 +5,8 @@ from typing import Any, Dict
 import requests
 
 from common.common_consts.network_consts import ES_SERVICE
-from infection_monkey.i_puppet import (
+from common.types import PingScanData
-    FingerprintData,
+from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData, PortStatus
-    IFingerprinter,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
 
 DISPLAY_NAME = "ElasticSearch"
 ES_PORT = 9200

monkey/infection_monkey/network_scanning/http_fingerprinter.py

@@ -5,13 +5,8 @@ from typing import Any, Dict, Iterable, Optional, Set, Tuple
 from requests import head
 from requests.exceptions import ConnectionError, Timeout
 
-from infection_monkey.i_puppet import (
+from common.types import PingScanData
-    FingerprintData,
+from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData, PortStatus
-    IFingerprinter,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
 
 logger = logging.getLogger(__name__)
 

monkey/infection_monkey/network_scanning/mssql_fingerprinter.py

@@ -3,7 +3,8 @@ import logging
 import socket
 from typing import Any, Dict, Optional
 
-from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PingScanData, PortScanData
+from common.types import PingScanData
+from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData
 
 MSSQL_SERVICE = "MSSQL"
 DISPLAY_NAME = MSSQL_SERVICE

monkey/infection_monkey/network_scanning/ping_scanner.py

@@ -6,7 +6,7 @@ import subprocess
 import sys
 
 from common import OperatingSystem
-from infection_monkey.i_puppet import PingScanData
+from common.types import PingScanData
 from infection_monkey.utils.environment import is_windows_os
 
 TTL_REGEX = re.compile(r"TTL=([0-9]+)\b", re.IGNORECASE)

monkey/infection_monkey/network_scanning/smb_fingerprinter.py

@@ -6,13 +6,8 @@ from typing import Dict
 from odict import odict
 
 from common import OperatingSystem
-from infection_monkey.i_puppet import (
+from common.types import PingScanData
-    FingerprintData,
+from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData, PortStatus
-    IFingerprinter,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
 
 DISPLAY_NAME = "SMB"
 SMB_PORT = 445

monkey/infection_monkey/network_scanning/ssh_fingerprinter.py

@@ -2,7 +2,8 @@ import re
 from typing import Dict, Optional, Tuple
 
 from common import OperatingSystem
-from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PingScanData, PortScanData
+from common.types import PingScanData
+from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData
 
 SSH_REGEX = r"SSH-\d\.\d-OpenSSH"
 LINUX_DIST_SSH = ["ubuntu", "debian"]

monkey/infection_monkey/puppet/puppet.py

@@ -4,12 +4,12 @@ from typing import Dict, Iterable, Sequence
 
 from common.common_consts.timeouts import CONNECTION_TIMEOUT
 from common.credentials import Credentials
+from common.types import PingScanData
 from infection_monkey import network_scanning
 from infection_monkey.i_puppet import (
     ExploiterResultData,
     FingerprintData,
     IPuppet,
-    PingScanData,
     PluginType,
     PortScanData,
     PostBreachData,

monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py

@@ -1,14 +1,14 @@
 import logging
 import threading
-from typing import Dict, Iterable, List, Sequence
+from typing import Dict, Iterable, Sequence
 
 from common import OperatingSystem
 from common.credentials import Credentials, LMHash, Password, SSHKeypair, Username
+from common.types import PingScanData
 from infection_monkey.i_puppet import (
     ExploiterResultData,
     FingerprintData,
     IPuppet,
-    PingScanData,
     PluginType,
     PortScanData,
     PortStatus,
@@ -25,26 +25,34 @@ logger = logging.getLogger()
 
 
 class MockPuppet(IPuppet):
-    def load_plugin(self, plugin: object, plugin_type: PluginType) -> None:
+    def load_plugin(self, plugin_name: str, plugin: object, plugin_type: PluginType) -> None:
         logger.debug(f"load_plugin({plugin}, {plugin_type})")
 
     def run_credential_collector(self, name: str, options: Dict) -> Sequence[Credentials]:
         logger.debug(f"run_credential_collector({name})")
 
         if name == "SSHCollector":
-            ssh_credentials = Credentials(
+            ssh_credentials = [
-                [Username("m0nk3y")],
+                Credentials(
-                [
+                    identity=Username(username="m0nk3y"),
-                    SSHKeypair("Public_Key_0", "Private_Key_0"),
+                    secret=SSHKeypair(private_key="Public_Key_0", public_key="Private_Key_0"),
-                    SSHKeypair("Public_Key_1", "Private_Key_1"),
+                ),
-                ],
+                Credentials(
-            )
+                    identity=Username(username="m0nk3y"),
-            return [ssh_credentials]
+                    secret=SSHKeypair(private_key="Public_Key_1", public_key="Private_Key_1"),
+                ),
+            ]
+            return ssh_credentials
         elif name == "MimikatzCollector":
-            windows_credentials = Credentials(
+            windows_credentials = [
-                [Username("test_user")], [Password("1234"), LMHash("DEADBEEF")]
+                Credentials(
-            )
+                    identity=Username(username="test_user"), secret=Password(password="1234")
-            return [windows_credentials]
+                ),
+                Credentials(
+                    identity=Username(username="test_user"), secret=LMHash(lm_hash="DEADBEEF")
+                ),
+            ]
+            return windows_credentials
 
         return []
 
@@ -59,13 +67,13 @@ class MockPuppet(IPuppet):
     def ping(self, host: str, timeout: float = 1) -> PingScanData:
         logger.debug(f"run_ping({host}, {timeout})")
         if host == DOT_1:
-            return PingScanData(True, "windows")
+            return PingScanData(True, OperatingSystem.WINDOWS)
 
         if host == DOT_2:
             return PingScanData(False, None)
 
         if host == DOT_3:
-            return PingScanData(True, "linux")
+            return PingScanData(True, OperatingSystem.LINUX)
 
         if host == DOT_4:
             return PingScanData(False, None)
@@ -73,7 +81,7 @@ class MockPuppet(IPuppet):
         return PingScanData(False, None)
 
     def scan_tcp_ports(
-        self, host: str, ports: List[int], timeout: float = 3
+        self, host: str, ports: Sequence[int], timeout: float = 3
     ) -> Dict[int, PortScanData]:
         logger.debug(f"run_scan_tcp_port({host}, {ports}, {timeout})")
         dot_1_results = {
@@ -139,6 +147,7 @@ class MockPuppet(IPuppet):
         name: str,
         host: VictimHost,
         current_depth: int,
+        servers: Sequence[str],
         options: Dict,
         interrupt: threading.Event,
     ) -> ExploiterResultData:
@@ -186,19 +195,19 @@ class MockPuppet(IPuppet):
         successful_exploiters = {
             DOT_1: {
                 "ZerologonExploiter": ExploiterResultData(
-                    False, False, False, OperatingSystem.WINDOWS, {}, [], "Zerologon failed"
+                    False, False, False, OperatingSystem.WINDOWS.value, {}, [], "Zerologon failed"
                 ),
                 "SSHExploiter": ExploiterResultData(
                     False,
                     False,
                     False,
-                    OperatingSystem.LINUX,
+                    OperatingSystem.LINUX.value,
                     info_ssh,
                     attempts,
                     "Failed exploiting",
                 ),
                 "WmiExploiter": ExploiterResultData(
-                    True, True, False, OperatingSystem.WINDOWS, info_wmi, attempts, None
+                    True, True, False, OperatingSystem.WINDOWS.value, info_wmi, attempts
                 ),
             },
             DOT_3: {
@@ -206,7 +215,7 @@ class MockPuppet(IPuppet):
                     False,
                     False,
                     False,
-                    OperatingSystem.WINDOWS,
+                    OperatingSystem.WINDOWS.value,
                     info_wmi,
                     attempts,
                     "PowerShell Exploiter Failed",
@@ -215,13 +224,13 @@ class MockPuppet(IPuppet):
                     False,
                     False,
                     False,
-                    OperatingSystem.LINUX,
+                    OperatingSystem.LINUX.value,
                     info_ssh,
                     attempts,
                     "Failed exploiting",
                 ),
                 "ZerologonExploiter": ExploiterResultData(
-                    True, False, False, OperatingSystem.WINDOWS, {}, [], None
+                    True, False, False, OperatingSystem.WINDOWS.value, {}, []
                 ),
             },
         }
@@ -233,7 +242,7 @@ class MockPuppet(IPuppet):
                 False,
                 False,
                 False,
-                OperatingSystem.LINUX,
+                OperatingSystem.LINUX.value,
                 {},
                 [],
                 f"{name} failed for host {host}",

monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py

@@ -5,12 +5,14 @@ from unittest.mock import MagicMock
 import pytest
 from tests.unit_tests.infection_monkey.master.mock_puppet import MockPuppet
 
+from common import OperatingSystem
 from common.agent_configuration.agent_sub_configurations import (
     ICMPScanConfiguration,
     NetworkScanConfiguration,
     PluginConfiguration,
     TCPScanConfiguration,
 )
+from common.types import PingScanData
 from infection_monkey.i_puppet import FingerprintData, PortScanData, PortStatus
 from infection_monkey.master import IPScanner
 from infection_monkey.network import NetworkAddress
@@ -78,10 +80,15 @@ def assert_scan_results(address, scan_results):
         assert_scan_results_host_down(address, ping_scan_data, port_scan_data, fingerprint_data)
 
 
-def assert_scan_results_no_1(domain, ping_scan_data, port_scan_data, fingerprint_data):
+def assert_scan_results_no_1(
+    domain,
+    ping_scan_data: PingScanData,
+    port_scan_data: PortScanData,
+    fingerprint_data: FingerprintData,
+):
     assert domain == "d1"
     assert ping_scan_data.response_received is True
-    assert ping_scan_data.os == WINDOWS_OS
+    assert ping_scan_data.os == OperatingSystem.WINDOWS
 
     assert len(port_scan_data.keys()) == 6
 
@@ -100,7 +107,7 @@ def assert_scan_results_no_1(domain, ping_scan_data, port_scan_data, fingerprint
     assert_fingerprint_results_no_1(fingerprint_data)
 
 
-def assert_fingerprint_results_no_1(fingerprint_data):
+def assert_fingerprint_results_no_1(fingerprint_data: FingerprintData):
     assert len(fingerprint_data.keys()) == 3
     assert fingerprint_data["SSHFinger"].services == {}
     assert fingerprint_data["HTTPFinger"].services == {}
@@ -112,11 +119,16 @@ def assert_fingerprint_results_no_1(fingerprint_data):
     assert fingerprint_data["SMBFinger"].services["tcp-445"]["name"] == "smb_service_name"
 
 
-def assert_scan_results_no_3(domain, ping_scan_data, port_scan_data, fingerprint_data):
+def assert_scan_results_no_3(
+    domain,
+    ping_scan_data: PingScanData,
+    port_scan_data: PortScanData,
+    fingerprint_data: FingerprintData,
+):
     assert domain == "d3"
 
     assert ping_scan_data.response_received is True
-    assert ping_scan_data.os == LINUX_OS
+    assert ping_scan_data.os == OperatingSystem.LINUX
     assert len(port_scan_data.keys()) == 6
 
     psd_443 = port_scan_data[443]
@@ -134,7 +146,7 @@ def assert_scan_results_no_3(domain, ping_scan_data, port_scan_data, fingerprint
     assert_fingerprint_results_no_3(fingerprint_data)
 
 
-def assert_fingerprint_results_no_3(fingerprint_data):
+def assert_fingerprint_results_no_3(fingerprint_data: FingerprintData):
     assert len(fingerprint_data.keys()) == 3
     assert fingerprint_data["SMBFinger"].services == {}
 
@@ -152,7 +164,12 @@ def assert_fingerprint_results_no_3(fingerprint_data):
     assert fingerprint_data["HTTPFinger"].services["tcp-443"]["data"] == ("SERVER_HEADERS_2", True)
 
 
-def assert_scan_results_host_down(address, ping_scan_data, port_scan_data, fingerprint_data):
+def assert_scan_results_host_down(
+    address,
+    ping_scan_data: PingScanData,
+    port_scan_data: PortScanData,
+    fingerprint_data: FingerprintData,
+):
     assert address.ip not in {"10.0.0.1", "10.0.0.3"}
     assert address.domain is None
 

monkey/tests/unit_tests/infection_monkey/master/test_propagator.py

@@ -9,13 +9,8 @@ from common.agent_configuration.agent_sub_configurations import (
     PropagationConfiguration,
     ScanTargetConfiguration,
 )
-from infection_monkey.i_puppet import (
+from common.types import PingScanData
-    ExploiterResultData,
+from infection_monkey.i_puppet import ExploiterResultData, FingerprintData, PortScanData, PortStatus
-    FingerprintData,
-    PingScanData,
-    PortScanData,
-    PortStatus,
-)
 from infection_monkey.master import IPScanResults, Propagator
 from infection_monkey.model import VictimHost, VictimHostFactory
 from infection_monkey.network import NetworkAddress

monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py

@@ -1,7 +1,8 @@
 import threading
 from unittest.mock import MagicMock
 
-from infection_monkey.i_puppet import PingScanData, PluginType
+from common.types import PingScanData
+from infection_monkey.i_puppet import PluginType
 from infection_monkey.puppet.puppet import EMPTY_FINGERPRINT, Puppet