166 lines
7.2 KiB
Python
166 lines
7.2 KiB
Python
import io
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from rpcscope.progress import ProgressEvent, ProgressReporter
|
|
|
|
|
|
class Clock:
|
|
def __init__(self):
|
|
self.now = 0.0
|
|
|
|
def __call__(self):
|
|
return self.now
|
|
|
|
|
|
class Terminal(io.StringIO):
|
|
def isatty(self):
|
|
return True
|
|
|
|
|
|
class ProgressTests(unittest.TestCase):
|
|
def test_redirected_output_has_results_and_no_terminal_codes(self):
|
|
output = io.StringIO()
|
|
clock = Clock()
|
|
reporter = ProgressReporter(output, clock=clock)
|
|
reporter(ProgressEvent("start", "ms17-010", "10.0.0.1", total=2))
|
|
clock.now = 1.25
|
|
reporter(ProgressEvent("finish", "ms17-010", "10.0.0.1", "vulnerable", 1, 2))
|
|
reporter(ProgressEvent("summary", "scan", status="complete", completed=2, total=2))
|
|
self.assertIn("[0/2] START ms17-010 10.0.0.1", output.getvalue())
|
|
self.assertIn("[1/2] DONE ms17-010 10.0.0.1 vulnerable (1.2s)", output.getvalue())
|
|
self.assertIn("[2/2] SUMMARY scan complete", output.getvalue())
|
|
self.assertNotIn("\x1b", output.getvalue())
|
|
self.assertNotIn("\r", output.getvalue())
|
|
|
|
def test_heartbeat_requires_inactivity_and_an_active_task(self):
|
|
output = io.StringIO()
|
|
clock = Clock()
|
|
reporter = ProgressReporter(output, clock=clock, interval=10)
|
|
reporter._render()
|
|
self.assertEqual(output.getvalue(), "")
|
|
reporter(ProgressEvent("start", "discovery", total=3))
|
|
initial = output.getvalue()
|
|
clock.now = 9
|
|
reporter._render()
|
|
self.assertEqual(output.getvalue(), initial)
|
|
clock.now = 10
|
|
reporter._render()
|
|
self.assertIn("[0/3] RUNNING discovery (elapsed 10s)", output.getvalue())
|
|
first_heartbeat = output.getvalue()
|
|
clock.now = 19
|
|
reporter._render()
|
|
self.assertEqual(output.getvalue(), first_heartbeat)
|
|
reporter(ProgressEvent("summary", "scan", status="complete", completed=3, total=3))
|
|
final = output.getvalue()
|
|
clock.now = 99
|
|
reporter._render()
|
|
self.assertEqual(output.getvalue(), final)
|
|
|
|
def test_parallel_tasks_and_out_of_order_counts(self):
|
|
output = io.StringIO()
|
|
clock = Clock()
|
|
reporter = ProgressReporter(output, clock=clock)
|
|
reporter(ProgressEvent("start", "rpcdump", "10.0.0.1", total=4))
|
|
reporter(ProgressEvent("start", "rpcdump", "10.0.0.2", total=4))
|
|
clock.now = 10
|
|
reporter._render()
|
|
self.assertIn("10.0.0.1 (+1 active)", output.getvalue())
|
|
reporter(ProgressEvent("finish", "rpcdump", "10.0.0.2", "failed", 2, 4))
|
|
reporter(ProgressEvent("finish", "rpcdump", "10.0.0.1", "ok", 1, 4))
|
|
self.assertIn("[2/4] DONE rpcdump 10.0.0.1 ok", output.getvalue())
|
|
|
|
def test_silent_discovery_advance_keeps_counts_without_printing_dead_hosts(self):
|
|
output = io.StringIO()
|
|
clock = Clock()
|
|
reporter = ProgressReporter(output, clock=clock, interval=10)
|
|
reporter(ProgressEvent("start", "tcp-discovery", total=254))
|
|
initial = output.getvalue()
|
|
clock.now = 9
|
|
reporter(ProgressEvent("advance", "tcp-discovery", "192.0.2.9", "closed", 9, 254))
|
|
self.assertEqual(output.getvalue(), initial)
|
|
clock.now = 10
|
|
reporter._render()
|
|
self.assertIn("[9/254] RUNNING tcp-discovery (elapsed 10s)", output.getvalue())
|
|
self.assertNotIn("192.0.2.9", output.getvalue())
|
|
reporter(ProgressEvent("finish", "tcp-discovery", "192.0.2.10", "alive", 10, 254))
|
|
self.assertIn("[10/254] DONE tcp-discovery 192.0.2.10 alive", output.getvalue())
|
|
|
|
def test_tty_advance_does_not_draw_dead_target_or_keep_it_active(self):
|
|
output = Terminal()
|
|
reporter = ProgressReporter(output)
|
|
reporter(ProgressEvent("start", "tcp-discovery", total=254))
|
|
initial = output.getvalue()
|
|
reporter(ProgressEvent("advance", "tcp-discovery", "192.0.2.9", "timeout", 1, 254))
|
|
self.assertEqual(output.getvalue(), initial)
|
|
self.assertNotIn("192.0.2.9", reporter._dashboard(reporter.clock()))
|
|
self.assertIn("[1/254]", reporter._dashboard(reporter.clock()))
|
|
|
|
def test_hostless_discovery_advance_preserves_aggregate_heartbeat(self):
|
|
output = io.StringIO()
|
|
clock = Clock()
|
|
reporter = ProgressReporter(output, clock=clock, interval=10)
|
|
reporter(ProgressEvent("start", "tcp-discovery", total=254))
|
|
initial = output.getvalue()
|
|
reporter(ProgressEvent("advance", "tcp-discovery", completed=1, total=254))
|
|
self.assertEqual(output.getvalue(), initial)
|
|
clock.now = 10
|
|
reporter._render()
|
|
self.assertIn("[1/254] RUNNING tcp-discovery (elapsed 10s)", output.getvalue())
|
|
|
|
def test_untrusted_text_cannot_inject_control_sequences_or_lines(self):
|
|
output = io.StringIO()
|
|
reporter = ProgressReporter(output, mode="plain")
|
|
reporter(ProgressEvent(
|
|
"finish", "check", "10.0.0.1\r\nFAKE", "\x1b[31mvulnerable\x1b[0m",
|
|
detail="one\tline\x00\u202e\nonly",
|
|
))
|
|
self.assertEqual(len(output.getvalue().splitlines()), 1)
|
|
self.assertIn("10.0.0.1 FAKE vulnerable", output.getvalue())
|
|
self.assertIn("one line only", output.getvalue())
|
|
self.assertNotIn("\x1b", output.getvalue())
|
|
self.assertNotIn("\u202e", output.getvalue())
|
|
|
|
def test_tty_dashboard_retains_completed_result(self):
|
|
output = Terminal()
|
|
clock = Clock()
|
|
reporter = ProgressReporter(output, clock=clock, color="never")
|
|
reporter(ProgressEvent("start", "check", "10.0.0.1", total=1))
|
|
self.assertIn("\r\x1b[2K[0/1] RUNNING check 10.0.0.1", output.getvalue())
|
|
reporter(ProgressEvent("finish", "check", "10.0.0.1", "inconclusive", 1, 1))
|
|
self.assertIn("DONE check 10.0.0.1 inconclusive (0.0s)\n", output.getvalue())
|
|
reporter.__exit__(None, None, None)
|
|
self.assertTrue(output.getvalue().endswith("\n"))
|
|
|
|
def test_plain_mode_never_uses_ansi_even_on_a_terminal(self):
|
|
output = Terminal()
|
|
reporter = ProgressReporter(output, mode="plain")
|
|
reporter(ProgressEvent("start", "discovery"))
|
|
self.assertNotIn("\x1b", output.getvalue())
|
|
|
|
def test_context_starts_and_joins_heartbeat_thread(self):
|
|
output = io.StringIO()
|
|
with patch("rpcscope.progress.threading.Thread") as thread:
|
|
with ProgressReporter(output) as reporter:
|
|
reporter(ProgressEvent("start", "discovery"))
|
|
thread.return_value.start.assert_called_once()
|
|
thread.return_value.join.assert_called_once()
|
|
self.assertTrue(reporter._stop.is_set())
|
|
final = output.getvalue()
|
|
reporter(ProgressEvent("start", "ignored"))
|
|
reporter._render()
|
|
self.assertEqual(output.getvalue(), final)
|
|
|
|
def test_quiet_does_not_create_a_thread_or_output(self):
|
|
output = io.StringIO()
|
|
with patch("rpcscope.progress.threading.Thread") as thread:
|
|
with ProgressReporter(output, mode="quiet") as reporter:
|
|
reporter(ProgressEvent("start", "discovery"))
|
|
reporter._render()
|
|
thread.assert_not_called()
|
|
self.assertEqual(output.getvalue(), "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|