diff --git a/bbot/modules/gowitness.py b/bbot/modules/gowitness.py index 5d0538d834..b26018d58c 100644 --- a/bbot/modules/gowitness.py +++ b/bbot/modules/gowitness.py @@ -1,4 +1,5 @@ import os +import sys import asyncio import aiosqlite import multiprocessing @@ -243,7 +244,7 @@ async def handle_batch(self, *events): context=f"{{module}} visited {{event.type}}: {url}", ) if url_event and ip: - url_event._resolved_hosts.add(ip) + url_event._resolved_hosts.add(sys.intern(ip)) await self.emit_event(url_event) # emit technologies diff --git a/bbot/modules/httpx.py b/bbot/modules/httpx.py index 628ef4bb66..dd996a2b58 100644 --- a/bbot/modules/httpx.py +++ b/bbot/modules/httpx.py @@ -1,4 +1,5 @@ import re +import sys import orjson import tempfile import subprocess @@ -200,7 +201,7 @@ async def handle_batch(self, *events): if url_event: httpx_ip = j.get("host", "") if httpx_ip: - url_event._resolved_hosts.add(httpx_ip) + url_event._resolved_hosts.add(sys.intern(httpx_ip)) url_event.data["status_code"] = status_code title = j.get("title", "") if title: diff --git a/bbot/modules/internal/dnsresolve.py b/bbot/modules/internal/dnsresolve.py index 680cedc605..f23a3bea76 100644 --- a/bbot/modules/internal/dnsresolve.py +++ b/bbot/modules/internal/dnsresolve.py @@ -1,3 +1,4 @@ +import sys import ipaddress from contextlib import suppress @@ -231,7 +232,7 @@ def check_scope(self, event): for rdtype in ("A", "AAAA", "CNAME"): hosts = dns_children.get(rdtype, []) # update resolved hosts - event.resolved_hosts.update(hosts) + event.resolved_hosts.update(sys.intern(h) for h in hosts) for host in hosts: # having a CNAME to an in-scope host doesn't make you in-scope if rdtype != "CNAME": @@ -258,6 +259,7 @@ async def resolve_event(self, event, types): queries = [(event_host, rdtype) for rdtype in types] dns_errors = {} async for (query, rdtype), (answers, errors) in self.helpers.dns.resolve_raw_batch(queries): + rdtype = sys.intern(rdtype) # errors try: dns_errors[rdtype].update(errors) @@ -272,6 +274,8 @@ async def resolve_event(self, event, types): event.raw_dns_records[rdtype] = {answer} # hosts for _rdtype, host in extract_targets(answer): + _rdtype = sys.intern(_rdtype) + host = sys.intern(host) try: event.dns_children[_rdtype].add(host) except KeyError: diff --git a/bbot/scripts/benchmark_report.py b/bbot/scripts/benchmark_report.py index 675c99bc64..50ca6f3384 100644 --- a/bbot/scripts/benchmark_report.py +++ b/bbot/scripts/benchmark_report.py @@ -180,6 +180,7 @@ def generate_comparison_table(current_data: Dict, base_data: Dict, current_branc |--------------|---------|------------|-----------|-----------|""" significant_changes = [] + new_tests = [] performance_summary = [] for current_bench in current_benchmarks: @@ -245,10 +246,16 @@ def generate_comparison_table(current_data: Dict, base_data: Dict, current_branc else: base_ops = 1 / base_mean # Default: single operation - # Use per-event memory if available, otherwise use time + # Use memory metrics if available, otherwise use time + current_mb = current_extra.get("total_memory_mb") + base_mb = base_extra.get("total_memory_mb") current_peb = current_extra.get("per_event_bytes") base_peb = base_extra.get("per_event_bytes") - if current_peb is not None and base_peb is not None: + if current_mb is not None and base_mb is not None and current_peb is None: + change_percent, emoji = calculate_change_percentage(base_mb, current_mb) + base_label = f"{base_mb:.1f} MB" + current_label = f"{current_mb:.1f} MB" + elif current_peb is not None and base_peb is not None: change_percent, emoji = calculate_change_percentage(base_peb, current_peb) base_label = f"{base_peb:.0f} B/event" current_label = f"{current_peb:.0f} B/event" @@ -269,7 +276,10 @@ def generate_comparison_table(current_data: Dict, base_data: Dict, current_branc # Track significant changes if abs(change_percent) > 10: - is_memory = current_extra.get("per_event_bytes") is not None + is_memory = ( + current_extra.get("per_event_bytes") is not None + or current_extra.get("total_memory_mb") is not None + ) if is_memory: direction = "🐌 more memory" if change_percent > 0 else "🚀 less memory" else: @@ -295,9 +305,7 @@ def generate_comparison_table(current_data: Dict, base_data: Dict, current_branc else: table += f"\n| **{test_name}** | `-` | `{format_time(current_mean)}` | **New** 🆕 | 🆕 |" - significant_changes.append( - f"- **{test_name}**: New test 🆕 ({format_time(current_mean)}, {format_ops(current_ops)})" - ) + new_tests.append(f"- **{test_name}**: {format_time(current_mean)}, {format_ops(current_ops)}") table += "\n\n\n\n" @@ -323,6 +331,13 @@ def generate_comparison_table(current_data: Dict, base_data: Dict, current_branc table += f"{change}\n" table += "\n" + # Add new tests section + if new_tests: + table += "### 🆕 New Tests\n\n" + for new_test in new_tests: + table += f"{new_test}\n" + table += "\n" + return table diff --git a/bbot/test/benchmarks/_scan_memory_subdomain_enum.py b/bbot/test/benchmarks/_scan_memory_subdomain_enum.py new file mode 100644 index 0000000000..33b4b997f6 --- /dev/null +++ b/bbot/test/benchmarks/_scan_memory_subdomain_enum.py @@ -0,0 +1,62 @@ +""" +Subprocess script for subdomain enumeration memory benchmark. + +Injects SUBDOMAIN_ENUM_COUNT synthetic DNS_NAME events into a scan +and prints peak tracemalloc memory to stdout. + +Invoked by test_scan_memory.py — not meant to be run directly. +""" + +import gc +import sys +import asyncio +import tracemalloc + +from bbot.scanner import Scanner + +SUBDOMAIN_ENUM_COUNT = int(sys.argv[1]) + +scan = Scanner( + "blacklanternsecurity.com", + modules=[], + output_modules=["python"], + config={ + "dns": {"disable": True}, + "scope": {"search_distance": 0}, + "web": {"spider_distance": 0, "spider_depth": 0}, + "speculate": False, + "excavate": True, + "aggregate": False, + "cloudcheck": False, + }, + force_start=True, +) + + +async def run(): + await scan._prep() + gc.collect() + if tracemalloc.is_tracing(): + tracemalloc.stop() + tracemalloc.start() + events = [] + injected = False + async for event in scan.async_start(): + events.append(event) + if event.type == "SCAN" and not injected: + injected = True + root_event = scan.root_event + for i in range(SUBDOMAIN_ENUM_COUNT): + dns_event = scan.make_event( + f"sub{i}.blacklanternsecurity.com", + "DNS_NAME", + parent=root_event, + context=f"benchmark DNS_NAME {i}", + ) + await scan.ingress_module.queue_event(dns_event, {}) + + +asyncio.run(run()) +_, peak = tracemalloc.get_traced_memory() +tracemalloc.stop() +print(f"PEAK_MB:{round(peak / 1024 / 1024, 2)}") diff --git a/bbot/test/benchmarks/_scan_memory_web_crawl.py b/bbot/test/benchmarks/_scan_memory_web_crawl.py new file mode 100644 index 0000000000..e609220c74 --- /dev/null +++ b/bbot/test/benchmarks/_scan_memory_web_crawl.py @@ -0,0 +1,86 @@ +""" +Subprocess script for web crawl memory benchmark. + +Launches a local HTTP server with NUM_PAGES pages (each BODY_SIZE bytes), +runs a BBOT scan against it, and prints peak tracemalloc memory to stdout. + +Invoked by test_scan_memory.py — not meant to be run directly. +""" + +import gc +import sys +import asyncio +import threading +import tracemalloc +import importlib.util +from http.server import HTTPServer, BaseHTTPRequestHandler + +from bbot.scanner import Scanner + +NUM_PAGES = int(sys.argv[1]) +BODY_SIZE = int(sys.argv[2]) + +HTTP_MODULE = "httpx" if importlib.util.find_spec("bbot.modules.httpx") else "http" + + +class H(BaseHTTPRequestHandler): + def do_GET(self): + if self.path == "/": + links = "".join(f'page{i}' for i in range(NUM_PAGES)) + body = "" + links + "" + elif self.path.startswith("/page"): + i = self.path.replace("/page", "") + links = f'infodetails' + body = "

