Merge branch '1597-integrate-automated-master' into agent-refactor
This commit is contained in:
commit
e392915b26
|
@ -1,4 +1,5 @@
|
|||
import re
|
||||
from typing import Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
||||
|
@ -20,3 +21,11 @@ def remove_port(url):
|
|||
with_port = f"{parsed.scheme}://{parsed.netloc}"
|
||||
without_port = re.sub(":[0-9]+(?=$|/)", "", with_port)
|
||||
return without_port
|
||||
|
||||
|
||||
def address_to_ip_port(address: str) -> Tuple[str, Optional[str]]:
|
||||
if ":" in address:
|
||||
ip, port = address.split(":")
|
||||
return ip, port or None
|
||||
else:
|
||||
return address, None
|
||||
|
|
|
@ -7,6 +7,7 @@ from infection_monkey.i_control_channel import IControlChannel, IslandCommunicat
|
|||
from infection_monkey.i_master import IMaster
|
||||
from infection_monkey.i_puppet import IPuppet
|
||||
from infection_monkey.model import VictimHostFactory
|
||||
from infection_monkey.network import NetworkInterface
|
||||
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
|
||||
from infection_monkey.telemetry.post_breach_telem import PostBreachTelem
|
||||
from infection_monkey.telemetry.system_info_telem import SystemInfoTelem
|
||||
|
@ -33,6 +34,7 @@ class AutomatedMaster(IMaster):
|
|||
telemetry_messenger: ITelemetryMessenger,
|
||||
victim_host_factory: VictimHostFactory,
|
||||
control_channel: IControlChannel,
|
||||
local_network_interfaces: List[NetworkInterface],
|
||||
):
|
||||
self._puppet = puppet
|
||||
self._telemetry_messenger = telemetry_messenger
|
||||
|
@ -41,7 +43,11 @@ class AutomatedMaster(IMaster):
|
|||
ip_scanner = IPScanner(self._puppet, NUM_SCAN_THREADS)
|
||||
exploiter = Exploiter(self._puppet, NUM_EXPLOIT_THREADS)
|
||||
self._propagator = Propagator(
|
||||
self._telemetry_messenger, ip_scanner, exploiter, victim_host_factory
|
||||
self._telemetry_messenger,
|
||||
ip_scanner,
|
||||
exploiter,
|
||||
victim_host_factory,
|
||||
local_network_interfaces,
|
||||
)
|
||||
|
||||
self._stop = threading.Event()
|
||||
|
|
|
@ -12,14 +12,14 @@ from infection_monkey.i_puppet import (
|
|||
PortScanData,
|
||||
PortStatus,
|
||||
)
|
||||
from infection_monkey.network import NetworkAddress
|
||||
|
||||
from . import IPScanResults
|
||||
from .threading_utils import run_worker_threads
|
||||
|
||||
logger = logging.getLogger()
|
||||
|
||||
IP = str
|
||||
Callback = Callable[[IP, IPScanResults], None]
|
||||
Callback = Callable[[NetworkAddress, IPScanResults], None]
|
||||
|
||||
|
||||
class IPScanner:
|
||||
|
@ -27,22 +27,33 @@ class IPScanner:
|
|||
self._puppet = puppet
|
||||
self._num_workers = num_workers
|
||||
|
||||
def scan(self, ips_to_scan: List[str], options: Dict, results_callback: Callback, stop: Event):
|
||||
def scan(
|
||||
self,
|
||||
addresses_to_scan: List[NetworkAddress],
|
||||
options: Dict,
|
||||
results_callback: Callback,
|
||||
stop: Event,
|
||||
):
|
||||
# Pre-fill a Queue with all IPs to scan so that threads know they can safely exit when the
|
||||
# queue is empty.
|
||||
ips = Queue()
|
||||
for ip in ips_to_scan:
|
||||
ips.put(ip)
|
||||
addresses = Queue()
|
||||
for address in addresses_to_scan:
|
||||
addresses.put(address)
|
||||
|
||||
scan_ips_args = (ips, options, results_callback, stop)
|
||||
run_worker_threads(target=self._scan_ips, args=scan_ips_args, num_workers=self._num_workers)
|
||||
scan_ips_args = (addresses, options, results_callback, stop)
|
||||
run_worker_threads(
|
||||
target=self._scan_addresses, args=scan_ips_args, num_workers=self._num_workers
|
||||
)
|
||||
|
||||
def _scan_ips(self, ips: Queue, options: Dict, results_callback: Callback, stop: Event):
|
||||
def _scan_addresses(
|
||||
self, addresses: Queue, options: Dict, results_callback: Callback, stop: Event
|
||||
):
|
||||
logger.debug(f"Starting scan thread -- Thread ID: {threading.get_ident()}")
|
||||
|
||||
try:
|
||||
while not stop.is_set():
|
||||
ip = ips.get_nowait()
|
||||
address = addresses.get_nowait()
|
||||
ip = address.ip
|
||||
logger.info(f"Scanning {ip}")
|
||||
|
||||
icmp_timeout = options["icmp"]["timeout_ms"] / 1000
|
||||
|
@ -60,7 +71,7 @@ class IPScanner:
|
|||
)
|
||||
|
||||
scan_results = IPScanResults(ping_scan_data, port_scan_data, fingerprint_data)
|
||||
results_callback(ip, scan_results)
|
||||
results_callback(address, scan_results)
|
||||
|
||||
logger.debug(
|
||||
f"Detected the stop signal, scanning thread {threading.get_ident()} exiting"
|
||||
|
|
|
@ -1,7 +1,7 @@
|
|||
import logging
|
||||
from queue import Queue
|
||||
from threading import Event
|
||||
from typing import Dict
|
||||
from typing import Dict, List
|
||||
|
||||
from infection_monkey.i_puppet import (
|
||||
ExploiterResultData,
|
||||
|
@ -11,6 +11,8 @@ from infection_monkey.i_puppet import (
|
|||
PortStatus,
|
||||
)
|
||||
from infection_monkey.model import VictimHost, VictimHostFactory
|
||||
from infection_monkey.network import NetworkAddress, NetworkInterface
|
||||
from infection_monkey.network.scan_target_generator import compile_scan_target_list
|
||||
from infection_monkey.telemetry.exploit_telem import ExploitTelem
|
||||
from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
|
||||
from infection_monkey.telemetry.scan_telem import ScanTelem
|
||||
|
@ -28,11 +30,13 @@ class Propagator:
|
|||
ip_scanner: IPScanner,
|
||||
exploiter: Exploiter,
|
||||
victim_host_factory: VictimHostFactory,
|
||||
local_network_interfaces: List[NetworkInterface],
|
||||
):
|
||||
self._telemetry_messenger = telemetry_messenger
|
||||
self._ip_scanner = ip_scanner
|
||||
self._exploiter = exploiter
|
||||
self._victim_host_factory = victim_host_factory
|
||||
self._local_network_interfaces = local_network_interfaces
|
||||
self._hosts_to_exploit = None
|
||||
|
||||
def propagate(self, propagation_config: Dict, stop: Event):
|
||||
|
@ -62,16 +66,30 @@ class Propagator:
|
|||
def _scan_network(self, propagation_config: Dict, stop: Event):
|
||||
logger.info("Starting network scan")
|
||||
|
||||
# TODO: Generate list of IPs to scan from propagation targets config
|
||||
ips_to_scan = propagation_config["targets"]["subnet_scan_list"]
|
||||
|
||||
target_config = propagation_config["targets"]
|
||||
scan_config = propagation_config["network_scan"]
|
||||
self._ip_scanner.scan(ips_to_scan, scan_config, self._process_scan_results, stop)
|
||||
|
||||
addresses_to_scan = self._compile_scan_target_list(target_config)
|
||||
self._ip_scanner.scan(addresses_to_scan, scan_config, self._process_scan_results, stop)
|
||||
|
||||
logger.info("Finished network scan")
|
||||
|
||||
def _process_scan_results(self, ip: str, scan_results: IPScanResults):
|
||||
victim_host = self._victim_host_factory.build_victim_host(ip)
|
||||
def _compile_scan_target_list(self, target_config: Dict) -> List[NetworkAddress]:
|
||||
ranges_to_scan = target_config["subnet_scan_list"]
|
||||
inaccessible_subnets = target_config["inaccessible_subnets"]
|
||||
blocklisted_ips = target_config["blocked_ips"]
|
||||
enable_local_network_scan = target_config["local_network_scan"]
|
||||
|
||||
return compile_scan_target_list(
|
||||
self._local_network_interfaces,
|
||||
ranges_to_scan,
|
||||
inaccessible_subnets,
|
||||
blocklisted_ips,
|
||||
enable_local_network_scan,
|
||||
)
|
||||
|
||||
def _process_scan_results(self, address: NetworkAddress, scan_results: IPScanResults):
|
||||
victim_host = self._victim_host_factory.build_victim_host(address)
|
||||
|
||||
Propagator._process_ping_scan_results(victim_host, scan_results.ping_scan_data)
|
||||
Propagator._process_tcp_scan_results(victim_host, scan_results.port_scan_data)
|
||||
|
|
|
@ -1,5 +1,8 @@
|
|||
from typing import Optional
|
||||
|
||||
|
||||
class VictimHost(object):
|
||||
def __init__(self, ip_addr, domain_name=""):
|
||||
def __init__(self, ip_addr: str, domain_name: str = ""):
|
||||
self.ip_addr = ip_addr
|
||||
self.domain_name = str(domain_name)
|
||||
self.os = {}
|
||||
|
@ -42,5 +45,5 @@ class VictimHost(object):
|
|||
victim += "target monkey: %s" % self.monkey_exe
|
||||
return victim
|
||||
|
||||
def set_default_server(self, default_server):
|
||||
self.default_server = default_server
|
||||
def set_island_address(self, ip: str, port: Optional[str]):
|
||||
self.default_server = f"{ip}:{port}" if port else f"{ip}"
|
||||
|
|
|
@ -1,28 +1,50 @@
|
|||
import logging
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from infection_monkey.model import VictimHost
|
||||
from infection_monkey.network import NetworkAddress
|
||||
from infection_monkey.network.tools import get_interface_to_target
|
||||
from infection_monkey.tunnel import MonkeyTunnel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VictimHostFactory:
|
||||
def __init__(self):
|
||||
pass
|
||||
def __init__(
|
||||
self,
|
||||
tunnel: Optional[MonkeyTunnel],
|
||||
island_ip: Optional[str],
|
||||
island_port: Optional[str],
|
||||
on_island: bool,
|
||||
):
|
||||
self.tunnel = tunnel
|
||||
self.island_ip = island_ip
|
||||
self.island_port = island_port
|
||||
self.on_island = on_island
|
||||
|
||||
def build_victim_host(self, ip: str):
|
||||
victim_host = VictimHost(ip)
|
||||
def build_victim_host(self, network_address: NetworkAddress) -> VictimHost:
|
||||
domain = network_address.domain or ""
|
||||
victim_host = VictimHost(network_address.ip, domain)
|
||||
|
||||
# TODO: Reimplement the below logic from the old monkey.py
|
||||
"""
|
||||
if self._monkey_tunnel:
|
||||
self._monkey_tunnel.set_tunnel_for_host(machine)
|
||||
if self._default_server:
|
||||
if self._network.on_island(self._default_server):
|
||||
machine.set_default_server(
|
||||
get_interface_to_target(machine.ip_addr)
|
||||
+ (":" + self._default_server_port if self._default_server_port else "")
|
||||
)
|
||||
else:
|
||||
machine.set_default_server(self._default_server)
|
||||
logger.debug(
|
||||
f"Default server for machine: {machine} set to {machine.default_server}"
|
||||
)
|
||||
"""
|
||||
if self.tunnel:
|
||||
victim_host.default_tunnel = self.tunnel.get_tunnel_for_ip(victim_host.ip_addr)
|
||||
|
||||
if self.island_ip:
|
||||
ip, port = self._choose_island_address(victim_host.ip_addr)
|
||||
victim_host.set_island_address(ip, port)
|
||||
|
||||
logger.debug(f"Default tunnel for {victim_host} set to {victim_host.default_tunnel}")
|
||||
logger.debug(f"Default server for {victim_host} set to {victim_host.default_server}")
|
||||
|
||||
return victim_host
|
||||
|
||||
def _choose_island_address(self, victim_ip: str) -> Tuple[str, Optional[str]]:
|
||||
# Victims need to connect back to the interface they can reach
|
||||
# On island, choose the right interface to pass to children monkeys
|
||||
if self.on_island:
|
||||
default_server_port = self.island_port if self.island_port else None
|
||||
interface = get_interface_to_target(victim_ip)
|
||||
|
||||
return interface, default_server_port
|
||||
else:
|
||||
return self.island_ip, self.island_port
|
||||
|
|
|
@ -4,16 +4,20 @@ import os
|
|||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
import infection_monkey.tunnel as tunnel
|
||||
from common.network.network_utils import address_to_ip_port
|
||||
from common.utils.attack_utils import ScanStatus, UsageEnum
|
||||
from common.version import get_version
|
||||
from infection_monkey.config import GUID, WormConfiguration
|
||||
from infection_monkey.control import ControlClient
|
||||
from infection_monkey.master import AutomatedMaster
|
||||
from infection_monkey.master.control_channel import ControlChannel
|
||||
from infection_monkey.master.mock_master import MockMaster
|
||||
from infection_monkey.model import DELAY_DELETE_CMD
|
||||
from infection_monkey.model import DELAY_DELETE_CMD, VictimHostFactory
|
||||
from infection_monkey.network import NetworkInterface
|
||||
from infection_monkey.network.firewall import app as firewall
|
||||
from infection_monkey.network.info import get_local_network_interfaces
|
||||
from infection_monkey.puppet.mock_puppet import MockPuppet
|
||||
from infection_monkey.system_singleton import SystemSingleton
|
||||
from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
|
||||
|
@ -35,11 +39,10 @@ logger = logging.getLogger(__name__)
|
|||
class InfectionMonkey:
|
||||
def __init__(self, args):
|
||||
logger.info("Monkey is initializing...")
|
||||
self._master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
|
||||
self._singleton = SystemSingleton()
|
||||
self._opts = self._get_arguments(args)
|
||||
# TODO Used in propagation phase to set the default server for the victim
|
||||
self._default_server_port = None
|
||||
self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.server)
|
||||
self._default_server = self._opts.server
|
||||
# TODO used in propogation phase
|
||||
self._monkey_inbound_tunnel = None
|
||||
|
||||
|
@ -52,6 +55,7 @@ class InfectionMonkey:
|
|||
arg_parser.add_argument("-d", "--depth", type=int)
|
||||
opts, _ = arg_parser.parse_known_args(args)
|
||||
InfectionMonkey._log_arguments(opts)
|
||||
|
||||
return opts
|
||||
|
||||
@staticmethod
|
||||
|
@ -108,26 +112,23 @@ class InfectionMonkey:
|
|||
|
||||
def _connect_to_island(self):
|
||||
# Sets island's IP and port for monkey to communicate to
|
||||
if not self._is_default_server_set():
|
||||
if self._current_server_is_set():
|
||||
self._default_server = WormConfiguration.current_server
|
||||
logger.debug("Default server set to: %s" % self._default_server)
|
||||
else:
|
||||
raise Exception(
|
||||
"Monkey couldn't find server with {} default tunnel.".format(self._opts.tunnel)
|
||||
)
|
||||
self._set_default_port()
|
||||
|
||||
ControlClient.wakeup(parent=self._opts.parent)
|
||||
ControlClient.load_control_config()
|
||||
|
||||
def _is_default_server_set(self) -> bool:
|
||||
"""
|
||||
Sets the default server for the Monkey to communicate back to.
|
||||
:return
|
||||
"""
|
||||
if not ControlClient.find_server(default_tunnel=self._opts.tunnel):
|
||||
return False
|
||||
self._opts.server = WormConfiguration.current_server
|
||||
logger.debug("default server set to: %s" % self._opts.server)
|
||||
def _current_server_is_set(self) -> bool:
|
||||
if ControlClient.find_server(default_tunnel=self._opts.tunnel):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_upgrade_to_64_needed():
|
||||
return WindowsUpgrader.should_upgrade()
|
||||
|
@ -151,17 +152,48 @@ class InfectionMonkey:
|
|||
StateTelem(is_done=False, version=get_version()).send()
|
||||
TunnelTelem().send()
|
||||
|
||||
self._build_master()
|
||||
|
||||
register_signal_handlers(self._master)
|
||||
|
||||
@staticmethod
|
||||
def _get_local_network_interfaces():
|
||||
local_network_interfaces = get_local_network_interfaces()
|
||||
for i in local_network_interfaces:
|
||||
logger.debug(f"Found local interface {i.address}{i.netmask}")
|
||||
|
||||
return local_network_interfaces
|
||||
|
||||
def _build_master(self):
|
||||
local_network_interfaces = InfectionMonkey._get_local_network_interfaces()
|
||||
|
||||
victim_host_factory = self._build_victim_host_factory(local_network_interfaces)
|
||||
|
||||
self._master = AutomatedMaster(
|
||||
MockPuppet(),
|
||||
LegacyTelemetryMessengerAdapter(),
|
||||
victim_host_factory,
|
||||
ControlChannel(self._default_server, GUID),
|
||||
local_network_interfaces,
|
||||
)
|
||||
|
||||
def _build_victim_host_factory(
|
||||
self, local_network_interfaces: List[NetworkInterface]
|
||||
) -> VictimHostFactory:
|
||||
on_island = self._running_on_island(local_network_interfaces)
|
||||
logger.debug(f"This agent is running on the island: {on_island}")
|
||||
|
||||
return VictimHostFactory(
|
||||
self._monkey_inbound_tunnel, self._cmd_island_ip, self._cmd_island_port, on_island
|
||||
)
|
||||
|
||||
def _running_on_island(self, local_network_interfaces: List[NetworkInterface]) -> bool:
|
||||
server_ip, _ = address_to_ip_port(self._default_server)
|
||||
return server_ip in {interface.address for interface in local_network_interfaces}
|
||||
|
||||
def _is_another_monkey_running(self):
|
||||
return not self._singleton.try_lock()
|
||||
|
||||
def _set_default_port(self):
|
||||
try:
|
||||
self._default_server_port = self._opts.server.split(":")[1]
|
||||
except KeyError:
|
||||
self._default_server_port = ""
|
||||
|
||||
def cleanup(self):
|
||||
logger.info("Monkey cleanup started")
|
||||
self._wait_for_exploited_machine_connection()
|
||||
|
|
|
@ -0,0 +1 @@
|
|||
from .scan_target_generator import NetworkAddress, NetworkInterface
|
|
@ -1,7 +1,9 @@
|
|||
import itertools
|
||||
import socket
|
||||
import struct
|
||||
from ipaddress import IPv4Network
|
||||
from random import randint # noqa: DUO102
|
||||
from typing import List
|
||||
|
||||
import netifaces
|
||||
import psutil
|
||||
|
@ -9,6 +11,8 @@ import psutil
|
|||
from common.network.network_range import CidrRange
|
||||
from infection_monkey.utils.environment import is_windows_os
|
||||
|
||||
from . import NetworkInterface
|
||||
|
||||
# Timeout for monkey connections
|
||||
TIMEOUT = 15
|
||||
LOOPBACK_NAME = b"lo"
|
||||
|
@ -18,6 +22,16 @@ RTF_UP = 0x0001 # Route usable
|
|||
RTF_REJECT = 0x0200
|
||||
|
||||
|
||||
def get_local_network_interfaces() -> List[NetworkInterface]:
|
||||
network_interfaces = []
|
||||
for i in get_host_subnets():
|
||||
netmask_bits = IPv4Network(f"{i['addr']}/{i['netmask']}", strict=False).prefixlen
|
||||
cidr_netmask = f"/{netmask_bits}"
|
||||
network_interfaces.append(NetworkInterface(i["addr"], cidr_netmask))
|
||||
|
||||
return network_interfaces
|
||||
|
||||
|
||||
def get_host_subnets():
|
||||
"""
|
||||
Returns a list of subnets visible to host (omitting loopback and auto conf networks)
|
||||
|
|
|
@ -4,7 +4,6 @@ import struct
|
|||
import time
|
||||
from threading import Thread
|
||||
|
||||
from infection_monkey.model import VictimHost
|
||||
from infection_monkey.network.firewall import app as firewall
|
||||
from infection_monkey.network.info import get_free_tcp_port, local_ips
|
||||
from infection_monkey.network.tools import check_tcp_port, get_interface_to_target
|
||||
|
@ -188,14 +187,13 @@ class MonkeyTunnel(Thread):
|
|||
proxy.stop()
|
||||
proxy.join()
|
||||
|
||||
def set_tunnel_for_host(self, host):
|
||||
assert isinstance(host, VictimHost)
|
||||
def get_tunnel_for_ip(self, ip: str):
|
||||
|
||||
if not self.local_port:
|
||||
return
|
||||
|
||||
ip_match = get_interface_to_target(host.ip_addr)
|
||||
host.default_tunnel = "%s:%d" % (ip_match, self.local_port)
|
||||
ip_match = get_interface_to_target(ip)
|
||||
return "%s:%d" % (ip_match, self.local_port)
|
||||
|
||||
def stop(self):
|
||||
self._stopped = True
|
||||
|
|
|
@ -1,6 +1,10 @@
|
|||
from unittest import TestCase
|
||||
|
||||
from common.network.network_utils import get_host_from_network_location, remove_port
|
||||
from common.network.network_utils import (
|
||||
address_to_ip_port,
|
||||
get_host_from_network_location,
|
||||
remove_port,
|
||||
)
|
||||
|
||||
|
||||
class TestNetworkUtils(TestCase):
|
||||
|
@ -15,3 +19,17 @@ class TestNetworkUtils(TestCase):
|
|||
assert remove_port("https://google.com:80") == "https://google.com"
|
||||
assert remove_port("https://8.8.8.8:65336") == "https://8.8.8.8"
|
||||
assert remove_port("ftp://ftpserver.com:21/hello/world") == "ftp://ftpserver.com"
|
||||
|
||||
|
||||
def test_address_to_ip_port():
|
||||
ip, port = address_to_ip_port("192.168.65.1:5000")
|
||||
assert ip == "192.168.65.1"
|
||||
assert port == "5000"
|
||||
|
||||
|
||||
def test_address_to_ip_port_no_port():
|
||||
ip, port = address_to_ip_port("192.168.65.1")
|
||||
assert port is None
|
||||
|
||||
ip, port = address_to_ip_port("192.168.65.1:")
|
||||
assert port is None
|
||||
|
|
|
@ -14,7 +14,7 @@ INTERVAL = 0.001
|
|||
|
||||
|
||||
def test_terminate_without_start():
|
||||
m = AutomatedMaster(None, None, None, None)
|
||||
m = AutomatedMaster(None, None, None, None, [])
|
||||
|
||||
# Test that call to terminate does not raise exception
|
||||
m.terminate()
|
||||
|
@ -34,7 +34,7 @@ def test_stop_if_cant_get_config_from_island(monkeypatch):
|
|||
monkeypatch.setattr(
|
||||
"infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL
|
||||
)
|
||||
m = AutomatedMaster(None, None, None, cc)
|
||||
m = AutomatedMaster(None, None, None, cc, [])
|
||||
m.start()
|
||||
|
||||
assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT
|
||||
|
@ -73,7 +73,7 @@ def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, sleep_and_return_
|
|||
"infection_monkey.master.automated_master.CHECK_FOR_TERMINATE_INTERVAL_SEC", INTERVAL
|
||||
)
|
||||
|
||||
m = AutomatedMaster(None, None, None, cc)
|
||||
m = AutomatedMaster(None, None, None, cc, [])
|
||||
m.start()
|
||||
|
||||
assert cc.should_agent_stop.call_count == CHECK_FOR_STOP_AGENT_COUNT
|
||||
|
|
|
@ -6,6 +6,7 @@ import pytest
|
|||
|
||||
from infection_monkey.i_puppet import FingerprintData, PortScanData, PortStatus
|
||||
from infection_monkey.master import IPScanner
|
||||
from infection_monkey.network import NetworkAddress
|
||||
from infection_monkey.puppet.mock_puppet import MockPuppet
|
||||
|
||||
WINDOWS_OS = "windows"
|
||||
|
@ -51,20 +52,21 @@ def assert_port_status(port_scan_data, expected_open_ports: Set[int]):
|
|||
assert psd.status == PortStatus.CLOSED
|
||||
|
||||
|
||||
def assert_scan_results(ip, scan_results):
|
||||
def assert_scan_results(address, scan_results):
|
||||
ping_scan_data = scan_results.ping_scan_data
|
||||
port_scan_data = scan_results.port_scan_data
|
||||
fingerprint_data = scan_results.fingerprint_data
|
||||
|
||||
if ip == "10.0.0.1":
|
||||
assert_scan_results_no_1(ping_scan_data, port_scan_data, fingerprint_data)
|
||||
elif ip == "10.0.0.3":
|
||||
assert_scan_results_no_3(ping_scan_data, port_scan_data, fingerprint_data)
|
||||
if address.ip == "10.0.0.1":
|
||||
assert_scan_results_no_1(address.domain, ping_scan_data, port_scan_data, fingerprint_data)
|
||||
elif address.ip == "10.0.0.3":
|
||||
assert_scan_results_no_3(address.domain, ping_scan_data, port_scan_data, fingerprint_data)
|
||||
else:
|
||||
assert_scan_results_host_down(ip, ping_scan_data, port_scan_data, fingerprint_data)
|
||||
assert_scan_results_host_down(address, ping_scan_data, port_scan_data, fingerprint_data)
|
||||
|
||||
|
||||
def assert_scan_results_no_1(ping_scan_data, port_scan_data, fingerprint_data):
|
||||
def assert_scan_results_no_1(domain, ping_scan_data, port_scan_data, fingerprint_data):
|
||||
assert domain == "d1"
|
||||
assert ping_scan_data.response_received is True
|
||||
assert ping_scan_data.os == WINDOWS_OS
|
||||
|
||||
|
@ -97,7 +99,9 @@ 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(ping_scan_data, port_scan_data, fingerprint_data):
|
||||
def assert_scan_results_no_3(domain, ping_scan_data, port_scan_data, fingerprint_data):
|
||||
assert domain == "d3"
|
||||
|
||||
assert ping_scan_data.response_received is True
|
||||
assert ping_scan_data.os == LINUX_OS
|
||||
assert len(port_scan_data.keys()) == 6
|
||||
|
@ -135,8 +139,9 @@ 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(ip, ping_scan_data, port_scan_data, fingerprint_data):
|
||||
assert ip not in {"10.0.0.1", "10.0.0.3"}
|
||||
def assert_scan_results_host_down(address, ping_scan_data, port_scan_data, fingerprint_data):
|
||||
assert address.ip not in {"10.0.0.1", "10.0.0.3"}
|
||||
assert address.domain is None
|
||||
|
||||
assert ping_scan_data.response_received is False
|
||||
assert len(port_scan_data.keys()) == 6
|
||||
|
@ -146,44 +151,49 @@ def assert_scan_results_host_down(ip, ping_scan_data, port_scan_data, fingerprin
|
|||
|
||||
|
||||
def test_scan_single_ip(callback, scan_config, stop):
|
||||
ips = ["10.0.0.1"]
|
||||
addresses = [NetworkAddress("10.0.0.1", "d1")]
|
||||
|
||||
ns = IPScanner(MockPuppet(), num_workers=1)
|
||||
ns.scan(ips, scan_config, callback, stop)
|
||||
ns.scan(addresses, scan_config, callback, stop)
|
||||
|
||||
callback.assert_called_once()
|
||||
|
||||
(ip, scan_results) = callback.call_args_list[0][0]
|
||||
assert_scan_results(ip, scan_results)
|
||||
(address, scan_results) = callback.call_args_list[0][0]
|
||||
assert_scan_results(address, scan_results)
|
||||
|
||||
|
||||
def test_scan_multiple_ips(callback, scan_config, stop):
|
||||
ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"]
|
||||
addresses = [
|
||||
NetworkAddress("10.0.0.1", "d1"),
|
||||
NetworkAddress("10.0.0.2", None),
|
||||
NetworkAddress("10.0.0.3", "d3"),
|
||||
NetworkAddress("10.0.0.4", None),
|
||||
]
|
||||
|
||||
ns = IPScanner(MockPuppet(), num_workers=4)
|
||||
ns.scan(ips, scan_config, callback, stop)
|
||||
ns.scan(addresses, scan_config, callback, stop)
|
||||
|
||||
assert callback.call_count == 4
|
||||
|
||||
(ip, scan_results) = callback.call_args_list[0][0]
|
||||
assert_scan_results(ip, scan_results)
|
||||
(address, scan_results) = callback.call_args_list[0][0]
|
||||
assert_scan_results(address, scan_results)
|
||||
|
||||
(ip, scan_results) = callback.call_args_list[1][0]
|
||||
assert_scan_results(ip, scan_results)
|
||||
(address, scan_results) = callback.call_args_list[1][0]
|
||||
assert_scan_results(address, scan_results)
|
||||
|
||||
(ip, scan_results) = callback.call_args_list[2][0]
|
||||
assert_scan_results(ip, scan_results)
|
||||
(address, scan_results) = callback.call_args_list[2][0]
|
||||
assert_scan_results(address, scan_results)
|
||||
|
||||
(ip, scan_results) = callback.call_args_list[3][0]
|
||||
assert_scan_results(ip, scan_results)
|
||||
(address, scan_results) = callback.call_args_list[3][0]
|
||||
assert_scan_results(address, scan_results)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_scan_lots_of_ips(callback, scan_config, stop):
|
||||
ips = [f"10.0.0.{i}" for i in range(0, 255)]
|
||||
addresses = [NetworkAddress(f"10.0.0.{i}", None) for i in range(0, 255)]
|
||||
|
||||
ns = IPScanner(MockPuppet(), num_workers=4)
|
||||
ns.scan(ips, scan_config, callback, stop)
|
||||
ns.scan(addresses, scan_config, callback, stop)
|
||||
|
||||
assert callback.call_count == 255
|
||||
|
||||
|
@ -199,10 +209,15 @@ def test_stop_after_callback(scan_config, stop):
|
|||
|
||||
stoppable_callback = MagicMock(side_effect=_callback)
|
||||
|
||||
ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"]
|
||||
addresses = [
|
||||
NetworkAddress("10.0.0.1", None),
|
||||
NetworkAddress("10.0.0.2", None),
|
||||
NetworkAddress("10.0.0.3", None),
|
||||
NetworkAddress("10.0.0.4", None),
|
||||
]
|
||||
|
||||
ns = IPScanner(MockPuppet(), num_workers=2)
|
||||
ns.scan(ips, scan_config, stoppable_callback, stop)
|
||||
ns.scan(addresses, scan_config, stoppable_callback, stop)
|
||||
|
||||
assert stoppable_callback.call_count == 2
|
||||
|
||||
|
@ -221,10 +236,15 @@ def test_interrupt_port_scanning(callback, scan_config, stop):
|
|||
puppet = MockPuppet()
|
||||
puppet.scan_tcp_port = MagicMock(side_effect=stoppable_scan_tcp_port)
|
||||
|
||||
ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"]
|
||||
addresses = [
|
||||
NetworkAddress("10.0.0.1", None),
|
||||
NetworkAddress("10.0.0.2", None),
|
||||
NetworkAddress("10.0.0.3", None),
|
||||
NetworkAddress("10.0.0.4", None),
|
||||
]
|
||||
|
||||
ns = IPScanner(puppet, num_workers=2)
|
||||
ns.scan(ips, scan_config, callback, stop)
|
||||
ns.scan(addresses, scan_config, callback, stop)
|
||||
|
||||
assert puppet.scan_tcp_port.call_count == 2
|
||||
|
||||
|
@ -243,9 +263,14 @@ def test_interrupt_fingerprinting(callback, scan_config, stop):
|
|||
puppet = MockPuppet()
|
||||
puppet.fingerprint = MagicMock(side_effect=stoppable_fingerprint)
|
||||
|
||||
ips = ["10.0.0.1", "10.0.0.2", "10.0.0.3", "10.0.0.4"]
|
||||
addresses = [
|
||||
NetworkAddress("10.0.0.1", None),
|
||||
NetworkAddress("10.0.0.2", None),
|
||||
NetworkAddress("10.0.0.3", None),
|
||||
NetworkAddress("10.0.0.4", None),
|
||||
]
|
||||
|
||||
ns = IPScanner(puppet, num_workers=2)
|
||||
ns.scan(ips, scan_config, callback, stop)
|
||||
ns.scan(addresses, scan_config, callback, stop)
|
||||
|
||||
assert puppet.fingerprint.call_count == 2
|
||||
|
|
|
@ -1,4 +1,7 @@
|
|||
from threading import Event
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from infection_monkey.i_puppet import (
|
||||
ExploiterResultData,
|
||||
|
@ -8,8 +11,26 @@ from infection_monkey.i_puppet import (
|
|||
PortStatus,
|
||||
)
|
||||
from infection_monkey.master import IPScanResults, Propagator
|
||||
from infection_monkey.model import VictimHostFactory
|
||||
from infection_monkey.network import NetworkInterface
|
||||
from infection_monkey.telemetry.exploit_telem import ExploitTelem
|
||||
from infection_monkey.model import VictimHost, VictimHostFactory
|
||||
from infection_monkey.network import NetworkAddress
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_victim_host_factory():
|
||||
class MockVictimHostFactory(VictimHostFactory):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def build_victim_host(self, network_address: NetworkAddress) -> VictimHost:
|
||||
domain = network_address.domain or ""
|
||||
return VictimHost(network_address.ip, domain)
|
||||
|
||||
return MockVictimHostFactory()
|
||||
|
||||
|
||||
empty_fingerprint_data = FingerprintData(None, None, {})
|
||||
|
||||
|
@ -83,15 +104,21 @@ dot_3_services = {
|
|||
}
|
||||
|
||||
|
||||
class MockIPScanner:
|
||||
def scan(self, ips_to_scan, _, results_callback, stop):
|
||||
for ip in ips_to_scan:
|
||||
if ip.endswith(".1"):
|
||||
results_callback(ip, dot_1_scan_results)
|
||||
elif ip.endswith(".3"):
|
||||
results_callback(ip, dot_3_scan_results)
|
||||
@pytest.fixture
|
||||
def mock_ip_scanner():
|
||||
def scan(adresses_to_scan, _, results_callback, stop):
|
||||
for address in adresses_to_scan:
|
||||
if address.ip.endswith(".1"):
|
||||
results_callback(address, dot_1_scan_results)
|
||||
elif address.ip.endswith(".3"):
|
||||
results_callback(address, dot_3_scan_results)
|
||||
else:
|
||||
results_callback(ip, dead_host_scan_results)
|
||||
results_callback(address, dead_host_scan_results)
|
||||
|
||||
ip_scanner = MagicMock()
|
||||
ip_scanner.scan = MagicMock(side_effect=scan)
|
||||
|
||||
return ip_scanner
|
||||
|
||||
|
||||
class StubExploiter:
|
||||
|
@ -101,11 +128,18 @@ class StubExploiter:
|
|||
pass
|
||||
|
||||
|
||||
def test_scan_result_processing(telemetry_messenger_spy):
|
||||
p = Propagator(telemetry_messenger_spy, MockIPScanner(), StubExploiter(), VictimHostFactory())
|
||||
def test_scan_result_processing(telemetry_messenger_spy, mock_ip_scanner, mock_victim_host_factory):
|
||||
p = Propagator(
|
||||
telemetry_messenger_spy, mock_ip_scanner, StubExploiter(), mock_victim_host_factory, []
|
||||
)
|
||||
p.propagate(
|
||||
{
|
||||
"targets": {"subnet_scan_list": ["10.0.0.1", "10.0.0.2", "10.0.0.3"]},
|
||||
"targets": {
|
||||
"subnet_scan_list": ["10.0.0.1", "10.0.0.2", "10.0.0.3"],
|
||||
"local_network_scan": False,
|
||||
"inaccessible_subnets": [],
|
||||
"blocked_ips": [],
|
||||
},
|
||||
"network_scan": {}, # This is empty since MockIPscanner ignores it
|
||||
"exploiters": {}, # This is empty since StubExploiter ignores it
|
||||
},
|
||||
|
@ -141,10 +175,13 @@ class MockExploiter:
|
|||
def exploit_hosts(
|
||||
self, exploiter_config, hosts_to_exploit, results_callback, scan_completed, stop
|
||||
):
|
||||
scan_completed.wait()
|
||||
hte = []
|
||||
for _ in range(0, 2):
|
||||
hte.append(hosts_to_exploit.get())
|
||||
|
||||
assert hosts_to_exploit.empty()
|
||||
|
||||
for host in hte:
|
||||
if host.ip_addr.endswith(".1"):
|
||||
results_callback(
|
||||
|
@ -157,7 +194,7 @@ class MockExploiter:
|
|||
host,
|
||||
ExploiterResultData(False, {}, {}, "SSH FAILED for .1"),
|
||||
)
|
||||
if host.ip_addr.endswith(".2"):
|
||||
elif host.ip_addr.endswith(".2"):
|
||||
results_callback(
|
||||
"PowerShellExploiter",
|
||||
host,
|
||||
|
@ -168,7 +205,7 @@ class MockExploiter:
|
|||
host,
|
||||
ExploiterResultData(False, {}, {}, "SSH FAILED for .2"),
|
||||
)
|
||||
if host.ip_addr.endswith(".3"):
|
||||
elif host.ip_addr.endswith(".3"):
|
||||
results_callback(
|
||||
"PowerShellExploiter",
|
||||
host,
|
||||
|
@ -181,11 +218,20 @@ class MockExploiter:
|
|||
)
|
||||
|
||||
|
||||
def test_exploiter_result_processing(telemetry_messenger_spy):
|
||||
p = Propagator(telemetry_messenger_spy, MockIPScanner(), MockExploiter(), VictimHostFactory())
|
||||
def test_exploiter_result_processing(
|
||||
telemetry_messenger_spy, mock_ip_scanner, mock_victim_host_factory
|
||||
):
|
||||
p = Propagator(
|
||||
telemetry_messenger_spy, mock_ip_scanner, MockExploiter(), mock_victim_host_factory, []
|
||||
)
|
||||
p.propagate(
|
||||
{
|
||||
"targets": {"subnet_scan_list": ["10.0.0.1", "10.0.0.2", "10.0.0.3"]},
|
||||
"targets": {
|
||||
"subnet_scan_list": ["10.0.0.1", "10.0.0.2", "10.0.0.3"],
|
||||
"local_network_scan": False,
|
||||
"inaccessible_subnets": [],
|
||||
"blocked_ips": [],
|
||||
},
|
||||
"network_scan": {}, # This is empty since MockIPscanner ignores it
|
||||
"exploiters": {}, # This is empty since MockExploiter ignores it
|
||||
},
|
||||
|
@ -211,3 +257,48 @@ def test_exploiter_result_processing(telemetry_messenger_spy):
|
|||
assert not data["result"]
|
||||
else:
|
||||
assert data["result"]
|
||||
|
||||
|
||||
def test_scan_target_generation(telemetry_messenger_spy, mock_ip_scanner, mock_victim_host_factory):
|
||||
local_network_interfaces = [NetworkInterface("10.0.0.9", "/29")]
|
||||
p = Propagator(
|
||||
telemetry_messenger_spy,
|
||||
mock_ip_scanner,
|
||||
StubExploiter(),
|
||||
mock_victim_host_factory,
|
||||
local_network_interfaces,
|
||||
)
|
||||
p.propagate(
|
||||
{
|
||||
"targets": {
|
||||
"subnet_scan_list": ["10.0.0.0/29", "172.10.20.30"],
|
||||
"local_network_scan": True,
|
||||
"blocked_ips": ["10.0.0.3"],
|
||||
"inaccessible_subnets": ["10.0.0.128/30", "10.0.0.8/29"],
|
||||
},
|
||||
"network_scan": {}, # This is empty since MockIPscanner ignores it
|
||||
"exploiters": {}, # This is empty since MockExploiter ignores it
|
||||
},
|
||||
Event(),
|
||||
)
|
||||
expected_ip_scan_list = [
|
||||
"10.0.0.0",
|
||||
"10.0.0.1",
|
||||
"10.0.0.2",
|
||||
"10.0.0.4",
|
||||
"10.0.0.5",
|
||||
"10.0.0.6",
|
||||
"10.0.0.8",
|
||||
"10.0.0.10",
|
||||
"10.0.0.11",
|
||||
"10.0.0.12",
|
||||
"10.0.0.13",
|
||||
"10.0.0.14",
|
||||
"10.0.0.128",
|
||||
"10.0.0.129",
|
||||
"10.0.0.130",
|
||||
"172.10.20.30",
|
||||
]
|
||||
|
||||
actual_ip_scan_list = [address.ip for address in mock_ip_scanner.scan.call_args_list[0][0][0]]
|
||||
assert actual_ip_scan_list == expected_ip_scan_list
|
||||
|
|
|
@ -0,0 +1,83 @@
|
|||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from infection_monkey.model import VictimHostFactory
|
||||
from infection_monkey.network.scan_target_generator import NetworkAddress
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tunnel():
|
||||
tunnel = MagicMock()
|
||||
tunnel.get_tunnel_for_ip = lambda _: "1.2.3.4:1234"
|
||||
return tunnel
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_get_interface_to_target(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"infection_monkey.model.victim_host_factory.get_interface_to_target", lambda _: "1.1.1.1"
|
||||
)
|
||||
|
||||
|
||||
def test_factory_no_tunnel():
|
||||
factory = VictimHostFactory(
|
||||
tunnel=None, island_ip="192.168.56.1", island_port="5000", on_island=False
|
||||
)
|
||||
network_address = NetworkAddress("192.168.56.2", None)
|
||||
|
||||
victim = factory.build_victim_host(network_address)
|
||||
|
||||
assert victim.default_server == "192.168.56.1:5000"
|
||||
assert victim.ip_addr == "192.168.56.2"
|
||||
assert victim.default_tunnel is None
|
||||
assert victim.domain_name == ""
|
||||
|
||||
|
||||
def test_factory_with_tunnel(mock_tunnel):
|
||||
factory = VictimHostFactory(
|
||||
tunnel=mock_tunnel, island_ip="192.168.56.1", island_port="5000", on_island=False
|
||||
)
|
||||
network_address = NetworkAddress("192.168.56.2", None)
|
||||
|
||||
victim = factory.build_victim_host(network_address)
|
||||
|
||||
assert victim.default_server == "192.168.56.1:5000"
|
||||
assert victim.ip_addr == "192.168.56.2"
|
||||
assert victim.default_tunnel == "1.2.3.4:1234"
|
||||
assert victim.domain_name == ""
|
||||
|
||||
|
||||
def test_factory_on_island(mock_tunnel):
|
||||
factory = VictimHostFactory(
|
||||
tunnel=mock_tunnel, island_ip="192.168.56.1", island_port="99", on_island=True
|
||||
)
|
||||
network_address = NetworkAddress("192.168.56.2", "www.bogus.monkey")
|
||||
|
||||
victim = factory.build_victim_host(network_address)
|
||||
|
||||
assert victim.default_server == "1.1.1.1:99"
|
||||
assert victim.domain_name == "www.bogus.monkey"
|
||||
assert victim.ip_addr == "192.168.56.2"
|
||||
assert victim.default_tunnel == "1.2.3.4:1234"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("default_port", ["", None])
|
||||
def test_factory_no_port(mock_tunnel, default_port):
|
||||
factory = VictimHostFactory(
|
||||
tunnel=mock_tunnel, island_ip="192.168.56.1", island_port=default_port, on_island=True
|
||||
)
|
||||
network_address = NetworkAddress("192.168.56.2", "www.bogus.monkey")
|
||||
|
||||
victim = factory.build_victim_host(network_address)
|
||||
|
||||
assert victim.default_server == "1.1.1.1"
|
||||
|
||||
|
||||
def test_factory_no_default_server(mock_tunnel):
|
||||
factory = VictimHostFactory(tunnel=mock_tunnel, island_ip=None, island_port="", on_island=True)
|
||||
network_address = NetworkAddress("192.168.56.2", "www.bogus.monkey")
|
||||
|
||||
victim = factory.build_victim_host(network_address)
|
||||
|
||||
assert victim.default_server is None
|
|
@ -96,9 +96,9 @@ def test_get_monkey_commandline_linux():
|
|||
|
||||
def test_build_monkey_commandline():
|
||||
example_host = VictimHost(ip_addr="bla")
|
||||
example_host.set_default_server("101010")
|
||||
example_host.set_island_address("101010", "5000")
|
||||
|
||||
expected = f" -p {GUID} -s 101010 -d 0 -l /home/bla"
|
||||
expected = f" -p {GUID} -s 101010:5000 -d 0 -l /home/bla"
|
||||
actual = build_monkey_commandline(target_host=example_host, depth=0, location="/home/bla")
|
||||
|
||||
assert expected == actual
|
||||
|
|
Loading…
Reference in New Issue