231 lines
11 KiB
Python
231 lines
11 KiB
Python
import errno
|
|
import socket
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from rpcscope.targets import _connect, discover_services, expand_targets
|
|
|
|
|
|
class TargetExpansionTests(unittest.TestCase):
|
|
def test_local_socket_exhaustion_is_not_a_dead_target_or_user_cancellation(self):
|
|
cancel = threading.Event()
|
|
def exhausted(address, port, timeout):
|
|
raise OSError(errno.EMFILE, "too many open files")
|
|
with self.assertRaisesRegex(OSError, "local socket resources exhausted"):
|
|
discover_services(["192.0.2.1"], jobs=1, cancel=cancel, probe=exhausted)
|
|
self.assertFalse(cancel.is_set())
|
|
|
|
def test_canonical_unique_sorted_addresses_and_host_semantics(self):
|
|
targets = ["2001:db8::/126", "192.0.2.0/30", "192.0.2.2",
|
|
"192.0.2.4/31", "192.0.2.7/32", "2001:db8::4/127", "::1/128"]
|
|
self.assertEqual(expand_targets(targets), ["192.0.2.1", "192.0.2.2", "192.0.2.4",
|
|
"192.0.2.5", "192.0.2.7", "::1",
|
|
"2001:db8::1", "2001:db8::2", "2001:db8::3",
|
|
"2001:db8::4", "2001:db8::5"])
|
|
|
|
def test_non_strict_network_and_overlapping_limit(self):
|
|
self.assertEqual(expand_targets(["192.0.2.2/30", "192.0.2.0/30"], max_hosts=2),
|
|
["192.0.2.1", "192.0.2.2"])
|
|
with self.assertRaisesRegex(ValueError, "limit"):
|
|
expand_targets(["192.0.2.1", "192.0.2.2"], max_hosts=1)
|
|
|
|
def test_huge_network_is_rejected_before_iteration(self):
|
|
for cidr in ("0.0.0.0/0", "::/0"):
|
|
with self.subTest(cidr=cidr), self.assertRaisesRegex(ValueError, "contains.*limit"):
|
|
expand_targets([cidr])
|
|
|
|
def test_file_comments_bom_hostname_and_ipv6_resolution(self):
|
|
resolver = Mock(return_value=[
|
|
(socket.AF_INET6, socket.SOCK_STREAM, 6, "", ("2001:0db8::9", 0, 0, 0)),
|
|
(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("192.0.2.9", 0)),
|
|
])
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
target_file = Path(temporary) / "targets.txt"
|
|
target_file.write_text("\ufeff# lab\n192.0.2.1 # first\n\nserver.example\n", encoding="utf-8")
|
|
self.assertEqual(expand_targets([str(target_file), "192.0.2.9"], resolver=resolver),
|
|
["192.0.2.1", "192.0.2.9", "2001:db8::9"])
|
|
resolver.assert_called_once_with("server.example", None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
|
|
|
def test_nested_files_rejected_with_line_number(self):
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
parent = Path(temporary) / "parent.txt"
|
|
child = Path(temporary) / "child.txt"
|
|
child.write_text("192.0.2.1\n", encoding="utf-8")
|
|
parent.write_text("# comment\nchild.txt\n", encoding="utf-8")
|
|
with self.assertRaisesRegex(ValueError, r"parent\.txt:2: nested target files"):
|
|
expand_targets([str(parent)])
|
|
|
|
def test_invalid_inputs_do_not_call_dns(self):
|
|
resolver = Mock()
|
|
for target in ("-M", "", "https://example.org", "999.2.3.4", "192.0.2.1/99", "bad name", "::x"):
|
|
with self.subTest(target=target), self.assertRaises(ValueError):
|
|
expand_targets([target], resolver=resolver)
|
|
resolver.assert_not_called()
|
|
|
|
def test_dns_errors_and_empty_file_are_useful(self):
|
|
with self.assertRaisesRegex(ValueError, "could not resolve.*example"):
|
|
expand_targets(["example"], resolver=Mock(side_effect=socket.gaierror("missing")))
|
|
with self.assertRaisesRegex(ValueError, "did not resolve"):
|
|
expand_targets(["example"], resolver=Mock(return_value=[]))
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
target_file = Path(temporary) / "empty.txt"
|
|
target_file.write_text("# none\n", encoding="utf-8")
|
|
with self.assertRaisesRegex(ValueError, "did not contain"):
|
|
expand_targets([str(target_file)])
|
|
|
|
|
|
class ServiceDiscoveryTests(unittest.TestCase):
|
|
def test_probe_states_are_independent_and_callbacks_use_caller_thread(self):
|
|
callback = Mock()
|
|
caller_thread = threading.get_ident()
|
|
|
|
def completed(address, ports):
|
|
self.assertEqual(threading.get_ident(), caller_thread)
|
|
callback(address, ports)
|
|
|
|
def probe(address, port, timeout):
|
|
self.assertEqual(timeout, 0.25)
|
|
if address == "192.0.2.1" and port == 135:
|
|
return "open"
|
|
if address == "192.0.2.1":
|
|
raise ConnectionRefusedError()
|
|
if port == 135:
|
|
raise TimeoutError()
|
|
raise OSError(errno.ENETUNREACH, "unreachable")
|
|
|
|
result = discover_services(["192.0.2.2", "192.0.2.1", "192.0.2.1"], jobs=2,
|
|
timeout=0.25, probe=probe, on_result=completed)
|
|
self.assertEqual(result, {"192.0.2.1": {135: "open", 445: "closed"},
|
|
"192.0.2.2": {135: "timeout", 445: "error"}})
|
|
self.assertEqual(callback.call_count, 2)
|
|
|
|
def test_host_concurrency_is_bounded(self):
|
|
barrier = threading.Barrier(2)
|
|
lock = threading.Lock()
|
|
active = 0
|
|
peak = 0
|
|
order = {}
|
|
|
|
def probe(address, port, timeout):
|
|
nonlocal active, peak
|
|
with lock:
|
|
order.setdefault(address, []).append(port)
|
|
if port == 135:
|
|
active += 1
|
|
peak = max(peak, active)
|
|
if port == 135:
|
|
barrier.wait(timeout=2)
|
|
else:
|
|
with lock:
|
|
active -= 1
|
|
return "open"
|
|
|
|
results = discover_services([f"192.0.2.{index}" for index in range(1, 7)], jobs=2, probe=probe)
|
|
self.assertEqual(len(results), 6)
|
|
self.assertEqual(peak, 2)
|
|
self.assertTrue(all(ports == [135, 445] for ports in order.values()))
|
|
|
|
def test_cancel_stops_new_hosts_and_second_port(self):
|
|
cancel = threading.Event()
|
|
probe = Mock(side_effect=lambda *_: (cancel.set() or "open"))
|
|
result = discover_services(["192.0.2.1", "192.0.2.2"], jobs=1, cancel=cancel, probe=probe)
|
|
self.assertEqual(result, {"192.0.2.1": {135: "open"}})
|
|
probe.assert_called_once()
|
|
self.assertEqual(discover_services(["192.0.2.1"], cancel=cancel, probe=probe), {})
|
|
|
|
def test_interrupt_sets_cancel_event_and_propagates(self):
|
|
cancel = threading.Event()
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
discover_services(["192.0.2.1"], cancel=cancel, probe=Mock(side_effect=KeyboardInterrupt))
|
|
self.assertTrue(cancel.is_set())
|
|
|
|
def test_expired_deadline_makes_no_probe_and_does_not_cancel_caller(self):
|
|
cancel = threading.Event()
|
|
probe = Mock()
|
|
with patch("rpcscope.targets.time.monotonic", return_value=10.0):
|
|
self.assertEqual(discover_services(["192.0.2.1"], cancel=cancel, probe=probe, deadline=9.0), {})
|
|
probe.assert_not_called()
|
|
self.assertFalse(cancel.is_set())
|
|
|
|
def test_deadline_stops_next_port_and_host_but_retains_completed_evidence(self):
|
|
cancel = threading.Event()
|
|
now = [10.0]
|
|
callback = Mock()
|
|
|
|
def probe(*_):
|
|
now[0] = 20.0
|
|
return "open"
|
|
|
|
with patch("rpcscope.targets.time.monotonic", side_effect=lambda: now[0]):
|
|
result = discover_services(["192.0.2.1", "192.0.2.2"], jobs=1, cancel=cancel,
|
|
probe=probe, deadline=20.0, on_result=callback)
|
|
self.assertEqual(result, {"192.0.2.1": {135: "open"}})
|
|
callback.assert_called_once_with("192.0.2.1", {135: "open"})
|
|
self.assertFalse(cancel.is_set())
|
|
|
|
def test_each_port_timeout_is_capped_to_remaining_deadline_budget(self):
|
|
now = [10.0]
|
|
calls = []
|
|
|
|
def probe(address, port, timeout):
|
|
calls.append((address, port, timeout))
|
|
now[0] += 0.25
|
|
return "closed"
|
|
|
|
with patch("rpcscope.targets.time.monotonic", side_effect=lambda: now[0]):
|
|
result = discover_services(["192.0.2.1", "192.0.2.2"], jobs=1, timeout=2.0,
|
|
deadline=10.5, probe=probe)
|
|
self.assertEqual(result, {"192.0.2.1": {135: "closed", 445: "closed"}})
|
|
self.assertEqual(calls, [("192.0.2.1", 135, 0.5), ("192.0.2.1", 445, 0.25)])
|
|
|
|
def test_worker_slot_is_reused_without_waiting_for_slower_host(self):
|
|
slow_started = threading.Event()
|
|
next_host_started = threading.Event()
|
|
callback = Mock()
|
|
|
|
def probe(address, port, timeout):
|
|
if address == "192.0.2.1" and port == 135:
|
|
slow_started.set()
|
|
if not next_host_started.wait(timeout=2):
|
|
raise AssertionError("completed worker did not take the next host")
|
|
if address == "192.0.2.2" and port == 135:
|
|
if not slow_started.wait(timeout=2):
|
|
raise AssertionError("first worker did not start")
|
|
if address == "192.0.2.3":
|
|
next_host_started.set()
|
|
return "closed"
|
|
|
|
result = discover_services(["192.0.2.1", "192.0.2.2", "192.0.2.3"], jobs=2,
|
|
probe=probe, on_result=callback)
|
|
self.assertTrue(next_host_started.is_set())
|
|
self.assertEqual(result, {address: {135: "closed", 445: "closed"}
|
|
for address in ("192.0.2.1", "192.0.2.2", "192.0.2.3")})
|
|
self.assertEqual(callback.call_count, 3)
|
|
self.assertEqual({call.args[0] for call in callback.call_args_list}, set(result))
|
|
|
|
def test_default_probe_connects_and_closes_without_sending(self):
|
|
with patch("rpcscope.targets.socket.create_connection") as connection:
|
|
self.assertEqual(_connect("192.0.2.1", 135, 0.2), "open")
|
|
connection.assert_called_once_with(("192.0.2.1", 135), timeout=0.2)
|
|
connection.return_value.__exit__.assert_called_once()
|
|
connection.return_value.__enter__.return_value.send.assert_not_called()
|
|
connection.return_value.__enter__.return_value.sendall.assert_not_called()
|
|
|
|
def test_invalid_configuration_and_unexpected_probe_result(self):
|
|
for kwargs in ({"jobs": 0}, {"jobs": True}, {"timeout": 0}, {"timeout": float("inf")},
|
|
{"timeout": float("nan")}, {"deadline": float("nan")},
|
|
{"deadline": float("inf")}, {"deadline": True}):
|
|
with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
|
|
discover_services(["192.0.2.1"], **kwargs)
|
|
with self.assertRaises(ValueError):
|
|
discover_services(["server.example"])
|
|
self.assertEqual(discover_services(["192.0.2.1"], probe=Mock(return_value="unexpected")),
|
|
{"192.0.2.1": {135: "error", 445: "error"}})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|