From aeb6630ebcf2a4f0026cf0482d355f3d1748cf65 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 18:52:23 +0000 Subject: [PATCH] Agent: Move PingScanData to common --- monkey/common/types.py | 9 +++ monkey/infection_monkey/i_puppet/__init__.py | 1 - .../i_puppet/i_fingerprinter.py | 4 +- monkey/infection_monkey/i_puppet/i_puppet.py | 8 +-- .../master/ip_scan_results.py | 3 +- monkey/infection_monkey/master/ip_scanner.py | 9 +-- monkey/infection_monkey/master/propagator.py | 9 +-- .../elasticsearch_fingerprinter.py | 9 +-- .../network_scanning/http_fingerprinter.py | 9 +-- .../network_scanning/mssql_fingerprinter.py | 3 +- .../network_scanning/ping_scanner.py | 2 +- .../network_scanning/smb_fingerprinter.py | 9 +-- .../network_scanning/ssh_fingerprinter.py | 3 +- monkey/infection_monkey/puppet/puppet.py | 2 +- .../infection_monkey/master/mock_puppet.py | 59 +++++++++++-------- .../master/test_ip_scanner.py | 31 +++++++--- .../master/test_propagator.py | 9 +-- .../infection_monkey/puppet/test_puppet.py | 3 +- 18 files changed, 93 insertions(+), 89 deletions(-) diff --git a/monkey/common/types.py b/monkey/common/types.py index 5f86d5060..0c14aeec9 100644 --- a/monkey/common/types.py +++ b/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] diff --git a/monkey/infection_monkey/i_puppet/__init__.py b/monkey/infection_monkey/i_puppet/__init__.py index 5aeb265f3..d38183700 100644 --- a/monkey/infection_monkey/i_puppet/__init__.py +++ b/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, diff --git a/monkey/infection_monkey/i_puppet/i_fingerprinter.py b/monkey/infection_monkey/i_puppet/i_fingerprinter.py index e6f177021..748b8b10d 100644 --- a/monkey/infection_monkey/i_puppet/i_fingerprinter.py +++ b/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: diff --git a/monkey/infection_monkey/i_puppet/i_puppet.py b/monkey/infection_monkey/i_puppet/i_puppet.py index 73afa1227..57e62c5bd 100644 --- a/monkey/infection_monkey/i_puppet/i_puppet.py +++ b/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"]) diff --git a/monkey/infection_monkey/master/ip_scan_results.py b/monkey/infection_monkey/master/ip_scan_results.py index 98f7b6646..06bf7cd2f 100644 --- a/monkey/infection_monkey/master/ip_scan_results.py +++ b/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 diff --git a/monkey/infection_monkey/master/ip_scanner.py b/monkey/infection_monkey/master/ip_scanner.py index 6f25f1f4b..d1534a6f5 100644 --- a/monkey/infection_monkey/master/ip_scanner.py +++ b/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 ( - FingerprintData, - IPuppet, - PingScanData, - PortScanData, - PortStatus, -) +from common.types import PingScanData +from infection_monkey.i_puppet import FingerprintData, IPuppet, PortScanData, PortStatus from infection_monkey.network import NetworkAddress from infection_monkey.utils.threading import interruptible_iter, run_worker_threads diff --git a/monkey/infection_monkey/master/propagator.py b/monkey/infection_monkey/master/propagator.py index 9b14bef2a..bf9e830e7 100644 --- a/monkey/infection_monkey/master/propagator.py +++ b/monkey/infection_monkey/master/propagator.py @@ -10,13 +10,8 @@ from common.agent_configuration import ( PropagationConfiguration, ScanTargetConfiguration, ) -from infection_monkey.i_puppet import ( - ExploiterResultData, - FingerprintData, - PingScanData, - PortScanData, - PortStatus, -) +from common.types import PingScanData +from infection_monkey.i_puppet import ExploiterResultData, FingerprintData, 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 diff --git a/monkey/infection_monkey/network_scanning/elasticsearch_fingerprinter.py b/monkey/infection_monkey/network_scanning/elasticsearch_fingerprinter.py index cbdce5812..0d9b69e9e 100644 --- a/monkey/infection_monkey/network_scanning/elasticsearch_fingerprinter.py +++ b/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 ( - FingerprintData, - IFingerprinter, - PingScanData, - PortScanData, - PortStatus, -) +from common.types import PingScanData +from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData, PortStatus DISPLAY_NAME = "ElasticSearch" ES_PORT = 9200 diff --git a/monkey/infection_monkey/network_scanning/http_fingerprinter.py b/monkey/infection_monkey/network_scanning/http_fingerprinter.py index 4c738e1ba..b2238217a 100644 --- a/monkey/infection_monkey/network_scanning/http_fingerprinter.py +++ b/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 ( - FingerprintData, - IFingerprinter, - PingScanData, - PortScanData, - PortStatus, -) +from common.types import PingScanData +from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData, PortStatus logger = logging.getLogger(__name__) diff --git a/monkey/infection_monkey/network_scanning/mssql_fingerprinter.py b/monkey/infection_monkey/network_scanning/mssql_fingerprinter.py index 18630aaa7..dee742101 100644 --- a/monkey/infection_monkey/network_scanning/mssql_fingerprinter.py +++ b/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 diff --git a/monkey/infection_monkey/network_scanning/ping_scanner.py b/monkey/infection_monkey/network_scanning/ping_scanner.py index c181d784b..13cafe7d1 100644 --- a/monkey/infection_monkey/network_scanning/ping_scanner.py +++ b/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) diff --git a/monkey/infection_monkey/network_scanning/smb_fingerprinter.py b/monkey/infection_monkey/network_scanning/smb_fingerprinter.py index b7293bf73..a80aaf6b4 100644 --- a/monkey/infection_monkey/network_scanning/smb_fingerprinter.py +++ b/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 ( - FingerprintData, - IFingerprinter, - PingScanData, - PortScanData, - PortStatus, -) +from common.types import PingScanData +from infection_monkey.i_puppet import FingerprintData, IFingerprinter, PortScanData, PortStatus DISPLAY_NAME = "SMB" SMB_PORT = 445 diff --git a/monkey/infection_monkey/network_scanning/ssh_fingerprinter.py b/monkey/infection_monkey/network_scanning/ssh_fingerprinter.py index e17271817..c28570fdd 100644 --- a/monkey/infection_monkey/network_scanning/ssh_fingerprinter.py +++ b/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"] diff --git a/monkey/infection_monkey/puppet/puppet.py b/monkey/infection_monkey/puppet/puppet.py index 151783653..80293d44d 100644 --- a/monkey/infection_monkey/puppet/puppet.py +++ b/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, diff --git a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py b/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py index 67b4fa30f..33ae899ef 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/mock_puppet.py +++ b/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( - [Username("m0nk3y")], - [ - SSHKeypair("Public_Key_0", "Private_Key_0"), - SSHKeypair("Public_Key_1", "Private_Key_1"), - ], - ) - return [ssh_credentials] + ssh_credentials = [ + Credentials( + identity=Username(username="m0nk3y"), + secret=SSHKeypair(private_key="Public_Key_0", public_key="Private_Key_0"), + ), + Credentials( + identity=Username(username="m0nk3y"), + secret=SSHKeypair(private_key="Public_Key_1", public_key="Private_Key_1"), + ), + ] + return ssh_credentials elif name == "MimikatzCollector": - windows_credentials = Credentials( - [Username("test_user")], [Password("1234"), LMHash("DEADBEEF")] - ) - return [windows_credentials] + windows_credentials = [ + Credentials( + identity=Username(username="test_user"), secret=Password(password="1234") + ), + 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}", diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py b/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py index c55fdb85c..3175b0f63 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_ip_scanner.py +++ b/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 diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py b/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py index 8f1b51274..b662462ac 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_propagator.py +++ b/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 ( - ExploiterResultData, - FingerprintData, - PingScanData, - PortScanData, - PortStatus, -) +from common.types import PingScanData +from infection_monkey.i_puppet import ExploiterResultData, FingerprintData, PortScanData, PortStatus from infection_monkey.master import IPScanResults, Propagator from infection_monkey.model import VictimHost, VictimHostFactory from infection_monkey.network import NetworkAddress diff --git a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py b/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py index 39273faee..a0df06fd6 100644 --- a/monkey/tests/unit_tests/infection_monkey/puppet/test_puppet.py +++ b/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