From df1933b3458d56455f5e3e1c37b0d4dc8d5dfbfc Mon Sep 17 00:00:00 2001 From: NotoriousRebel <36310667+NotoriousRebel@users.noreply.github.com> Date: Sat, 15 Aug 2026 13:20:26 -0400 Subject: [PATCH] fix: preserve provider sessions across pagination --- AGENTS.md | 1 + CHANGELOG.md | 1 + docs/wiki/How-to-add-a-new-module.md | 13 ++ tests/discovery/test_censys.py | 55 ++++++- tests/discovery/test_githubcode_contract.py | 11 +- .../test_provider_session_lifecycle.py | 104 +++++++++++++ tests/lib/test_core.py | 130 +++++++++++++++- theHarvester/discovery/censysearch.py | 141 ++++++++++-------- theHarvester/discovery/githubcode.py | 98 ++++++------ theHarvester/lib/core.py | 107 +++++++------ 10 files changed, 500 insertions(+), 161 deletions(-) create mode 100644 tests/discovery/test_provider_session_lifecycle.py diff --git a/AGENTS.md b/AGENTS.md index e972aca3..a688dc75 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -18,6 +18,7 @@ Read [CONTEXT.md](CONTEXT.md) when changing discovery terminology, evidence clas - **External compatibility:** Flag changes that remove or rename CLI flags, output formats or fields, REST API response fields, or discovery source identifiers without a backward-compatible path and regression coverage. Preserve the existing contract or document and test the migration. - **Sensitive-data boundary:** Flag committed credentials, real target or operator data, reconnaissance results, or unsanitized provider payloads, including in logs, fixtures, and examples. Keep only the diagnostic metadata needed, redact sensitive values, and use RFC-reserved domains and TEST-NET IP ranges. - **Reconnaissance boundary:** Flag routine tests or CI that contact live third-party targets or providers. Use mocks or local fixtures; live reconnaissance belongs only in intentionally configured integration checks against explicitly authorized targets. +- **State-lifetime audit:** When changing network, pagination, retry, proxy, cookie, or cancellation behavior, trace the entire provider conversation using [How to add a new discovery module](docs/wiki/How-to-add-a-new-module.md#own-the-provider-conversation). Account for every owned resource and every shared-helper caller before declaring the change complete. ## Verification diff --git a/CHANGELOG.md b/CHANGELOG.md index a87e144e..e1d3acad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -79,6 +79,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Removed the nonfunctional ThreatCrowd source because its service hostnames terminate at deleted AWS load balancers and return NXDOMAIN; OTX remains available through its separate adapter. ### Fixed +- Reused one connection pool, proxy identity, and cookie jar across Censys and GitHub Code pagination while keeping provider sessions isolated and cancellation-safe. - Sent a stable, versioned theHarvester identity with provider and API requests while preserving explicit browser identities for sources that require them. - Kept API endpoint scan URLs canonical instead of prefixing targets onto already complete URLs. - Made DeHashed pagination honor the CLI limit, retain only normalized email and IP evidence, and discard raw breach rows; aligned LeakIX with its authenticated subdomain endpoint and documented rate-limit retry. diff --git a/docs/wiki/How-to-add-a-new-module.md b/docs/wiki/How-to-add-a-new-module.md index 069742e6..41117613 100644 --- a/docs/wiki/How-to-add-a-new-module.md +++ b/docs/wiki/How-to-add-a-new-module.md @@ -25,6 +25,19 @@ An adapter normally provides: Do not return fields the provider did not supply. Normalize and deduplicate before returning results. +### Own the provider conversation + +A provider conversation is the related request sequence for one source execution: initial request, pagination, retries or polling, and final response handling. Give that sequence one explicit owner. + +- Reuse one `AsyncFetcher.open_session()` for related requests so the connection pool, headers, cookie jar, and chosen proxy identity remain stable. Pass the borrowed session to shared fetch methods with `session=` and let only the outer owner close it. +- Keep the default cookie jar when later provider requests may depend on earlier responses. Use `aiohttp.DummyCookieJar()` for deliberately independent probes, such as takeover candidates, so one target cannot influence another. +- Scope a session to one provider and authorized target. Never share cookies, authentication state, or proxy identity across source executions or unrelated targets. +- Preserve cancellation while closing every owned session, response, task, and connector. Cover both successful completion and interruption in focused tests. +- Treat session construction and teardown as adapter lifecycle stages. Preserve the existing TLS and timeout policy unless the source contract explicitly changes, classify ordinary lifecycle failures through the adapter status fields, and let native cancellation propagate. +- Before extending a shared fetcher interface, audit positional callers and every owned-versus-borrowed branch. New optional parameters must not reinterpret existing calls. + +The completion check is an offline test in which a later page depends on state established by an earlier page, plus a cleanup assertion proving the provider session closes. + ## 3. Register the source Add one catalog entry in [`theHarvester/lib/source_catalog.py`](https://github.com/laramies/theHarvester/blob/dev/theHarvester/lib/source_catalog.py) and one factory entry in [`theHarvester/lib/source_runner.py`](https://github.com/laramies/theHarvester/blob/dev/theHarvester/lib/source_runner.py). The catalog supplies CLI help, source selection, and activity classification; the factory constructs the adapter; the runner collects declared result routes and persists them through the existing completed-result flow. diff --git a/tests/discovery/test_censys.py b/tests/discovery/test_censys.py index a2471052..50e17bec 100644 --- a/tests/discovery/test_censys.py +++ b/tests/discovery/test_censys.py @@ -1,7 +1,10 @@ import asyncio +import contextlib import sys import types +from collections.abc import AsyncIterator from pathlib import Path +from typing import Any import pytest @@ -41,9 +44,11 @@ def test_legacy_search_api_credentials_fail_closed(monkeypatch) -> None: @pytest.mark.asyncio -async def test_search_calls_platform_api_directly_and_follows_page_tokens(monkeypatch) -> None: +async def test_search_calls_platform_api_directly_and_follows_page_tokens(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(censysearch.Core, 'censys_key', lambda: ('platform-token', 'org-id')) calls: list[dict[str, object]] = [] + session_options: list[dict[str, object]] = [] + session = object() responses = [ FetcherResponse( { @@ -86,11 +91,17 @@ async def test_search_calls_platform_api_directly_and_follows_page_tokens(monkey ), ] - async def fake_post_fetch(url: str, **kwargs): + async def fake_post_fetch(url: str, **kwargs: Any) -> FetcherResponse: calls.append({'url': url, **kwargs}) return responses.pop(0) + @contextlib.asynccontextmanager + async def fake_open_session(**kwargs: object) -> AsyncIterator[object]: + session_options.append(kwargs) + yield session + monkeypatch.setattr(censysearch.AsyncFetcher, 'post_fetch', fake_post_fetch) + monkeypatch.setattr(censysearch.AsyncFetcher, 'open_session', fake_open_session) search = censysearch.SearchCensys('example.com', limit=250) await search.process(proxy=True) @@ -116,11 +127,51 @@ async def test_search_calls_platform_api_directly_and_follows_page_tokens(monkey assert calls[0]['json'] is True assert calls[0]['proxy'] is True assert calls[0]['include_metadata'] is True + assert session_options == [ + { + 'headers': {'Accept': 'application/json', 'Authorization': 'Bearer platform-token'}, + 'proxy': True, + 'request_timeout': 720, + } + ] + assert all(call['session'] is session for call in calls) assert await search.get_hostnames() == {'a.example.com', 'b.example.com'} assert await search.get_emails() == {'admin@example.com', 'ops@example.com'} assert search.execution_status == 'completed' +@pytest.mark.asyncio +async def test_session_setup_failure_reports_transport_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(censysearch.Core, 'censys_key', lambda: ('platform-token', 'org-id')) + + @contextlib.asynccontextmanager + async def failed_open_session(**_kwargs: object) -> AsyncIterator[object]: + raise OSError('sensitive provider detail') + yield object() + + monkeypatch.setattr(censysearch.AsyncFetcher, 'open_session', failed_open_session) + search = censysearch.SearchCensys('example.com') + + await search.process() + + assert search.execution_status == 'failed' + assert search.stop_reason == 'transport-error' + + +@pytest.mark.asyncio +async def test_unexpected_adapter_failure_is_not_misclassified(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(censysearch.Core, 'censys_key', lambda: ('platform-token', 'org-id')) + search = censysearch.SearchCensys('example.com') + + async def failed_search() -> None: + raise RuntimeError('adapter defect') + + monkeypatch.setattr(search, 'do_search', failed_search) + + with pytest.raises(RuntimeError, match='adapter defect'): + await search.process() + + @pytest.mark.asyncio async def test_search_uses_free_wallet_and_respects_result_limit(monkeypatch) -> None: monkeypatch.setattr(censysearch.Core, 'censys_key', lambda: ('platform-token', None)) diff --git a/tests/discovery/test_githubcode_contract.py b/tests/discovery/test_githubcode_contract.py index b2367fc8..1e2e55e4 100644 --- a/tests/discovery/test_githubcode_contract.py +++ b/tests/discovery/test_githubcode_contract.py @@ -1,3 +1,5 @@ +import contextlib +from collections.abc import AsyncIterator from typing import Any import pytest @@ -42,11 +44,16 @@ def install_github_responses(monkeypatch: pytest.MonkeyPatch): async def __aexit__(self, *_args: Any) -> None: return None - def get(self, url: str, *, proxy: str | None) -> FakeResponse: + def get(self, url: str, *, proxy: str | None = None) -> FakeResponse: requested_urls.append(url) return next(response_iterator) - monkeypatch.setattr(githubcode.aiohttp, 'ClientSession', FakeSession) + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: object) -> AsyncIterator[FakeSession]: + async with FakeSession(headers={}) as session: + yield session + + monkeypatch.setattr(githubcode.AsyncFetcher, 'open_session', fake_open_session) return requested_urls monkeypatch.setattr(githubcode.Core, 'github_key', staticmethod(lambda: 'test-token')) diff --git a/tests/discovery/test_provider_session_lifecycle.py b/tests/discovery/test_provider_session_lifecycle.py new file mode 100644 index 00000000..e1d6a11d --- /dev/null +++ b/tests/discovery/test_provider_session_lifecycle.py @@ -0,0 +1,104 @@ +import pytest +from aiohttp import web + +from theHarvester.discovery import censysearch, githubcode + + +@pytest.mark.asyncio +async def test_censys_pagination_preserves_provider_cookies( + monkeypatch: pytest.MonkeyPatch, + unused_tcp_port: int, +) -> None: + requests: list[tuple[str | None, str | None]] = [] + + async def search(request: web.Request) -> web.Response: + body = await request.json() + page_token = body.get('page_token') + requests.append((page_token, request.cookies.get('provider-session'))) + if page_token is None: + response = web.json_response( + { + 'result': { + 'hits': [{'certificate_v1': {'resource': {'names': ['one.example.com']}}}], + 'next_page_token': 'page-two', + } + } + ) + response.set_cookie('provider-session', 'ready') + return response + if request.cookies.get('provider-session') != 'ready': + return web.json_response({'error': 'missing provider session'}, status=403) + return web.json_response( + { + 'result': { + 'hits': [{'certificate_v1': {'resource': {'names': ['two.example.com']}}}], + 'next_page_token': '', + } + } + ) + + app = web.Application() + app.router.add_post('/search', search) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, '127.0.0.1', unused_tcp_port) + await site.start() + monkeypatch.setattr(censysearch.Core, 'censys_key', lambda: ('platform-token', None)) + monkeypatch.setattr(censysearch.SearchCensys, 'SERVER', f'http://localhost:{unused_tcp_port}/search') + + try: + source = censysearch.SearchCensys('example.com', limit=2) + await source.process() + finally: + await runner.cleanup() + + assert requests == [(None, None), ('page-two', 'ready')] + assert await source.get_hostnames() == {'one.example.com', 'two.example.com'} + assert source.execution_status == 'completed' + + +@pytest.mark.asyncio +async def test_github_code_pagination_preserves_provider_cookies( + monkeypatch: pytest.MonkeyPatch, + unused_tcp_port: int, +) -> None: + requests: list[tuple[str | None, str | None]] = [] + + async def search(request: web.Request) -> web.Response: + page = request.query.get('page') + requests.append((page, request.cookies.get('provider-session'))) + if page == '1': + response = web.json_response( + {'items': [{'text_matches': [{'fragment': 'first@example.com'}]}]}, + headers={'Link': f'; rel="next"'}, + ) + response.set_cookie('provider-session', 'ready') + return response + if request.cookies.get('provider-session') != 'ready': + return web.json_response({'error': 'missing provider session'}, status=403) + return web.json_response({'items': [{'text_matches': [{'fragment': 'second@example.com'}]}]}) + + app = web.Application() + app.router.add_get('/search/code', search) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, '127.0.0.1', unused_tcp_port) + await site.start() + monkeypatch.setattr(githubcode.Core, 'github_key', lambda: 'github-token') + monkeypatch.setattr(githubcode, 'get_delay', lambda: 0) + + async def no_sleep(_seconds: float) -> None: + return None + + monkeypatch.setattr(githubcode.asyncio, 'sleep', no_sleep) + + try: + source = githubcode.SearchGithubCode('example.com', limit=2) + source.base_url = f'http://localhost:{unused_tcp_port}/search/code?q=example.com' + await source.process() + finally: + await runner.cleanup() + + assert requests == [('1', None), ('2', 'ready')] + assert source.counter == 2 + assert await source.get_emails() == {'first@example.com', 'second@example.com'} diff --git a/tests/lib/test_core.py b/tests/lib/test_core.py index 3bc749b0..e9e81745 100644 --- a/tests/lib/test_core.py +++ b/tests/lib/test_core.py @@ -228,10 +228,20 @@ class DummyResponse: class DummySession: instances: list[DummySession] = [] - def __init__(self, *, headers=None, timeout=None, connector=None): + def __init__( + self, + *, + headers: dict[str, str] | None = None, + timeout: Any = None, + connector: Any = None, + proxy: str | None = None, + cookie_jar: Any = None, + ) -> None: self.headers = headers self.timeout = timeout self.connector = connector + self.proxy = proxy + self.cookie_jar = cookie_jar self.closed = False self.requests: list[tuple[str, str, dict[str, Any]]] = [] DummySession.instances.append(self) @@ -406,6 +416,99 @@ async def test_fetch_reused_session_uses_a_stable_explicit_ssl_policy(monkeypatc assert ssl_policies[0] is ssl_policies[1] +@pytest.mark.asyncio +async def test_open_session_owns_one_proxy_and_cookie_policy(monkeypatch: pytest.MonkeyPatch) -> None: + reset_dummy_sessions() + cookie_jar = core_module.aiohttp.DummyCookieJar() + monkeypatch.setattr(core_module.aiohttp, 'ClientSession', DummySession) + monkeypatch.setattr(AsyncFetcher, '_ssl_context', staticmethod(lambda: 'ssl-context')) + + async def fake_connector(*_args: object) -> str: + return 'provider-connector' + + monkeypatch.setattr(AsyncFetcher, '_create_connector', fake_connector) + + async with AsyncFetcher.open_session( + headers={'Accept': 'application/json'}, + proxy='http://proxy.example:8080', + request_timeout=45, + cookie_jar=cookie_jar, + ) as session: + assert session.closed is False + + assert session.closed is True + assert session.headers['Accept'] == 'application/json' + assert session.proxy == 'http://proxy.example:8080' + assert session.connector == 'provider-connector' + assert session.cookie_jar is cookie_jar + assert session.timeout.total == 45 + + +@pytest.mark.asyncio +async def test_open_session_preserves_project_ca_and_aiohttp_default_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + reset_dummy_sessions() + connector_options: list[tuple[str | None, str | None, object]] = [] + monkeypatch.setattr(core_module.aiohttp, 'ClientSession', DummySession) + monkeypatch.setattr(AsyncFetcher, '_ssl_context', staticmethod(lambda: 'ssl-context')) + + async def fake_connector( + proxy_url: str | None, + proxy_type: str | None, + ssl_context: object, + ) -> str: + connector_options.append((proxy_url, proxy_type, ssl_context)) + return 'direct-provider-connector' + + monkeypatch.setattr(AsyncFetcher, '_create_connector', fake_connector) + + async with AsyncFetcher.open_session() as session: + assert session.timeout is None + + assert connector_options == [(None, None, 'ssl-context')] + assert session.connector == 'direct-provider-connector' + + +@pytest.mark.asyncio +async def test_open_session_finishes_close_and_preserves_the_first_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + close_started = asyncio.Event() + release_close = asyncio.Event() + cancelled = asyncio.CancelledError('provider-cancelled') + + class BlockingCloseSession: + closed = False + + async def close(self) -> None: + close_started.set() + await release_close.wait() + self.closed = True + + session = BlockingCloseSession() + + async def fake_build_session(*_args: object, **_kwargs: object) -> BlockingCloseSession: + return session + + monkeypatch.setattr(AsyncFetcher, '_build_session', fake_build_session) + + async def use_session() -> None: + async with AsyncFetcher.open_session(): + raise cancelled + + task = asyncio.create_task(use_session()) + await close_started.wait() + task.cancel('second-cancellation') + release_close.set() + + with pytest.raises(asyncio.CancelledError) as raised: + await task + + assert raised.value is cancelled + assert session.closed is True + + def test_default_headers_add_project_identity_without_mutating_caller_headers(monkeypatch) -> None: monkeypatch.setattr(Core, 'get_user_agent', staticmethod(lambda: 'test-agent')) supplied = {'Accept': 'application/json'} @@ -1057,6 +1160,15 @@ async def test_post_fetch_decodes_string_payload_and_posts_params(monkeypatch) - monkeypatch.setattr(core_module.certifi, 'where', lambda: '/tmp/cacert.pem') monkeypatch.setattr(Core, 'get_user_agent', staticmethod(lambda: 'test-agent')) + async def fake_create_connector( + _proxy_url: str | None, + _proxy_type: str | None, + ssl_context: object, + ) -> str: + return f'connector:{ssl_context}' + + monkeypatch.setattr(AsyncFetcher, '_create_connector', fake_create_connector) + result = await AsyncFetcher.post_fetch( 'https://example.com/api', data='{"query": "example"}', @@ -1067,8 +1179,9 @@ async def test_post_fetch_decodes_string_payload_and_posts_params(monkeypatch) - assert result == {'ok': True} session = DummySession.instances[0] assert session.headers == {'User-Agent': 'test-agent'} + assert session.connector == 'connector:ssl-context' assert session.requests == [ - ('POST', 'https://example.com/api', {'data': {'query': 'example'}, 'ssl': 'ssl-context', 'params': {'page': 2}}) + ('POST', 'https://example.com/api', {'data': {'query': 'example'}, 'params': {'page': 2}}) ] @@ -1080,6 +1193,15 @@ async def test_post_fetch_sends_json_body(monkeypatch) -> None: monkeypatch.setattr(core_module.ssl, 'create_default_context', lambda cafile=None: 'ssl-context') monkeypatch.setattr(core_module.certifi, 'where', lambda: '/tmp/cacert.pem') + async def fake_create_connector( + _proxy_url: str | None, + _proxy_type: str | None, + ssl_context: object, + ) -> str: + return f'connector:{ssl_context}' + + monkeypatch.setattr(AsyncFetcher, '_create_connector', fake_create_connector) + result = await AsyncFetcher.post_fetch( 'https://example.com/api', json_body={'scan': 'example'}, @@ -1088,6 +1210,7 @@ async def test_post_fetch_sends_json_body(monkeypatch) -> None: assert result == {'ok': True} session = DummySession.instances[0] + assert session.connector == 'connector:ssl-context' assert session.requests == [ ('POST', 'https://example.com/api', {'json': {'scan': 'example'}}) ] @@ -1138,10 +1261,11 @@ async def test_post_fetch_proxy_branch_posts_body_and_params_with_http_proxy(mon assert created_connectors == [('http://proxy.local:8080', 'http', 'ssl-context')] session = DummySession.instances[0] assert session.connector == 'connector' + assert session.proxy == 'http://proxy.local:8080' assert session.requests == [ ( 'POST', 'https://example.com/resource', - {'json': {'scan': 'example'}, 'params': {'page': 2}, 'proxy': 'http://proxy.local:8080'}, + {'json': {'scan': 'example'}, 'params': {'page': 2}}, ) ] diff --git a/theHarvester/discovery/censysearch.py b/theHarvester/discovery/censysearch.py index 1fb11e16..d2279a76 100644 --- a/theHarvester/discovery/censysearch.py +++ b/theHarvester/discovery/censysearch.py @@ -1,5 +1,9 @@ +import asyncio +import ssl from typing import Any +import aiohttp + from theHarvester.discovery.constants import MissingKey from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse @@ -75,72 +79,74 @@ class SearchCensys: records_seen = 0 malformed = False - while records_seen < self.limit: - body = { - 'query': f'cert.names: "{self.word}"', - 'fields': ['cert.names', 'cert.parsed.subject.email_address'], - 'page_size': min(self.MAX_RESULTS_PER_PAGE, self.limit - records_seen), - } - if page_token is not None: - body['page_token'] = page_token - try: - response = await AsyncFetcher.post_fetch( - self.SERVER, - headers=headers, - params=params, - json=True, - proxy=self.proxy, - include_metadata=True, - json_body=body, - ) - except Exception: - self._stop('failed', 'transport-error') - return - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') - return - if not isinstance(response.body, dict) or not isinstance(response.body.get('result'), dict): - self._stop('failed', 'invalid-response') - return - result = response.body['result'] - hits = result.get('hits') - next_page_token = result.get('next_page_token') - if not isinstance(hits, list) or (next_page_token is not None and not isinstance(next_page_token, str)): - self._stop('failed', 'invalid-response') - return + async with AsyncFetcher.open_session(headers=headers, proxy=self.proxy, request_timeout=720) as session: + while records_seen < self.limit: + body = { + 'query': f'cert.names: "{self.word}"', + 'fields': ['cert.names', 'cert.parsed.subject.email_address'], + 'page_size': min(self.MAX_RESULTS_PER_PAGE, self.limit - records_seen), + } + if page_token is not None: + body['page_token'] = page_token + try: + response = await AsyncFetcher.post_fetch( + self.SERVER, + session=session, + headers=headers, + params=params, + json=True, + proxy=self.proxy, + include_metadata=True, + json_body=body, + ) + except Exception: + self._stop('failed', 'transport-error') + return + if not isinstance(response, FetcherResponse): + self._stop('failed', 'transport-error') + return + if response.status == 429: + self._stop('rate-limited', 'http-429') + return + if response.status in {401, 403}: + self._stop('failed', 'access-denied') + return + if not 200 <= response.status < 300: + self._stop('failed', f'http-{response.status}') + return + if not isinstance(response.body, dict) or not isinstance(response.body.get('result'), dict): + self._stop('failed', 'invalid-response') + return + result = response.body['result'] + hits = result.get('hits') + next_page_token = result.get('next_page_token') + if not isinstance(hits, list) or (next_page_token is not None and not isinstance(next_page_token, str)): + self._stop('failed', 'invalid-response') + return - for hit in hits: + for hit in hits: + if records_seen >= self.limit: + break + malformed = self._parse_hit(hit) or malformed + records_seen += 1 if records_seen >= self.limit: - break - malformed = self._parse_hit(hit) or malformed - records_seen += 1 - if records_seen >= self.limit: - if malformed: - self._stop('failed', 'invalid-response') - else: - self.execution_status = 'completed' - return - if not next_page_token: - if malformed: - self._stop('failed', 'invalid-response') - else: - self.execution_status = 'completed' - self.stop_reason = None if self._has_results() else 'no-results' - return - if next_page_token in seen_tokens: - self._stop('failed', 'repeated-cursor') - return - seen_tokens.add(next_page_token) - page_token = next_page_token + if malformed: + self._stop('failed', 'invalid-response') + else: + self.execution_status = 'completed' + return + if not next_page_token: + if malformed: + self._stop('failed', 'invalid-response') + else: + self.execution_status = 'completed' + self.stop_reason = None if self._has_results() else 'no-results' + return + if next_page_token in seen_tokens: + self._stop('failed', 'repeated-cursor') + return + seen_tokens.add(next_page_token) + page_token = next_page_token async def get_hostnames(self) -> set[str]: return self.totalhosts @@ -150,4 +156,9 @@ class SearchCensys: async def process(self, proxy: bool = False) -> None: self.proxy = proxy - await self.do_search() + try: + await self.do_search() + except asyncio.CancelledError: + raise + except (aiohttp.ClientError, TimeoutError, OSError, ssl.SSLError, ValueError): + self._stop('failed', 'transport-error') diff --git a/theHarvester/discovery/githubcode.py b/theHarvester/discovery/githubcode.py index 551a7ed7..5e9c941f 100644 --- a/theHarvester/discovery/githubcode.py +++ b/theHarvester/discovery/githubcode.py @@ -1,13 +1,12 @@ import asyncio import logging -import random import urllib.parse as urlparse from typing import Any, NamedTuple import aiohttp from theHarvester.discovery.constants import MissingKey, get_delay -from theHarvester.lib.core import Core +from theHarvester.lib.core import AsyncFetcher, Core from theHarvester.parsers import myparser logger = logging.getLogger(__name__) @@ -111,12 +110,18 @@ class SearchGithubCode: else: return result.last_page - async def do_search(self, page: int) -> tuple[str, dict, int, Any]: + async def do_search( + self, + page: int, + session: aiohttp.ClientSession | None = None, + ) -> tuple[str, dict, int, Any]: try: + if session is None: + async with AsyncFetcher.open_session(headers=self.headers, proxy=self.proxy) as owned_session: + return await self.do_search(page, owned_session) url = f'{self.base_url}&page={page}' if page else self.base_url - async with aiohttp.ClientSession(headers=self.headers) as sess: - async with sess.get(url, proxy=random.choice(Core.proxy_list()) if self.proxy else None) as resp: - return await resp.text(), await resp.json(), resp.status, resp.links + async with session.get(url) as resp: + return await resp.text(), await resp.json(), resp.status, resp.links except Exception as e: logger.info(f'Error performing search: {e}') return '', {}, 500, {} @@ -124,50 +129,51 @@ class SearchGithubCode: async def process(self, proxy: bool = False) -> None: try: self.proxy = proxy - while self.counter < self.limit and self.page != 0: - try: - api_response = await self.do_search(self.page) - result = await self.handle_response(api_response) + async with AsyncFetcher.open_session(headers=self.headers, proxy=self.proxy) as session: + while self.counter < self.limit and self.page != 0: + try: + api_response = await self.do_search(self.page, session) + result = await self.handle_response(api_response) - if isinstance(result, SuccessResult): - # Reset retry counter on any successful response - self.retry_count = 0 - logger.info(f'\tSearching {self.counter} results.') - remaining = self.limit - self.counter - fragments = result.fragments[:remaining] - if not fragments: + if isinstance(result, SuccessResult): + # Reset retry counter on any successful response + self.retry_count = 0 + logger.info(f'\tSearching {self.counter} results.') + remaining = self.limit - self.counter + fragments = result.fragments[:remaining] + if not fragments: + self.page = 0 + break + self.total_results += f'{" ".join(fragments)} ' + self.counter += len(fragments) + if self.counter >= self.limit: + self.page = 0 + break + next_or_last = result.next_page or result.last_page + # Break if pagination does not advance to avoid infinite loop + if next_or_last == self.page: + logger.info('\tNo page advancement detected; exiting to avoid infinite loop.') + self.page = 0 + break + self.page = next_or_last + await asyncio.sleep(get_delay()) + elif isinstance(result, RetryResult): + self.retry_count += 1 + if self.retry_count > self.max_retries: + logger.info('\tMaximum retries reached; exiting to avoid infinite loop.') + self.page = 0 + break + sleepy_time = get_delay() + result.time + logger.info(f'\tRetrying page in {sleepy_time} seconds...') + await asyncio.sleep(sleepy_time) + else: + # On error, stop to avoid endless retries on a bad state + logger.info(f'\tGitHub code API request failed with status {result.status_code}') self.page = 0 break - self.total_results += f'{" ".join(fragments)} ' - self.counter += len(fragments) - if self.counter >= self.limit: - self.page = 0 - break - next_or_last = result.next_page or result.last_page - # Break if pagination does not advance to avoid infinite loop - if next_or_last == self.page: - logger.info('\tNo page advancement detected; exiting to avoid infinite loop.') - self.page = 0 - break - self.page = next_or_last + except Exception as e: + logger.info(f'Error processing page: {e}') await asyncio.sleep(get_delay()) - elif isinstance(result, RetryResult): - self.retry_count += 1 - if self.retry_count > self.max_retries: - logger.info('\tMaximum retries reached; exiting to avoid infinite loop.') - self.page = 0 - break - sleepy_time = get_delay() + result.time - logger.info(f'\tRetrying page in {sleepy_time} seconds...') - await asyncio.sleep(sleepy_time) - else: - # On error, stop to avoid endless retries on a bad state - logger.info(f'\tGitHub code API request failed with status {result.status_code}') - self.page = 0 - break - except Exception as e: - logger.info(f'Error processing page: {e}') - await asyncio.sleep(get_delay()) except Exception as e: logger.info(f'An exception has occurred in githubcode process: {e}') diff --git a/theHarvester/lib/core.py b/theHarvester/lib/core.py index dadf6e5a..28fec96d 100644 --- a/theHarvester/lib/core.py +++ b/theHarvester/lib/core.py @@ -20,6 +20,7 @@ import yaml from aiohttp_socks import ProxyConnector from theHarvester import __version__ +from theHarvester.lib.cancellation import drain_tasks_after_cancellation from theHarvester.lib.output import output_logger from theHarvester.lib.source_catalog import SOURCE_SPECS, resolve_sources @@ -498,10 +499,47 @@ class AsyncFetcher: 'timeout': client_timeout, 'connector': connector, } + if proxy_url is not None and proxy_type == 'http': + session_kwargs['proxy'] = proxy_url if cookie_jar is not None: session_kwargs['cookie_jar'] = cookie_jar return aiohttp.ClientSession(**session_kwargs) + @classmethod + @contextlib.asynccontextmanager + async def open_session( + cls, + *, + headers: dict[str, str] | None = None, + proxy: str | bool | None = '', + request_timeout: int | None = None, + cookie_jar: aiohttp.abc.AbstractCookieJar | None = None, + ) -> AsyncIterator[aiohttp.ClientSession]: + """Own one connection pool, proxy identity, and cookie jar for a provider conversation.""" + proxy_url, proxy_type = cls._resolve_proxy(proxy) + session = await cls._build_session( + cls._default_headers(headers), + cls._request_timeout(request_timeout), + proxy_url, + proxy_type, + cls._ssl_context(), + cookie_jar, + ) + body_error: BaseException | None = None + try: + yield session + except BaseException as error: + body_error = error + close_task = asyncio.create_task(session.close(), name='provider-http-session-close') + interruptions = await drain_tasks_after_cancellation((close_task,), cancel=False) + close_error = None if close_task.cancelled() else close_task.exception() + if body_error is not None: + raise body_error + if interruptions: + raise interruptions[0] + if close_error is not None: + raise close_error + @staticmethod async def _read_response( response: aiohttp.ClientResponse, @@ -627,63 +665,46 @@ class AsyncFetcher: @classmethod async def post_fetch( cls, - url, - headers=None, + url: str, + headers: dict[str, str] | None = None, data: str | dict[str, Any] = '', params: Sized = '', json: bool = False, - proxy: bool = False, + proxy: str | bool | None = False, include_metadata: bool = False, json_body: dict[str, Any] | None = None, - ): + *, + session: aiohttp.ClientSession | None = None, + ) -> Any: headers = cls._default_headers(headers) - timeout = cls._request_timeout(720) # By default, timeout is 5 minutes, changed to 12-minutes # results are well worth the wait try: - if proxy: - proxy_url, proxy_type = cls._resolve_proxy(proxy) - sslcontext = cls._ssl_context() - request_kwargs: dict[str, Any] = { - 'data': cls._normalize_data(data) if json_body is None else None, - 'proxy': proxy_url if proxy_type == 'http' else None, - } - if params != '': - request_kwargs['params'] = params - async with await cls._build_session(headers, timeout, proxy_url, proxy_type, sslcontext) as session: - return await cls._request( - session, - 'POST', + if session is None: + async with cls.open_session(headers=headers, proxy=proxy, request_timeout=720) as owned_session: + return await cls.post_fetch( url, - json=json, - json_body=json_body, - include_metadata=include_metadata, - **request_kwargs, - ) - elif params == '': - async with await cls._build_session(headers, timeout) as session: - return await cls._request( - session, - 'POST', - url, - data=cls._normalize_data(data) if json_body is None else None, - json=json, - json_body=json_body, - include_metadata=include_metadata, - ) - else: - async with await cls._build_session(headers, timeout) as session: - return await cls._request( - session, - 'POST', - url, - data=cls._normalize_data(data) if json_body is None else None, - ssl=cls._ssl_context(), + session=owned_session, + data=data, params=params, json=json, - json_body=json_body, include_metadata=include_metadata, + json_body=json_body, ) + request_kwargs: dict[str, Any] = { + 'data': cls._normalize_data(data) if json_body is None else None, + } + if params != '': + request_kwargs['params'] = params + return await cls._request( + session, + 'POST', + url, + json=json, + json_body=json_body, + include_metadata=include_metadata, + **request_kwargs, + ) except (aiohttp.ClientError, TimeoutError, OSError, ssl.SSLError, UnicodeDecodeError, ValueError): return None if include_metadata else ''