From 4fef2c5e3be688448a270fbeef2fa50a4becefc6 Mon Sep 17 00:00:00 2001 From: Matt <36310667+NotoriousRebel@users.noreply.github.com> Date: Fri, 14 Aug 2026 23:10:34 -0400 Subject: [PATCH] refactor: bound screenshot lifecycle ownership (#2543) * refactor: bound screenshot lifecycle ownership (#293) * fix: retain HTTP error screenshots (#293) --- CHANGELOG.md | 1 + pyproject.toml | 1 - tests/lib/test_run_backend.py | 74 +++-- tests/test_main.py | 401 +++++++++++++++----------- tests/test_screenshot.py | 386 +++++++++++++++++++++++++ theHarvester/__main__.py | 45 +-- theHarvester/screenshot/screenshot.py | 261 +++++++++++++---- theHarvester/theHarvester.py | 6 - uv.lock | 11 - 9 files changed, 857 insertions(+), 329 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d817616..ce746692 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - Replaced Shodan's synchronous Python SDK with cancellable async Host API requests that honor configured proxies, query every unique resolved IPv4, paginate target-bound hostname and TLS-certificate searches without an adapter-specific result cap, retain successful partial results, and add no source-local deadline. Shodan now stores one canonical `shodan-host` result per IP with every normalized TCP or UDP service and scoped certificate CN/SAN metadata in native JSONL, SQLite, API, and HarvestView details instead of an escaped JSON value. +- Reworked screenshot scans to use one bounded aiohttp session and one shared browser, with isolated per-target contexts, status-based reachability, and deterministic async cleanup. - Migrated BuiltWith to the current v23 Domain API with privacy-preserving request controls, nested result parsing, and truthful partial or failed outcomes. - Migrated Censys discovery from the deprecated Python Search SDK to the Censys Platform API, using a Personal Access Token and optional organization ID. - API endpoint scans now attempt every configured path by default while retaining fixed concurrency, per-request timeouts, retries, and a per-response body limit; explicit library request and runtime limits still report partial coverage. diff --git a/pyproject.toml b/pyproject.toml index 505cfdb3..1136d0e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,6 @@ dependencies = [ "aiofiles==25.1.0", "aiohttp==3.14.1", "aiohttp-socks==0.11.0", - "aiomultiprocess==0.9.1", "aiosqlite==0.22.1", "beautifulsoup4==4.15.0", "certifi==2026.6.17", diff --git a/tests/lib/test_run_backend.py b/tests/lib/test_run_backend.py index 30d3c9c8..78d3e9ed 100644 --- a/tests/lib/test_run_backend.py +++ b/tests/lib/test_run_backend.py @@ -15,6 +15,38 @@ import pytest from fastapi.testclient import TestClient +class _FakeScreenshotBatch: + async def reachable_targets(self, targets: list[str]) -> list[tuple[str, str]]: + reachable: list[tuple[str, str]] = [] + for subject in targets: + final_url, status = await self.visit(subject) + if status: + reachable.append((subject, final_url)) + return reachable + + async def capture_targets(self, targets, record) -> None: + async def capture(subject: str, final_url: str) -> tuple[str, str, Path]: + output_path = self.screenshot_path(subject) + return subject, await self.take_screenshot(final_url, output_path=output_path), output_path + + tasks = [asyncio.create_task(capture(*target)) for target in targets] + try: + for task in asyncio.as_completed(tasks): + await record(*await task) + except asyncio.CancelledError: + for task in tasks: + task.cancel() + outcomes = await asyncio.gather(*tasks, return_exceptions=True) + for outcome in outcomes: + if isinstance(outcome, tuple): + await record(*outcome) + raise + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + def test_run_paths_use_one_expanded_database_and_artifact_root(tmp_path, monkeypatch) -> None: from theHarvester.lib.api.run_store import RunStore @@ -288,7 +320,7 @@ def test_child_screenshot_run_persists_downloadable_artifact_metadata(tmp_path, from theHarvester.lib.api.run_store import RunStore from theHarvester.lib.api.run_worker import _child_execute - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -297,16 +329,9 @@ def test_child_screenshot_run_persists_downloadable_artifact_metadata(tmp_path, def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return f'https://{host}', 'reachable' - @staticmethod - def chunk_list(values: list[str], _size: int) -> list[list[str]]: - return [values] - def screenshot_path(self, url: str) -> Path: return Path(self.output) / f'{url.removeprefix("https://")}.png' @@ -315,23 +340,9 @@ def test_child_screenshot_run_persists_downloadable_artifact_metadata(tmp_path, (output_path or self.screenshot_path(captured_url)).write_bytes(b'png') return captured_url - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - database = tmp_path / 'state' / 'runs.sqlite' monkeypatch.setenv('THEHARVESTER_RUN_ARTIFACTS', str(tmp_path / 'artifacts')) monkeypatch.setattr(main_module, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(main_module, 'Pool', FakePool) async def scenario(): store = RunStore(database) @@ -363,7 +374,7 @@ def test_child_screenshot_cancellation_reuses_the_checkpointed_evidence(tmp_path first_captured = asyncio.Event() - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -372,9 +383,6 @@ def test_child_screenshot_cancellation_reuses_the_checkpointed_evidence(tmp_path def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return f'https://{host}', 'reachable' @@ -400,24 +408,10 @@ def test_child_screenshot_cancellation_reuses_the_checkpointed_evidence(tmp_path async def get_hostnames(self) -> set[str]: return {'first.example.test', 'second.example.test'} - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - database = tmp_path / 'runs.sqlite' monkeypatch.setenv('THEHARVESTER_RUN_ARTIFACTS', str(tmp_path / 'artifacts')) monkeypatch.setattr(main_module.crtsh, 'SearchCrtsh', TwoHostSource) monkeypatch.setattr(main_module, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(main_module, 'Pool', FakePool) async def scenario(): store = RunStore(database) diff --git a/tests/test_main.py b/tests/test_main.py index c2c25d28..86e0726f 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -6,9 +6,11 @@ import xml.etree.ElementTree as ElementTree from datetime import UTC, datetime from pathlib import Path from types import ModuleType +from unittest.mock import AsyncMock, MagicMock import pytest +import theHarvester.screenshot.screenshot as screenshot_module from theHarvester import __main__ as theharvester_main from theHarvester.discovery.constants import MissingKey from theHarvester.lib.asn_attribution import AsnAttributionObservation @@ -1614,6 +1616,38 @@ async def test_dns_brute_resolver_configuration_does_not_enable_dns_resolution( assert 'dns-resolve' not in actions +class _FakeScreenshotBatch: + async def reachable_targets(self, targets: list[str]) -> list[tuple[str, str]]: + reachable: list[tuple[str, str]] = [] + for subject in targets: + final_url, status = await self.visit(subject) + if status: + reachable.append((subject, final_url)) + return reachable + + async def capture_targets(self, targets, record) -> None: + async def capture(subject: str, final_url: str) -> tuple[str, str, Path]: + output_path = self.screenshot_path(subject) + return subject, await self.take_screenshot(final_url, output_path=output_path), output_path + + tasks = [asyncio.create_task(capture(*target)) for target in targets] + try: + for task in asyncio.as_completed(tasks): + await record(*await task) + except asyncio.CancelledError: + for task in tasks: + task.cancel() + outcomes = await asyncio.gather(*tasks, return_exceptions=True) + for outcome in outcomes: + if isinstance(outcome, tuple): + await record(*outcome) + raise + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + @pytest.mark.asyncio async def test_dns_proven_cname_hosts_reach_screenshot_filter( monkeypatch: pytest.MonkeyPatch, @@ -1652,7 +1686,7 @@ async def test_dns_proven_cname_hosts_reach_screenshot_filter( ['192.0.2.1'], ) - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -1661,17 +1695,10 @@ async def test_dns_proven_cname_hosts_reach_screenshot_filter( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: visited.add(host) return host, 'https' - @staticmethod - def chunk_list(values: list[str], _size: int) -> list[list[str]]: - return [values] - async def take_screenshot(self, host: str, *, output_path: Path | None = None) -> str: path = output_path or self.screenshot_path(host) path.write_bytes(b'png') @@ -1680,24 +1707,10 @@ async def test_dns_proven_cname_hosts_reach_screenshot_filter( def screenshot_path(self, host: str) -> Path: return Path(self.output) / f'{host.removeprefix("https://")}.png' - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self) -> 'FakePool': - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', FakeResultStore) monkeypatch.setattr(theharvester_main.crtsh, 'SearchCrtsh', FakeCrtsh) monkeypatch.setattr(theharvester_main.hostchecker, 'Checker', FakeChecker) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) monkeypatch.setattr( sys, 'argv', @@ -1759,6 +1772,190 @@ def _recording_result_store(saved: list[CompletedResult]) -> type[_NoopResultSto return RecordingResultStore +@pytest.mark.asyncio +async def test_screenshot_scan_uses_one_bounded_reachability_and_capture_lifecycle( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + lifecycle_calls: list[str] = [] + + class FakeScreenShotter: + slash = '/' + + def __init__(self, output: str) -> None: + self.output = output + + def verify_path(self) -> bool: + return True + + async def verify_installation(self) -> None: + raise AssertionError('the scan must not launch a separate installation-check browser') + + async def reachable_targets(self, targets: list[str]) -> list[tuple[str, str]]: + lifecycle_calls.append('reachability') + assert targets == ['api.example.test'] + return [('api.example.test', 'https://login.example.net/session')] + + async def capture_targets(self, targets, record) -> None: + lifecycle_calls.append('capture') + assert targets == [('api.example.test', 'https://login.example.net/session')] + output_path = self.screenshot_path('api.example.test') + output_path.write_bytes(b'png') + await record('api.example.test', 'https://login.example.net/session', output_path) + + def screenshot_path(self, subject: str) -> Path: + return Path(self.output) / f'{subject}.png' + + monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) + monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) + + result = await theharvester_main.start( + EnumerationOptions(domain='api.example.test', screenshot=str(tmp_path), quiet=True), + return_completed_result=True, + ) + + assert lifecycle_calls == ['reachability', 'capture'] + execution = next(item for item in result[-1].active_evidence.executions if item.action == 'screenshot') + assert execution.status == 'completed' + assert execution.artifacts[0].subject_value == 'api.example.test' + + +@pytest.mark.asyncio +async def test_screenshot_cancellation_finishes_owned_artifact_record_before_checkpoint( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + saved: list[CompletedResult] = [] + record_started = asyncio.Event() + release_record = asyncio.Event() + original_read_bytes = theharvester_main.anyio.Path.read_bytes + + class RecordingResultStore(_NoopResultStore): + async def save_run(self, result: CompletedResult) -> None: + saved.append(result) + + class Page: + async def goto(self, *_args, **_kwargs) -> None: + return None + + async def screenshot(self, *, path: Path) -> None: + path.write_bytes(b'png') # noqa: ASYNC240 - tiny in-memory browser fixture + + async def close(self) -> None: + return None + + class Context: + async def new_page(self) -> Page: + return Page() + + async def close(self) -> None: + return None + + class Browser: + async def new_context(self) -> Context: + return Context() + + async def close(self) -> None: + return None + + class Playwright: + chromium = None + + def __init__(self) -> None: + self.chromium = self + + async def launch(self, **_kwargs) -> Browser: + return Browser() + + class Manager: + async def __aenter__(self) -> Playwright: + return Playwright() + + async def __aexit__(self, *_args) -> None: + return None + + async def reachable_targets(_self, targets: list[str]) -> list[tuple[str, str]]: + assert targets == ['api.example.test'] + return [('api.example.test', 'https://api.example.test')] + + async def delayed_read_bytes(path) -> bytes: + if path.name == 'api.example.test.png': + record_started.set() + await release_record.wait() + return await original_read_bytes(path) + + monkeypatch.setattr(theharvester_main, 'ResultStore', RecordingResultStore) + monkeypatch.setattr(screenshot_module, 'async_playwright', lambda: Manager()) + monkeypatch.setattr(screenshot_module.ScreenShotter, 'reachable_targets', reachable_targets) + monkeypatch.setattr(theharvester_main.anyio.Path, 'read_bytes', delayed_read_bytes) + + operation = asyncio.create_task( + theharvester_main.start( + EnumerationOptions(domain='api.example.test', screenshot=str(tmp_path), quiet=True), + return_completed_result=True, + ) + ) + await asyncio.wait_for(record_started.wait(), timeout=1) + operation.cancel('operator-stop') + await asyncio.sleep(0) + release_record.set() + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await operation + + execution = next(item for item in saved[-1].active_evidence.executions if item.action == 'screenshot') + assert execution.status == 'partial' + assert execution.stop_reason == 'cancelled' + assert execution.artifacts[0].subject_value == 'api.example.test' + assert not [task for task in asyncio.all_tasks() if task.get_name().startswith('screenshot-')] + + +@pytest.mark.asyncio +async def test_screenshot_cleanup_failure_marks_the_action_failed( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + page = MagicMock() + + async def write_screenshot(*, path: Path) -> None: + path.write_bytes(b'png') # noqa: ASYNC240 - tiny in-memory browser fixture + + page.goto = AsyncMock() + page.screenshot = AsyncMock(side_effect=write_screenshot) + page.close = AsyncMock(side_effect=RuntimeError('page close failed')) + context = MagicMock() + context.new_page = AsyncMock(return_value=page) + context.close = AsyncMock() + browser = MagicMock() + browser.new_context = AsyncMock(return_value=context) + browser.close = AsyncMock() + playwright = MagicMock() + playwright.chromium.launch = AsyncMock(return_value=browser) + manager = MagicMock() + manager.__aenter__ = AsyncMock(return_value=playwright) + manager.__aexit__ = AsyncMock(return_value=False) + + async def reachable_targets(_self, targets: list[str]) -> list[tuple[str, str]]: + return [(targets[0], f'https://{targets[0]}')] + + monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) + monkeypatch.setattr(screenshot_module, 'async_playwright', MagicMock(return_value=manager)) + monkeypatch.setattr(screenshot_module.ScreenShotter, 'reachable_targets', reachable_targets) + + result = await theharvester_main.start( + EnumerationOptions(domain='api.example.test', screenshot=str(tmp_path), quiet=True), + return_completed_result=True, + ) + + execution = next(item for item in result[-1].active_evidence.executions if item.action == 'screenshot') + assert execution.status == 'failed' + assert execution.stop_reason == 'capture-errors' + assert execution.error_type == 'RuntimeError' + page.close.assert_awaited_once() + context.close.assert_awaited_once() + browser.close.assert_awaited_once() + manager.__aexit__.assert_awaited_once() + + @pytest.mark.asyncio async def test_cli_can_capture_an_explicit_target_without_discovery_sources( monkeypatch: pytest.MonkeyPatch, @@ -1774,7 +1971,7 @@ async def test_cli_can_capture_an_explicit_target_without_discovery_sources( async def save_run(self, result: CompletedResult) -> None: saved.append(result) - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -1783,16 +1980,9 @@ async def test_cli_can_capture_an_explicit_target_without_discovery_sources( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return host, 'https' - @staticmethod - def chunk_list(values: list[str], _size: int) -> list[list[str]]: - return [values] - async def take_screenshot(self, host: str, *, output_path: Path | None = None) -> str: captured.append(host) (output_path or self.screenshot_path(host)).write_bytes(b'png') @@ -1801,22 +1991,8 @@ async def test_cli_can_capture_an_explicit_target_without_discovery_sources( def screenshot_path(self, host: str) -> Path: return Path(self.output) / f'{host.removeprefix("https://")}.png' - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self) -> 'FakePool': - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', FakeResultStore) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) monkeypatch.setattr( sys, 'argv', @@ -1847,7 +2023,7 @@ async def test_screenshot_reports_no_reachable_target_as_failed( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -1856,28 +2032,11 @@ async def test_screenshot_reports_no_reachable_target_as_failed( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, _host: str) -> tuple[str, str]: return '', '' - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) result = await theharvester_main.start( EnumerationOptions(domain='api.example.test', screenshot=str(tmp_path), quiet=True), @@ -1894,7 +2053,7 @@ async def test_screenshot_redirect_stays_attached_to_the_authorized_host( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -1903,9 +2062,6 @@ async def test_screenshot_redirect_stays_attached_to_the_authorized_host( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: assert host == 'api.example.test' return 'https://outside.example/path', 'reachable' @@ -1917,22 +2073,8 @@ async def test_screenshot_redirect_stays_attached_to_the_authorized_host( def screenshot_path(self, _url: str) -> Path: return Path(self.output) / 'outside.example.png' - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) result = await theharvester_main.start( EnumerationOptions(domain='api.example.test', screenshot=str(tmp_path), quiet=True), @@ -1950,7 +2092,7 @@ async def test_target_only_ip_screenshot_keeps_ip_result_and_artifact_subject( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -1959,9 +2101,6 @@ async def test_target_only_ip_screenshot_keeps_ip_result_and_artifact_subject( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return f'https://{host}', 'reachable' @@ -1973,22 +2112,8 @@ async def test_target_only_ip_screenshot_keeps_ip_result_and_artifact_subject( def screenshot_path(self, url: str) -> Path: return Path(self.output) / f'{url.removeprefix("https://")}.png' - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) result = await theharvester_main.start( EnumerationOptions(domain='192.0.2.1', screenshot=str(tmp_path), quiet=True), @@ -2008,7 +2133,7 @@ async def test_screenshot_redirects_to_one_login_keep_distinct_subject_artifacts monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -2017,9 +2142,6 @@ async def test_screenshot_redirects_to_one_login_keep_distinct_subject_artifacts def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, _host: str) -> tuple[str, str]: return 'https://login.example.net/session', 'reachable' @@ -2043,23 +2165,9 @@ async def test_screenshot_redirects_to_one_login_keep_distinct_subject_artifacts async def get_hostnames(self) -> set[str]: return {'first.example.test', 'second.example.test'} - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) monkeypatch.setattr(theharvester_main.crtsh, 'SearchCrtsh', TwoHostSource) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) result = await theharvester_main.start( EnumerationOptions( @@ -2090,7 +2198,7 @@ async def test_screenshot_cancellation_persists_failed_execution_and_propagates( async def save_run(self, result: CompletedResult) -> None: saved.append(result) - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -2099,9 +2207,6 @@ async def test_screenshot_cancellation_persists_failed_execution_and_propagates( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return f'https://{host}', 'reachable' @@ -2126,23 +2231,9 @@ async def test_screenshot_cancellation_persists_failed_execution_and_propagates( async def get_hostnames(self) -> set[str]: return {'first.example.test', 'second.example.test'} - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', RecordingResultStore) monkeypatch.setattr(theharvester_main.crtsh, 'SearchCrtsh', TwoHostSource) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) task = asyncio.create_task( theharvester_main.start( @@ -2173,7 +2264,7 @@ async def test_screenshot_capture_failure_cancels_sibling_tasks( ) -> None: sibling_cancelled = asyncio.Event() - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -2182,9 +2273,6 @@ async def test_screenshot_capture_failure_cancels_sibling_tasks( def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return f'https://{host}', 'reachable' @@ -2210,23 +2298,9 @@ async def test_screenshot_capture_failure_cancels_sibling_tasks( async def get_hostnames(self) -> set[str]: return {'first.example.test', 'second.example.test'} - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - monkeypatch.setattr(theharvester_main, 'ResultStore', _NoopResultStore) monkeypatch.setattr(theharvester_main.crtsh, 'SearchCrtsh', TwoHostSource) monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) result = await theharvester_main.start( EnumerationOptions( @@ -2292,7 +2366,7 @@ async def test_direct_action_evidence_reaches_completed_result(monkeypatch: pyte async def get_takeover_results(self) -> dict[str, list[dict[str, str]]]: return {'https://api.example.com': [{'No such app': 'Heroku'}]} - class FakeScreenShotter: + class FakeScreenShotter(_FakeScreenshotBatch): slash = '/' def __init__(self, output: str) -> None: @@ -2301,16 +2375,9 @@ async def test_direct_action_evidence_reaches_completed_result(monkeypatch: pyte def verify_path(self) -> bool: return True - async def verify_installation(self) -> None: - return None - async def visit(self, host: str) -> tuple[str, str]: return host, 'reachable' - @staticmethod - def chunk_list(values: list[str], _size: int) -> list[list[str]]: - return [values] - async def take_screenshot(self, host: str, *, output_path: Path | None = None) -> str: (output_path or self.screenshot_path(host)).write_bytes(b'png') return f'https://{host}' @@ -2386,19 +2453,6 @@ async def test_direct_action_evidence_reaches_completed_result(monkeypatch: pyte def get_status_codes(self) -> set[int]: return {200} - class FakePool: - def __init__(self, _workers: int) -> None: - pass - - async def __aenter__(self) -> 'FakePool': - return self - - async def __aexit__(self, *_args) -> None: - return None - - async def map(self, function, values): - return [await function(value) for value in values] - async def no_sleep(_seconds: float) -> None: return None @@ -2411,7 +2465,6 @@ async def test_direct_action_evidence_reaches_completed_result(monkeypatch: pyte monkeypatch.setattr(theharvester_main, 'ScreenShotter', FakeScreenShotter) monkeypatch.setattr(theharvester_main.shodansearch, 'SearchShodan', FakeShodan) monkeypatch.setattr(theharvester_main.api_endpoints, 'SearchApiEndpoints', FakeApiScanner) - monkeypatch.setattr(theharvester_main, 'Pool', FakePool) monkeypatch.setattr(theharvester_main.asyncio, 'sleep', no_sleep) result = await theharvester_main.start( diff --git a/tests/test_screenshot.py b/tests/test_screenshot.py index 5013ccd8..1747bcf5 100644 --- a/tests/test_screenshot.py +++ b/tests/test_screenshot.py @@ -1,7 +1,10 @@ from __future__ import annotations +import asyncio import stat +from contextlib import AbstractAsyncContextManager from pathlib import Path +from typing import TYPE_CHECKING, Self from unittest.mock import AsyncMock, MagicMock import pytest @@ -9,6 +12,386 @@ import pytest import theHarvester.screenshot.screenshot as screenshot_module from theHarvester.screenshot.screenshot import ScreenShotter +if TYPE_CHECKING: + from types import TracebackType + + +class _EmptyResponse(AbstractAsyncContextManager): + def __init__(self, url: str, active_requests: list[int], status: int = 204) -> None: + self.url = url + self.status = status + self._active_requests = active_requests + + async def __aenter__(self) -> Self: + self._active_requests[0] += 1 + self._active_requests[1] = max(self._active_requests) + await asyncio.sleep(0) + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self._active_requests[0] -= 1 + + +@pytest.mark.asyncio +async def test_reachable_targets_reuses_one_session_and_bounds_empty_responses( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + active_requests = [0, 0] + session = MagicMock() + session.get.side_effect = lambda url, **_kwargs: _EmptyResponse(url, active_requests) + session.__aenter__ = AsyncMock(return_value=session) + session.__aexit__ = AsyncMock(return_value=False) + session_factory = MagicMock(return_value=session) + monkeypatch.setattr(screenshot_module.aiohttp, 'ClientSession', session_factory) + monkeypatch.setattr(screenshot_module.aiohttp, 'TCPConnector', MagicMock()) + monkeypatch.setattr(screenshot_module.ssl, 'create_default_context', MagicMock()) + targets = [f'host-{index}.example.test' for index in range(25)] + + reachable = await ScreenShotter(str(tmp_path)).reachable_targets(targets) + + assert reachable == [(target, f'https://{target}') for target in targets] + assert session_factory.call_count == 1 + assert session.get.call_count == 25 + assert active_requests == [0, 20] + + +@pytest.mark.asyncio +async def test_reachable_targets_retains_completed_http_error_responses( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + active_requests = [0, 0] + statuses = { + 'blocked.example.test': 403, + 'unavailable.example.test': 503, + } + session = MagicMock() + session.get.side_effect = lambda url, **_kwargs: _EmptyResponse( + url, + active_requests, + statuses[url.removeprefix('https://').removeprefix('http://')], + ) + session.__aenter__ = AsyncMock(return_value=session) + session.__aexit__ = AsyncMock(return_value=False) + monkeypatch.setattr(screenshot_module.aiohttp, 'ClientSession', MagicMock(return_value=session)) + monkeypatch.setattr(screenshot_module.aiohttp, 'TCPConnector', MagicMock()) + monkeypatch.setattr(screenshot_module.ssl, 'create_default_context', MagicMock()) + targets = ['blocked.example.test', 'unavailable.example.test'] + + reachable = await ScreenShotter(str(tmp_path)).reachable_targets(targets) + + assert reachable == [(target, f'https://{target}') for target in targets] + + +@pytest.mark.asyncio +async def test_reachable_targets_cancellation_closes_session_and_leaves_no_tasks( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + started = asyncio.Event() + active_requests = 0 + + class PendingResponse(AbstractAsyncContextManager): + async def __aenter__(self) -> Self: + nonlocal active_requests + active_requests += 1 + try: + if active_requests == 20: + started.set() + await asyncio.Event().wait() + raise AssertionError('unreachable') + finally: + active_requests -= 1 + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + return None + + session = MagicMock() + session.get.return_value = PendingResponse() + session.__aenter__ = AsyncMock(return_value=session) + session.__aexit__ = AsyncMock(return_value=False) + monkeypatch.setattr(screenshot_module.aiohttp, 'ClientSession', MagicMock(return_value=session)) + monkeypatch.setattr(screenshot_module.aiohttp, 'TCPConnector', MagicMock()) + monkeypatch.setattr(screenshot_module.ssl, 'create_default_context', MagicMock()) + targets = [f'host-{index}.example.test' for index in range(25)] + + operation = asyncio.create_task(ScreenShotter(str(tmp_path)).reachable_targets(targets)) + await asyncio.wait_for(started.wait(), timeout=1) + operation.cancel() + with pytest.raises(asyncio.CancelledError): + await operation + + session.__aexit__.assert_awaited_once() + assert active_requests == 0 + assert not [task for task in asyncio.all_tasks() if task.get_name().startswith('screenshot-reachability-')] + + +@pytest.mark.asyncio +async def test_capture_targets_reuses_one_browser_and_bounds_isolated_contexts( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + context_counts = [0, 0] + page_close_count = 0 + context_close_count = 0 + browser_close_count = 0 + manager_exit_count = 0 + + class Page: + async def goto(self, _url: str, **kwargs: int) -> None: + assert kwargs == {'timeout': 35000} + await asyncio.sleep(0) + + async def screenshot(self, *, path: Path) -> None: + assert path.parent == tmp_path + + async def close(self) -> None: + nonlocal page_close_count + page_close_count += 1 + + class Context: + async def new_page(self) -> Page: + return Page() + + async def close(self) -> None: + nonlocal context_close_count + context_close_count += 1 + context_counts[0] -= 1 + + class Browser: + async def new_context(self) -> Context: + context_counts[0] += 1 + context_counts[1] = max(context_counts) + return Context() + + async def close(self) -> None: + nonlocal browser_close_count + browser_close_count += 1 + + class Chromium: + def __init__(self) -> None: + self.launch_count = 0 + + async def launch(self, *, headless: bool) -> Browser: + assert headless is True + self.launch_count += 1 + return Browser() + + class Playwright: + def __init__(self) -> None: + self.chromium = Chromium() + + playwright = Playwright() + + class Manager: + async def __aenter__(self) -> Playwright: + return playwright + + async def __aexit__(self, *_args: object) -> None: + nonlocal manager_exit_count + manager_exit_count += 1 + + monkeypatch.setattr(screenshot_module, 'async_playwright', lambda: Manager()) + monkeypatch.setattr(screenshot_module.os, 'chmod', MagicMock()) + captured: list[tuple[str, str, Path]] = [] + + async def record(subject: str, captured_url: str, path: Path) -> None: + captured.append((subject, captured_url, path)) + + targets = [(f'host-{index}.example.test', f'https://host-{index}.example.test') for index in range(8)] + await ScreenShotter(str(tmp_path)).capture_targets(targets, record) + + assert [item[:2] for item in captured] == targets + assert playwright.chromium.launch_count == 1 + assert context_counts == [0, 3] + assert page_close_count == 8 + assert context_close_count == 8 + assert browser_close_count == 1 + assert manager_exit_count == 1 + + +@pytest.mark.parametrize( + 'failure_stage', + [ + 'launch', + 'new-context', + 'new-page', + 'goto', + 'goto-timeout', + 'screenshot', + 'record', + 'page-close', + 'context-close', + 'browser-close', + ], +) +@pytest.mark.asyncio +async def test_capture_targets_closes_resources_after_each_failure_stage( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + failure_stage: str, +) -> None: + page = MagicMock() + goto_error: Exception | None = None + if failure_stage == 'goto': + goto_error = RuntimeError('goto failed') + elif failure_stage == 'goto-timeout': + goto_error = TimeoutError('goto timed out') + page.goto = AsyncMock(side_effect=goto_error) + page.screenshot = AsyncMock(side_effect=RuntimeError('screenshot failed') if failure_stage == 'screenshot' else None) + page.close = AsyncMock(side_effect=RuntimeError('page close failed') if failure_stage == 'page-close' else None) + context = MagicMock() + context.new_page = AsyncMock( + side_effect=RuntimeError('new page failed') if failure_stage == 'new-page' else None, + return_value=page, + ) + context.close = AsyncMock(side_effect=RuntimeError('context close failed') if failure_stage == 'context-close' else None) + browser = MagicMock() + browser.new_context = AsyncMock( + side_effect=RuntimeError('new context failed') if failure_stage == 'new-context' else None, + return_value=context, + ) + browser.close = AsyncMock(side_effect=RuntimeError('browser close failed') if failure_stage == 'browser-close' else None) + playwright = MagicMock() + playwright.chromium.launch = AsyncMock( + side_effect=RuntimeError('launch failed') if failure_stage == 'launch' else None, + return_value=browser, + ) + manager = MagicMock() + manager.__aenter__ = AsyncMock(return_value=playwright) + manager.__aexit__ = AsyncMock(return_value=False) + monkeypatch.setattr(screenshot_module, 'async_playwright', MagicMock(return_value=manager)) + monkeypatch.setattr(screenshot_module.os, 'chmod', MagicMock()) + + async def record(_subject: str, _captured_url: str, _path: Path) -> None: + if failure_stage == 'record': + raise RuntimeError('record failed') + + operation = ScreenShotter(str(tmp_path)).capture_targets([('example.test', 'https://example.test')], record) + if failure_stage in {'launch', 'record', 'page-close', 'context-close', 'browser-close'}: + with pytest.raises(RuntimeError): + await operation + else: + await operation + + manager.__aexit__.assert_awaited_once() + if failure_stage == 'launch': + browser.close.assert_not_awaited() + return + browser.close.assert_awaited_once() + if failure_stage == 'new-context': + context.close.assert_not_awaited() + page.close.assert_not_awaited() + elif failure_stage == 'new-page': + context.close.assert_awaited_once() + page.close.assert_not_awaited() + else: + context.close.assert_awaited_once() + page.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_capture_targets_cancellation_closes_resources_and_leaves_no_tasks( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + started = asyncio.Event() + pages: list[MagicMock] = [] + contexts: list[MagicMock] = [] + + async def wait_for_cancellation(*_args: object, **_kwargs: object) -> None: + if len(pages) == 3: + started.set() + await asyncio.Event().wait() + + async def new_context() -> MagicMock: + page = MagicMock() + page.goto = AsyncMock(side_effect=wait_for_cancellation) + page.screenshot = AsyncMock() + page.close = AsyncMock() + pages.append(page) + context = MagicMock() + context.new_page = AsyncMock(return_value=page) + context.close = AsyncMock() + contexts.append(context) + return context + + browser = MagicMock() + browser.new_context = AsyncMock(side_effect=new_context) + browser.close = AsyncMock() + playwright = MagicMock() + playwright.chromium.launch = AsyncMock(return_value=browser) + manager = MagicMock() + manager.__aenter__ = AsyncMock(return_value=playwright) + manager.__aexit__ = AsyncMock(return_value=False) + monkeypatch.setattr(screenshot_module, 'async_playwright', MagicMock(return_value=manager)) + + async def record(_subject: str, _captured_url: str, _path: Path) -> None: + raise AssertionError('cancelled captures must not be recorded') + + targets = [(f'host-{index}.example.test', f'https://host-{index}.example.test') for index in range(5)] + operation = asyncio.create_task(ScreenShotter(str(tmp_path)).capture_targets(targets, record)) + await asyncio.wait_for(started.wait(), timeout=1) + operation.cancel() + with pytest.raises(asyncio.CancelledError): + await operation + + assert len(pages) == len(contexts) == 3 + assert all(page.close.await_count == 1 for page in pages) + assert all(context.close.await_count == 1 for context in contexts) + browser.close.assert_awaited_once() + manager.__aexit__.assert_awaited_once() + assert not [task for task in asyncio.all_tasks() if task.get_name().startswith('screenshot-capture-')] + + +@pytest.mark.asyncio +async def test_capture_cleanup_preserves_cancellation_and_attempts_every_resource( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cancellation = asyncio.CancelledError('page-close-cancelled') + page = MagicMock() + page.goto = AsyncMock() + page.screenshot = AsyncMock() + page.close = AsyncMock(side_effect=cancellation) + context = MagicMock() + context.new_page = AsyncMock(return_value=page) + context.close = AsyncMock(side_effect=BaseException('context-close-failed')) + browser = MagicMock() + browser.new_context = AsyncMock(return_value=context) + browser.close = AsyncMock(side_effect=BaseException('browser-close-failed')) + playwright = MagicMock() + playwright.chromium.launch = AsyncMock(return_value=browser) + manager = MagicMock() + manager.__aenter__ = AsyncMock(return_value=playwright) + manager.__aexit__ = AsyncMock(side_effect=BaseException('playwright-close-failed')) + monkeypatch.setattr(screenshot_module, 'async_playwright', MagicMock(return_value=manager)) + monkeypatch.setattr(screenshot_module.os, 'chmod', MagicMock()) + + async def record(_subject: str, _captured_url: str, _path: Path) -> None: + raise AssertionError('cleanup cancellation must prevent recording') + + with pytest.raises(asyncio.CancelledError, match='page-close-cancelled') as error: + await ScreenShotter(str(tmp_path)).capture_targets([('example.test', 'https://example.test')], record) + + assert error.value is cancellation + page.close.assert_awaited_once() + context.close.assert_awaited_once() + browser.close.assert_awaited_once() + manager.__aexit__.assert_awaited_once() + assert not [task for task in asyncio.all_tasks() if task.get_name().startswith('screenshot-')] + @pytest.mark.parametrize( ('platform', 'expected_separator'), @@ -53,6 +436,7 @@ async def test_visit_prefers_https_and_normalizes_the_target( ) -> None: response = MagicMock() response.url = response_url + response.status = 200 response.__aenter__ = AsyncMock(return_value=response) response.__aexit__ = AsyncMock(return_value=False) response.text = AsyncMock(return_value='reachable') @@ -73,6 +457,7 @@ async def test_visit_prefers_https_and_normalizes_the_target( async def test_visit_falls_back_to_http_when_https_is_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: response = MagicMock() response.url = 'http://www.example.com' + response.status = 200 response.__aenter__ = AsyncMock(return_value=response) response.__aexit__ = AsyncMock(return_value=False) response.text = AsyncMock(return_value='reachable') @@ -160,6 +545,7 @@ async def test_take_screenshot_can_name_the_artifact_for_the_authorized_subject( async def test_screenshot_visit_uses_an_http_proxy(monkeypatch) -> None: response = MagicMock() response.url = 'https://example.com' + response.status = 200 response.__aenter__ = AsyncMock(return_value=response) response.__aexit__ = AsyncMock(return_value=False) response.text = AsyncMock(return_value='reachable') diff --git a/theHarvester/__main__.py b/theHarvester/__main__.py index dfef0139..bf99637a 100644 --- a/theHarvester/__main__.py +++ b/theHarvester/__main__.py @@ -22,7 +22,6 @@ from uuid import UUID, uuid4 import anyio import netaddr import ujson -from aiomultiprocess import Pool from theHarvester.discovery import ( api_endpoints, @@ -2686,11 +2685,6 @@ async def start( ) ) else: - try: - await screen_shotter.verify_installation() - except asyncio.CancelledError: - await persist_screenshot_cancellation() - raise output_logger.info(f'\nScreenshots can be found in: {screen_shotter.output}{screen_shotter.slash}') output_logger.info('Filtering domains for ones we can reach') if not engines: @@ -2708,26 +2702,6 @@ async def start( # First filter out ones that didn't resolve output_logger.info('Attempting to visit unique resolved domains, this is ACTIVE RECON') - async def visit_screenshot_target(host: str) -> tuple[str, str]: - final_url, body = await screen_shotter.visit(host) - return host, final_url if body else '' - - async with Pool(10) as pool: - try: - results = await pool.map(visit_screenshot_target, list(unique_resolved_domains)) - except asyncio.CancelledError: - await persist_screenshot_cancellation() - raise - reachable_targets = sorted((host, final_url) for host, final_url in results if final_url) - - semaphore = asyncio.Semaphore(3) - - async def capture_screenshot_target(target: tuple[str, str]) -> tuple[str, str, Path]: - subject, final_url = target - output_path = screen_shotter.screenshot_path(subject) - async with semaphore: - return subject, await screen_shotter.take_screenshot(final_url, output_path=output_path), output_path - async def record_screenshot_artifact(subject: str, captured_url: str, screenshot_path: Path) -> None: if not captured_url: capture_error_types.add('CaptureError') @@ -2770,26 +2744,16 @@ async def start( ) ) - capture_tasks = [asyncio.create_task(capture_screenshot_target(target)) for target in reachable_targets] try: - for capture_task in asyncio.as_completed(capture_tasks): - subject, captured_url, screenshot_path = await capture_task - await record_screenshot_artifact(subject, captured_url, screenshot_path) + reachable_targets = await screen_shotter.reachable_targets(sorted(unique_resolved_domains)) + if reachable_targets: + await screen_shotter.capture_targets(reachable_targets, record_screenshot_artifact) except asyncio.CancelledError: - for capture_task in capture_tasks: - capture_task.cancel() - outcomes = await asyncio.gather(*capture_tasks, return_exceptions=True) - for outcome in outcomes: - if isinstance(outcome, tuple): - await record_screenshot_artifact(*outcome) await persist_screenshot_cancellation() raise except Exception as ee: - for capture_task in capture_tasks: - capture_task.cancel() - await asyncio.gather(*capture_tasks, return_exceptions=True) capture_error_types.add(type(ee).__name__) - output_logger.info(f'An exception has occurred while mapping: {ee}') + output_logger.info(f'An exception has occurred while taking screenshots: {ee}') if not unique_resolved_domains: screenshot_status: ExecutionStatus = 'skipped' screenshot_stop_reason = 'no-input' @@ -2824,7 +2788,6 @@ async def start( hr, mon = divmod(mon, 60) total_time = f'{mon:02d}:{sec:02d}' output_logger.info(f'Finished taking screenshots in {total_time} seconds') - output_logger.info('[+] Note there may be leftover chrome processes you may have to kill manually\n') # Shodan shodanres: list[dict[str, object]] = [] diff --git a/theHarvester/screenshot/screenshot.py b/theHarvester/screenshot/screenshot.py index 6fd99a90..a9733f0e 100644 --- a/theHarvester/screenshot/screenshot.py +++ b/theHarvester/screenshot/screenshot.py @@ -2,11 +2,13 @@ take screenshots """ +import asyncio import logging import os import ssl import sys -from collections.abc import Collection +from collections.abc import AsyncIterator, Awaitable, Callable, Collection +from contextlib import asynccontextmanager from datetime import datetime from ipaddress import ip_address from pathlib import Path @@ -15,10 +17,13 @@ from urllib.parse import urlsplit import aiohttp import certifi from aiohttp_socks import ProxyConnector -from playwright.async_api import async_playwright +from playwright.async_api import Browser, async_playwright logger = logging.getLogger(__name__) +REACHABILITY_LIMIT = 20 +CAPTURE_LIMIT = 3 + def _target_url(value: str, scheme: str = 'https') -> str: if value.startswith(('http://', 'https://')): @@ -53,74 +58,218 @@ class ScreenShotter: return False @staticmethod - async def verify_installation() -> None: - # Helper function that verifies playwright & chromium is installed - try: - async with async_playwright() as p: - browser = await p.chromium.launch() - await browser.close() - logger.info('Playwright and Chromium are successfully installed.') - except Exception as e: - logger.info(f'An exception has occurred while attempting to verify installation: {e}') + async def _visit( + session: aiohttp.ClientSession, + url: str, + proxy: str | None = None, + ) -> tuple[str, str]: + urls = (url,) if url.startswith(('http://', 'https://')) else (_target_url(url), _target_url(url, 'http')) + for candidate in urls: + try: + async with session.get(candidate, proxy=proxy) as response: + return str(response.url), 'reachable' + except (aiohttp.ClientError, TimeoutError) as e: + logger.info(f'An exception has occurred while attempting to visit {candidate} : {e}') + return '', '' @staticmethod - def chunk_list(items: Collection, chunk_size: int) -> list: - # Based off of: https://github.com/apache/incubator-sdap-ingester - return [list(items)[i : i + chunk_size] for i in range(0, len(items), chunk_size)] + def _connector(proxy: str | None) -> tuple[ProxyConnector | aiohttp.TCPConnector, str | None]: + sslcontext = ssl.create_default_context(cafile=certifi.where()) + if proxy and proxy.startswith('socks5://'): + return ProxyConnector.from_url(proxy, ssl=sslcontext), None + return aiohttp.TCPConnector(ssl=sslcontext), proxy - @staticmethod - async def visit(url: str, proxy: str | None = None) -> tuple[str, str]: + @classmethod + async def visit(cls, url: str, proxy: str | None = None) -> tuple[str, str]: try: timeout = aiohttp.ClientTimeout(total=35) - urls = (url,) if url.startswith(('http://', 'https://')) else (_target_url(url), _target_url(url, 'http')) - sslcontext = ssl.create_default_context(cafile=certifi.where()) - connector: ProxyConnector | aiohttp.TCPConnector - proxy_param = None - if proxy and proxy.startswith('socks5://'): - connector = ProxyConnector.from_url(proxy, ssl=sslcontext) - else: - connector = aiohttp.TCPConnector(ssl=sslcontext) - proxy_param = proxy - + connector, proxy_param = cls._connector(proxy) async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: - for candidate in urls: - try: - async with session.get(candidate, proxy=proxy_param) as resp: - text = await resp.text('UTF-8') - return str(resp.url), text - except (aiohttp.ClientError, TimeoutError) as e: - logger.info(f'An exception has occurred while attempting to visit {candidate} : {e}') - return '', '' + return await cls._visit(session, url, proxy_param) except Exception as e: logger.info(f'An exception has occurred while attempting to visit {url} : {e}') return '', '' - async def take_screenshot(self, url: str, *, output_path: Path | None = None) -> str: - url = _target_url(url) - logger.info(f'Attempting to take a screenshot of: {url}') - async with async_playwright() as p: - browser = await p.chromium.launch(headless=True) - # New browser context + async def reachable_targets( + self, + targets: Collection[str], + proxy: str | None = None, + ) -> list[tuple[str, str]]: + indexed_targets = list(enumerate(targets)) + if not indexed_targets: + return [] + target_iterator = iter(indexed_targets) + reachable: dict[int, tuple[str, str]] = {} + timeout = aiohttp.ClientTimeout(total=35) + connector, proxy_param = self._connector(proxy) + + async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session: + + async def worker() -> None: + while True: + try: + index, target = next(target_iterator) + except StopIteration: + return + final_url, status = await self._visit(session, target, proxy_param) + if status: + reachable[index] = (target, final_url) + + tasks = [ + asyncio.create_task(worker(), name=f'screenshot-reachability-{index}') + for index in range(min(REACHABILITY_LIMIT, len(indexed_targets))) + ] + try: + await asyncio.gather(*tasks) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + return [reachable[index] for index in sorted(reachable)] + + @staticmethod + async def _close(resource: object, resource_name: str) -> BaseException | None: + try: + await resource.close() # type: ignore[attr-defined] + except BaseException as error: + logger.info(f'An exception occurred while closing screenshot {resource_name}: {error}') + return error + return None + + @staticmethod + def _cleanup_failure( + primary: BaseException | None, + cleanup_errors: Collection[BaseException], + ) -> BaseException | None: + if isinstance(primary, asyncio.CancelledError): + return primary + cancellation = next((error for error in cleanup_errors if isinstance(error, asyncio.CancelledError)), None) + return cancellation or primary or next(iter(cleanup_errors), None) + + @asynccontextmanager + async def _browser(self) -> AsyncIterator[Browser]: + manager = async_playwright() + browser = None + manager_entered = False + primary_error: BaseException | None = None + cleanup_errors: list[BaseException] = [] + try: + playwright = await manager.__aenter__() + manager_entered = True + browser = await playwright.chromium.launch(headless=True) + yield browser + except BaseException as error: + primary_error = error + finally: + if browser is not None: + if close_error := await self._close(browser, 'browser'): + cleanup_errors.append(close_error) + if manager_entered: + exit_error = self._cleanup_failure(primary_error, cleanup_errors) + try: + await manager.__aexit__( + type(exit_error) if exit_error is not None else None, + exit_error, + exit_error.__traceback__ if exit_error is not None else None, + ) + except BaseException as error: + logger.info(f'An exception occurred while closing screenshot Playwright: {error}') + cleanup_errors.append(error) + if final_error := self._cleanup_failure(primary_error, cleanup_errors): + raise final_error + + async def _capture(self, browser: Browser, url: str, output_path: Path) -> str: + normalized_url = _target_url(url) + logger.info(f'Attempting to take a screenshot of: {normalized_url}') + context = None + page = None + captured_url = '' + primary_error: BaseException | None = None + operation_error: Exception | None = None + cleanup_errors: list[BaseException] = [] + date = str(datetime.now()) + try: context = await browser.new_context() page = await context.new_page() - path: Path | None = output_path or self.screenshot_path(url) - date = str(datetime.now()) + await page.goto(normalized_url, timeout=35000) + await page.screenshot(path=output_path) + os.chmod(output_path, 0o600) + captured_url = normalized_url + except Exception as error: + operation_error = error + logger.info(f'An exception has occurred attempting to screenshot: {normalized_url} : {error}') + except BaseException as error: + primary_error = error + finally: + if page is not None: + if close_error := await self._close(page, 'page'): + cleanup_errors.append(close_error) + if context is not None: + if close_error := await self._close(context, 'context'): + cleanup_errors.append(close_error) + logger.info(f'{date} {normalized_url} {output_path if captured_url else None}') + if final_error := self._cleanup_failure(primary_error, cleanup_errors): + raise final_error + if operation_error is not None: + return '' + return captured_url + + async def capture_targets( + self, + targets: Collection[tuple[str, str]], + record: Callable[[str, str, Path], Awaitable[None]], + ) -> None: + target_iterator = iter(targets) + if not targets: + return + + async with self._browser() as browser: + + async def worker(worker_index: int) -> None: + async def record_capture(subject: str, captured_url: str, output_path: Path) -> None: + await record(subject, captured_url, output_path) + + while True: + try: + subject, url = next(target_iterator) + except StopIteration: + return + output_path = self.screenshot_path(subject) + captured_url = await self._capture(browser, url, output_path) + record_task: asyncio.Task[None] = asyncio.create_task( + record_capture(subject, captured_url, output_path), + name=f'screenshot-record-{worker_index}', + ) + try: + await asyncio.shield(record_task) + except asyncio.CancelledError as cancellation: + try: + while not record_task.done(): + try: + await asyncio.shield(record_task) + except asyncio.CancelledError: + continue + record_task.result() + except BaseException as error: + logger.info(f'An exception occurred while finishing screenshot artifact recording: {error}') + raise cancellation + + tasks = [ + asyncio.create_task(worker(index), name=f'screenshot-capture-{index}') + for index in range(min(CAPTURE_LIMIT, len(targets))) + ] try: - # Will fail if network idle or load event doesn't fire after - # 35s which should be handled - await page.goto(url, timeout=35000) - await page.screenshot(path=path) - if path is not None: - os.chmod(path, 0o600) - except Exception as e: - logger.info(f'An exception has occurred attempting to screenshot: {url} : {e}') - path = None + await asyncio.gather(*tasks) finally: - await page.close() - await context.close() - await browser.close() - logger.info(f'{date} {url} {path}') - return url if path else '' + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + async def take_screenshot(self, url: str, *, output_path: Path | None = None) -> str: + normalized_url = _target_url(url) + path = output_path or self.screenshot_path(normalized_url) + async with self._browser() as browser: + return await self._capture(browser, normalized_url, path) def screenshot_path(self, url: str) -> Path: parsed = urlsplit(_target_url(url)) diff --git a/theHarvester/theHarvester.py b/theHarvester/theHarvester.py index d462cfbd..c88b7d18 100644 --- a/theHarvester/theHarvester.py +++ b/theHarvester/theHarvester.py @@ -31,10 +31,4 @@ def main(): import uvloop uvloop.install() - - if 'linux' in platform: - import aiomultiprocess - - # As we are not using Windows, we can change the spawn method to fork for greater performance - aiomultiprocess.set_context('fork') asyncio.run(_run()) diff --git a/uv.lock b/uv.lock index e530d1f3..4c5f9419 100644 --- a/uv.lock +++ b/uv.lock @@ -157,15 +157,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/7d/4b633d709b8901d59444d2e512b93e72fe62d2b492a040097c3f7ba017bb/aiohttp_socks-0.11.0-py3-none-any.whl", hash = "sha256:9aacce57c931b8fbf8f6d333cf3cafe4c35b971b35430309e167a35a8aab9ec1", size = 10556, upload-time = "2025-12-09T13:35:50.18Z" }, ] -[[package]] -name = "aiomultiprocess" -version = "0.9.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/02/d4/1e69e17dda5df91734b70d03dbbf9f222ddb438e1f3bf4ea8fa135ce46de/aiomultiprocess-0.9.1.tar.gz", hash = "sha256:f0231dbe0291e15325d7896ebeae0002d95a4f2675426ca05eb35f24c60e495b", size = 24514, upload-time = "2024-04-23T08:26:04.223Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ad/14/c48c2f5c96960f5649a72b96a0a31d45384b37d89a63f7ccea76bf4fceba/aiomultiprocess-0.9.1-py3-none-any.whl", hash = "sha256:3a7b3bb3c38dbfb4d9d1194ece5934b6d32cf0280e8edbe64a7d215bba1322c6", size = 17517, upload-time = "2024-04-23T08:26:01.649Z" }, -] - [[package]] name = "aiosignal" version = "1.4.0" @@ -1571,7 +1562,6 @@ dependencies = [ { name = "aiofiles" }, { name = "aiohttp" }, { name = "aiohttp-socks" }, - { name = "aiomultiprocess" }, { name = "aiosqlite" }, { name = "beautifulsoup4" }, { name = "certifi" }, @@ -1614,7 +1604,6 @@ requires-dist = [ { name = "aiofiles", specifier = "==25.1.0" }, { name = "aiohttp", specifier = "==3.14.1" }, { name = "aiohttp-socks", specifier = "==0.11.0" }, - { name = "aiomultiprocess", specifier = "==0.9.1" }, { name = "aiosqlite", specifier = "==0.22.1" }, { name = "beautifulsoup4", specifier = "==4.15.0" }, { name = "certifi", specifier = "==2026.6.17" },