Page " + i + "

" + links + "A" * BODY_SIZE + "" + elif self.path.startswith("/data"): + body = "data endpoint" + else: + self.send_response(404) + self.end_headers() + return + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + self.wfile.write(body.encode()) + + def log_message(self, *a): + pass + + +server = HTTPServer(("127.0.0.1", 0), H) +port = server.server_address[1] +threading.Thread(target=server.serve_forever, daemon=True).start() + +scan = Scanner( + f"http://127.0.0.1:{port}/", + modules=[HTTP_MODULE], + output_modules=["python"], + config={ + "dns": {"disable": True}, + "scope": {"search_distance": 0}, + "web": {"spider_distance": 10, "spider_depth": 10, "spider_links_per_page": NUM_PAGES}, + "speculate": True, + "excavate": True, + "aggregate": False, + "cloudcheck": False, + }, + force_start=True, +) + + +async def run(): + await scan._prep() + gc.collect() + if tracemalloc.is_tracing(): + tracemalloc.stop() + tracemalloc.start() + events = [] + async for event in scan.async_start(): + events.append(event) + + +asyncio.run(run()) +_, peak = tracemalloc.get_traced_memory() +tracemalloc.stop() +server.shutdown() +print(f"PEAK_MB:{round(peak / 1024 / 1024, 2)}") diff --git a/bbot/test/benchmarks/test_scan_memory.py b/bbot/test/benchmarks/test_scan_memory.py new file mode 100644 index 0000000000..af4f76dfa6 --- /dev/null +++ b/bbot/test/benchmarks/test_scan_memory.py @@ -0,0 +1,66 @@ +""" +Memory benchmarks for BBOT scan patterns. + +Each benchmark launches a scan as a subprocess so tracemalloc measurements +are not contaminated by pytest's own allocations. The subprocess writes +peak memory (MB) to stdout, which the test reads and stores in +benchmark extra_info["total_memory_mb"]. +""" + +import subprocess +import sys +from pathlib import Path + +import pytest + + +NUM_PAGES = 500 +BODY_SIZE = 500_000 # 500 KB per page +SUBDOMAIN_ENUM_COUNT = 5000 + +_BENCHMARKS_DIR = Path(__file__).parent + + +def _run_scan_subprocess(script_path: Path, *args: str) -> float: + """Run a scan script in a clean subprocess, return peak memory in MB.""" + result = subprocess.run( + [sys.executable, str(script_path), *args], + capture_output=True, + text=True, + timeout=600, + ) + if result.returncode != 0: + raise RuntimeError(f"Scan subprocess failed:\n{result.stderr[-2000:]}") + for line in result.stdout.strip().splitlines(): + if line.startswith("PEAK_MB:"): + return float(line.split(":", 1)[1]) + raise RuntimeError(f"No PEAK_MB in subprocess output:\n{result.stdout[-2000:]}") + + +class TestWebCrawlMemory: + """Measures peak memory during a realistic web crawl with large pages.""" + + @pytest.mark.benchmark(group="memory_scan_patterns") + def test_memory_use_web_crawl(self, benchmark): + peak_mb = _run_scan_subprocess( + _BENCHMARKS_DIR / "_scan_memory_web_crawl.py", + str(NUM_PAGES), + str(BODY_SIZE), + ) + benchmark.extra_info["total_memory_mb"] = peak_mb + benchmark.extra_info["num_pages"] = NUM_PAGES + benchmark.pedantic(lambda: None, iterations=1, rounds=1, warmup_rounds=0) + + +class TestSubdomainEnumMemory: + """Measures peak memory during a large subdomain enumeration.""" + + @pytest.mark.benchmark(group="memory_scan_patterns") + def test_memory_use_subdomain_enum(self, benchmark): + peak_mb = _run_scan_subprocess( + _BENCHMARKS_DIR / "_scan_memory_subdomain_enum.py", + str(SUBDOMAIN_ENUM_COUNT), + ) + benchmark.extra_info["total_memory_mb"] = peak_mb + benchmark.extra_info["num_subdomains"] = SUBDOMAIN_ENUM_COUNT + benchmark.pedantic(lambda: None, iterations=1, rounds=1, warmup_rounds=0)