diff --git a/tests/discovery/test_bevigil.py b/tests/discovery/test_bevigil.py index daa803e0..02e11ead 100644 --- a/tests/discovery/test_bevigil.py +++ b/tests/discovery/test_bevigil.py @@ -1,3 +1,6 @@ +import asyncio +import contextlib +from collections.abc import AsyncIterator from typing import Any import pytest @@ -11,6 +14,9 @@ from theHarvester.lib.core import FetcherResponse @pytest.mark.asyncio async def test_process_collects_scoped_hostnames_and_urls(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(bevigil.Core, 'bevigil_key', staticmethod(lambda: 'test-key')) + session = object() + session_exited = False + open_calls: list[dict[str, Any]] = [] calls: list[tuple[list[str], dict[str, Any]]] = [] responses = [ FetcherResponse( @@ -25,10 +31,20 @@ async def test_process_collects_scoped_hostnames_and_urls(monkeypatch: pytest.Mo ), ] + @contextlib.asynccontextmanager + async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + open_calls.append(kwargs) + try: + yield session + finally: + session_exited = True + async def fake_fetch_all(urls: list[str], **kwargs: Any) -> list[FetcherResponse]: calls.append((urls, kwargs)) return [responses.pop(0)] + monkeypatch.setattr(bevigil.AsyncFetcher, 'open_session', fake_open_session) monkeypatch.setattr(bevigil.AsyncFetcher, 'fetch_all', fake_fetch_all) search = bevigil.SearchBeVigil('example.com') @@ -44,6 +60,15 @@ async def test_process_collects_scoped_hostnames_and_urls(monkeypatch: pytest.Mo assert all(kwargs['proxy'] is True for _urls, kwargs in calls) assert all(kwargs['json'] is True for _urls, kwargs in calls) assert all(kwargs['include_metadata'] is True for _urls, kwargs in calls) + assert all(kwargs['session'] is session for _urls, kwargs in calls) + assert open_calls == [ + { + 'headers': {'X-Access-Token': 'test-key'}, + 'proxy': True, + 'request_timeout': 60, + } + ] + assert session_exited is True assert search.execution_status == 'completed' assert search.stop_reason is None @@ -112,6 +137,30 @@ async def test_later_malformed_response_preserves_partial_results(monkeypatch: p assert search.stop_reason == 'invalid-response' +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(bevigil.Core, 'bevigil_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_fetch_all(*_args: Any, **_kwargs: Any) -> list[FetcherResponse]: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(bevigil.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(bevigil.AsyncFetcher, 'fetch_all', fake_fetch_all) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await bevigil.SearchBeVigil('example.com').process() + assert session_exited is True + + @pytest.mark.asyncio async def test_early_malformed_rows_and_later_valid_evidence_are_partial(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(bevigil.Core, 'bevigil_key', staticmethod(lambda: 'test-key')) diff --git a/tests/discovery/test_fofa.py b/tests/discovery/test_fofa.py index 7a929fb8..9c399066 100644 --- a/tests/discovery/test_fofa.py +++ b/tests/discovery/test_fofa.py @@ -16,6 +16,7 @@ from theHarvester.lib.core import FetcherResponse async def test_process_uses_cursor_api_to_limit_and_retains_scoped_results(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(fofa.Core, 'fofa_key', lambda: ('test-key', 'operator@example.com')) session = object() + session_exited = False calls: list[dict[str, Any]] = [] responses = [ FetcherResponse( @@ -35,8 +36,12 @@ async def test_process_uses_cursor_api_to_limit_and_retains_scoped_results(monke @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited assert kwargs['proxy'] is True - yield session + try: + yield session + finally: + session_exited = True async def fake_fetch(**kwargs: Any) -> FetcherResponse: calls.append(kwargs) @@ -56,6 +61,7 @@ async def test_process_uses_cursor_api_to_limit_and_retains_scoped_results(monke assert [call['params'].get('next') for call in calls] == [None, 'cursor-2'] assert all(call['url'] == 'https://fofa.info/api/v1/search/next' for call in calls) assert all(call['session'] is session for call in calls) + assert session_exited is True assert base64.b64decode(calls[0]['params']['qbase64']).decode() == 'domain="example.com"' @@ -178,10 +184,15 @@ async def test_malformed_url_does_not_discard_later_valid_rows(monkeypatch: pyte @pytest.mark.asyncio async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(fofa.Core, 'fofa_key', lambda: ('test-key', 'operator@example.com')) + session_exited = False @contextlib.asynccontextmanager async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: - yield object() + nonlocal session_exited + try: + yield object() + finally: + session_exited = True async def fake_fetch(**_kwargs: Any) -> FetcherResponse: raise asyncio.CancelledError @@ -190,3 +201,4 @@ async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(fofa.AsyncFetcher, 'fetch', fake_fetch) with pytest.raises(asyncio.CancelledError): await fofa.SearchFofa('example.com', 10).process() + assert session_exited is True diff --git a/tests/discovery/test_fullhuntsearch.py b/tests/discovery/test_fullhuntsearch.py index c23711b3..d0570b64 100644 --- a/tests/discovery/test_fullhuntsearch.py +++ b/tests/discovery/test_fullhuntsearch.py @@ -1,5 +1,7 @@ import asyncio +import contextlib import logging +from collections.abc import AsyncIterator from typing import Any import pytest @@ -9,6 +11,55 @@ from theHarvester.discovery.constants import MissingKey from theHarvester.lib.core import FetcherResponse +@pytest.mark.asyncio +async def test_process_reuses_one_session_for_fallback_requests(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(fullhuntsearch.Core, 'fullhunt_key', lambda: 'test-key') + session = object() + session_exited = False + open_calls: list[dict[str, Any]] = [] + calls: list[tuple[list[str], dict[str, Any]]] = [] + responses = [ + FetcherResponse({'hosts': []}, 200, {}), + FetcherResponse({'hosts': ['api.example.com']}, 200, {}), + ] + + @contextlib.asynccontextmanager + async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + open_calls.append(kwargs) + try: + yield session + finally: + session_exited = True + + async def fake_fetch_all(urls: list[str], **kwargs: Any) -> list[FetcherResponse]: + calls.append((urls, kwargs)) + return [responses.pop(0)] + + monkeypatch.setattr(fullhuntsearch.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(fullhuntsearch.AsyncFetcher, 'fetch_all', fake_fetch_all) + search = fullhuntsearch.SearchFullHunt('example.com') + + await search.process(proxy=True) + + assert await search.get_hostnames() == ['api.example.com'] + assert [urls for urls, _kwargs in calls] == [ + ['https://fullhunt.io/api/v1/domain/example.com/details'], + ['https://fullhunt.io/api/v1/domain/example.com/subdomains'], + ] + assert all(kwargs['session'] is session for _urls, kwargs in calls) + assert open_calls == [ + { + 'headers': {'User-Agent': fullhuntsearch.Core.get_user_agent(), 'X-API-KEY': 'test-key'}, + 'proxy': True, + 'request_timeout': 60, + } + ] + assert session_exited is True + assert search.execution_status == 'completed' + assert search.stop_reason is None + + @pytest.mark.asyncio async def test_http_failure_is_reported_without_results( monkeypatch: pytest.MonkeyPatch, @@ -236,13 +287,24 @@ async def test_failures_are_structured( @pytest.mark.asyncio async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(fullhuntsearch.Core, 'fullhunt_key', lambda: 'test-key') + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True async def fake_fetch_all(*_args: Any, **_kwargs: Any) -> list[FetcherResponse]: raise asyncio.CancelledError monkeypatch.setattr(fullhuntsearch.AsyncFetcher, 'fetch_all', fake_fetch_all) + monkeypatch.setattr(fullhuntsearch.AsyncFetcher, 'open_session', fake_open_session) with pytest.raises(asyncio.CancelledError): await fullhuntsearch.SearchFullHunt('example.com').process() + assert session_exited is True pytestmark = pytest.mark.provider_contract('fullhunt') diff --git a/tests/discovery/test_hunterhow.py b/tests/discovery/test_hunterhow.py index e9797e8d..8afd84c4 100644 --- a/tests/discovery/test_hunterhow.py +++ b/tests/discovery/test_hunterhow.py @@ -1,3 +1,4 @@ +import asyncio import base64 import contextlib from collections.abc import AsyncIterator @@ -15,6 +16,7 @@ from theHarvester.lib.core import FetcherResponse async def test_process_paginates_to_limit_and_keeps_scoped_hostnames(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(searchhunterhow.Core, 'hunterhow_key', staticmethod(lambda: 'test-key')) session = object() + session_exited = False session_options: list[dict[str, Any]] = [] calls: list[dict[str, Any]] = [] delays: list[float] = [] @@ -39,8 +41,12 @@ async def test_process_paginates_to_limit_and_keeps_scoped_hostnames(monkeypatch @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited session_options.append(kwargs) - yield session + try: + yield session + finally: + session_exited = True async def fake_fetch(**kwargs: Any) -> FetcherResponse: calls.append(kwargs) @@ -68,6 +74,7 @@ async def test_process_paginates_to_limit_and_keeps_scoped_hostnames(monkeypatch assert delays == [2.0] assert search.execution_status == 'completed' assert search.stop_reason is None + assert session_exited is True @pytest.mark.parametrize('key', [None, '', ' ']) @@ -152,3 +159,27 @@ async def test_early_malformed_rows_and_later_valid_evidence_are_partial(monkeyp assert await search.get_hostnames() == {'api.example.com'} assert search.execution_status == 'partial' assert search.stop_reason == 'invalid-response' + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(searchhunterhow.Core, 'hunterhow_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_fetch(**_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(searchhunterhow.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(searchhunterhow.AsyncFetcher, 'fetch', fake_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await searchhunterhow.SearchHunterHow('example.com', limit=10).process() + assert session_exited is True diff --git a/tests/discovery/test_netlas.py b/tests/discovery/test_netlas.py index 378a0376..af9259b7 100644 --- a/tests/discovery/test_netlas.py +++ b/tests/discovery/test_netlas.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib from typing import TYPE_CHECKING, Any @@ -18,6 +19,7 @@ if TYPE_CHECKING: async def test_process_uses_current_api_contract_and_keeps_scoped_results(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(netlas.Core, 'netlas_key', staticmethod(lambda: 'test-key')) session = object() + session_exited = False session_options: list[dict[str, Any]] = [] calls: list[dict[str, Any]] = [] responses = [ @@ -34,8 +36,12 @@ async def test_process_uses_current_api_contract_and_keeps_scoped_results(monkey @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited session_options.append(kwargs) - yield session + try: + yield session + finally: + session_exited = True async def fake_post_fetch(*args: Any, **kwargs: Any) -> FetcherResponse: calls.append({'url': args[0], **kwargs}) @@ -59,6 +65,7 @@ async def test_process_uses_current_api_contract_and_keeps_scoped_results(monkey 'fields': ['domain'], 'source_type': 'include', } + assert session_exited is True @pytest.mark.parametrize('key', [None, '', ' ']) @@ -136,3 +143,27 @@ async def test_malformed_download_rows_preserve_valid_partial_results(monkeypatc assert await search.get_hostnames() == {'ok.example.com'} assert search.execution_status == 'partial' assert search.stop_reason == 'invalid-response' + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(netlas.Core, 'netlas_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(netlas.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(netlas.AsyncFetcher, 'post_fetch', fake_post_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await netlas.SearchNetlas('example.com', limit=10).process() + assert session_exited is True diff --git a/tests/discovery/test_onyphe.py b/tests/discovery/test_onyphe.py index fb3f9b79..1fda9773 100644 --- a/tests/discovery/test_onyphe.py +++ b/tests/discovery/test_onyphe.py @@ -15,6 +15,7 @@ from theHarvester.lib.core import FetcherResponse async def test_process_paginates_to_limit_and_preserves_all_routes(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(onyphe.Core, 'onyphe_key', lambda: 'test-key') session = object() + session_exited = False calls: list[dict[str, Any]] = [] responses = [ FetcherResponse( @@ -51,9 +52,13 @@ async def test_process_paginates_to_limit_and_preserves_all_routes(monkeypatch: @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited assert kwargs['proxy'] is True assert kwargs['headers']['Authorization'] == 'bearer test-key' - yield session + try: + yield session + finally: + session_exited = True async def fake_fetch(**kwargs: Any) -> FetcherResponse: calls.append(kwargs) @@ -68,8 +73,7 @@ async def test_process_paginates_to_limit_and_preserves_all_routes(monkeypatch: assert await search.get_hostnames() == {'api.example.com', 'geo.example.com', 'mail.example.com', 'www.example.com'} assert await search.get_asns() == {'AS64496', 'AS64497'} assert { - (item.asn, item.organization_label, item.subject_kind, item.subject_value) - for item in await search.get_asn_attributions() + (item.asn, item.organization_label, item.subject_kind, item.subject_value) for item in await search.get_asn_attributions() } == { ('AS64496', 'Example Physical Network', 'ip', '192.0.2.10'), ('AS64497', 'Example Logical Network', 'ip', '192.0.2.10'), @@ -79,6 +83,7 @@ async def test_process_paginates_to_limit_and_preserves_all_routes(monkeypatch: assert all(call['session'] is session for call in calls) assert search.execution_status == 'completed' assert search.stop_reason is None + assert session_exited is True @pytest.mark.parametrize('key', [None, '', ' ']) @@ -259,15 +264,24 @@ async def test_malformed_items_preserve_valid_partial_results(monkeypatch: pytes @pytest.mark.asyncio async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(onyphe.Core, 'onyphe_key', lambda: 'test-key') + session_exited = False @contextlib.asynccontextmanager async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: - yield object() - - async def fake_fetch(**_kwargs: Any) -> FetcherResponse: - raise asyncio.CancelledError + nonlocal session_exited + try: + yield object() + finally: + session_exited = True monkeypatch.setattr(onyphe.AsyncFetcher, 'open_session', fake_open_session) + cancellation = asyncio.CancelledError('operator-stop') + + async def fake_fetch(**_kwargs: Any) -> FetcherResponse: + raise cancellation + monkeypatch.setattr(onyphe.AsyncFetcher, 'fetch', fake_fetch) - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as caught: await onyphe.SearchOnyphe('example.com', 10).process() + assert caught.value is cancellation + assert session_exited is True diff --git a/tests/discovery/test_securityscorecard.py b/tests/discovery/test_securityscorecard.py index 8e4156f2..243963e3 100644 --- a/tests/discovery/test_securityscorecard.py +++ b/tests/discovery/test_securityscorecard.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib from collections.abc import AsyncIterator from typing import Any @@ -17,6 +18,7 @@ async def test_process_paginates_documented_domain_and_ip_asset_routes(monkeypat monkeypatch.setattr(securityscorecard.Core, 'securityscorecard_key', staticmethod(lambda: 'test-key')) monkeypatch.setattr(securityscorecard.SearchSecurityScorecard, 'PAGE_SIZE', 2) session = object() + session_exited = False calls: list[dict[str, Any]] = [] post_responses = [ FetcherResponse( @@ -31,8 +33,12 @@ async def test_process_paginates_documented_domain_and_ip_asset_routes(monkeypat @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited assert kwargs == {'headers': search.headers, 'proxy': True} - yield session + try: + yield session + finally: + session_exited = True async def fake_fetch(**kwargs: Any) -> FetcherResponse: calls.append(kwargs) @@ -67,6 +73,7 @@ async def test_process_paginates_documented_domain_and_ip_asset_routes(monkeypat ('https://api.securityscorecard.io/parent-domains/example.com/ips', {'page': 1, 'page_size': 2}), ] assert all(call['session'] is session for call in calls[1:]) + assert session_exited is True @pytest.mark.parametrize('key', [None, '', ' ']) @@ -147,3 +154,27 @@ async def test_malformed_asset_rows_preserve_valid_partial_results(monkeypatch: assert await search.get_ips() == {'192.0.2.1'} assert search.execution_status == 'partial' assert search.stop_reason == 'invalid-response' + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(securityscorecard.Core, 'securityscorecard_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_fetch(**_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(securityscorecard.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(securityscorecard.AsyncFetcher, 'fetch', fake_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await securityscorecard.SearchSecurityScorecard('example.com', 10).process() + assert session_exited is True diff --git a/tests/discovery/test_securitytrails.py b/tests/discovery/test_securitytrails.py index 76a53821..3ca5882b 100644 --- a/tests/discovery/test_securitytrails.py +++ b/tests/discovery/test_securitytrails.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib from typing import TYPE_CHECKING, Any @@ -18,6 +19,7 @@ if TYPE_CHECKING: async def test_process_reuses_session_and_parses_scoped_evidence(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(securitytrailssearch.Core, 'security_trails_key', staticmethod(lambda: 'test-key')) session = object() + session_exited = False calls: list[dict[str, Any]] = [] responses = [ FetcherResponse( @@ -36,8 +38,12 @@ async def test_process_reuses_session_and_parses_scoped_evidence(monkeypatch: py @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited assert kwargs == {'headers': {'APIKEY': 'test-key', 'Accept': 'application/json'}, 'proxy': True} - yield session + try: + yield session + finally: + session_exited = True async def fake_fetch(**kwargs: Any) -> FetcherResponse: calls.append(kwargs) @@ -57,6 +63,7 @@ async def test_process_reuses_session_and_parses_scoped_evidence(monkeypatch: py 'https://api.securitytrails.com/v1/domain/example.com/subdomains', ] assert all(call['session'] is session for call in calls) + assert session_exited is True @pytest.mark.parametrize('key', [None, '', ' ']) @@ -99,3 +106,27 @@ async def test_first_request_failures_are_truthful( assert search.execution_status == status assert search.stop_reason == reason + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(securitytrailssearch.Core, 'security_trails_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_fetch(**_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(securitytrailssearch.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(securitytrailssearch.AsyncFetcher, 'fetch', fake_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await securitytrailssearch.SearchSecuritytrail('example.com').process() + assert session_exited is True diff --git a/tests/discovery/test_sherlockeye.py b/tests/discovery/test_sherlockeye.py index 29ad2db9..6a0b14a6 100644 --- a/tests/discovery/test_sherlockeye.py +++ b/tests/discovery/test_sherlockeye.py @@ -1,7 +1,10 @@ import asyncio +import contextlib import logging import sys import types +from collections.abc import AsyncIterator +from typing import Any import pytest @@ -18,6 +21,16 @@ if 'aiohttp_socks' not in sys.modules: from theHarvester.discovery import sherlockeye from theHarvester.discovery.constants import MissingKey +from theHarvester.lib.core import FetcherResponse + + +@pytest.fixture(autouse=True) +def provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + yield object() + + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'open_session', fake_open_session) @pytest.mark.parametrize('key', [None, '', ' ']) @@ -28,6 +41,58 @@ def test_missing_or_blank_key_raises(monkeypatch: pytest.MonkeyPatch, key: str | sherlockeye.SearchSherlockeye('example.com') +@pytest.mark.asyncio +async def test_process_uses_one_shared_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') + session = object() + exited = False + calls: list[dict[str, Any]] = [] + + @contextlib.asynccontextmanager + async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal exited + assert kwargs == { + 'headers': { + 'User-Agent': sherlockeye.Core.get_user_agent(), + 'Authorization': 'Bearer dummy-key', + 'Content-Type': 'application/json', + }, + 'proxy': True, + 'request_timeout': 90, + } + try: + yield session + finally: + exited = True + + async def fake_post_fetch(*args: Any, **kwargs: Any) -> FetcherResponse: + calls.append({'url': args[0], **kwargs}) + return FetcherResponse({'success': True, 'data': {'results': []}}, 200, {}) + + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', fake_post_fetch) + search = sherlockeye.SearchSherlockeye('example.com') + + await search.process(proxy=True) + + assert calls == [ + { + 'url': search.SYNC_SEARCH_URL, + 'session': session, + 'json': True, + 'include_metadata': True, + 'json_body': { + 'type': 'domain', + 'value': 'example.com', + 'timeoutSeconds': 60, + }, + } + ] + assert exited is True + assert search.execution_status == 'completed' + assert search.stop_reason == 'no-results' + + @pytest.mark.asyncio async def test_process_extracts_domain_intelligence(monkeypatch) -> None: monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') @@ -67,32 +132,10 @@ async def test_process_extracts_domain_intelligence(monkeypatch) -> None: 'balance': {'credits': 10}, } - class _FakeResponse: - status = 200 + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + return FetcherResponse(api_payload, 200, {}) - async def json(self): - return api_payload - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - pass - - class _FakeSession: - def __init__(self, **_kwargs): - pass - - def post(self, *_args, **_kwargs): - return _FakeResponse() - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - pass - - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', _FakeSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', fake_post_fetch) search = sherlockeye.SearchSherlockeye('example.com') await search.process() @@ -108,32 +151,10 @@ async def test_process_extracts_domain_intelligence(monkeypatch) -> None: async def test_process_handles_api_error(monkeypatch, caplog) -> None: monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') - class _FakeResponse: - status = 401 + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + return FetcherResponse({'secret': 'provider-secret-payload'}, 401, {}) - async def text(self): - return 'provider-secret-payload' - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - pass - - class _FakeSession: - def __init__(self, **_kwargs): - pass - - def post(self, *_args, **_kwargs): - return _FakeResponse() - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - pass - - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', _FakeSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', fake_post_fetch) caplog.set_level(logging.INFO, logger=sherlockeye.__name__) search = sherlockeye.SearchSherlockeye('example.com') @@ -152,32 +173,10 @@ async def test_process_handles_api_error(monkeypatch, caplog) -> None: async def test_process_does_not_log_provider_error_message(monkeypatch, caplog) -> None: monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') - class _FakeResponse: - status = 200 + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + return FetcherResponse({'success': False, 'message': 'provider-secret-payload'}, 200, {}) - async def json(self): - return {'success': False, 'message': 'provider-secret-payload'} - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - pass - - class _FakeSession: - def __init__(self, **_kwargs): - pass - - def post(self, *_args, **_kwargs): - return _FakeResponse() - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - pass - - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', _FakeSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', fake_post_fetch) caplog.set_level(logging.INFO, logger=sherlockeye.__name__) search = sherlockeye.SearchSherlockeye('example.com') @@ -202,28 +201,10 @@ async def test_http_failures_are_structured( ) -> None: monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') - class FakeResponse: - async def __aenter__(self): - self.status = status - return self + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + return FetcherResponse({}, status, {}) - async def __aexit__(self, *_args): - return None - - class FakeSession: - def __init__(self, **_kwargs): - pass - - def post(self, *_args, **_kwargs): - return FakeResponse() - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', FakeSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', fake_post_fetch) search = sherlockeye.SearchSherlockeye('example.com') await search.process() @@ -235,32 +216,10 @@ async def test_http_failures_are_structured( async def test_malformed_response_is_structured(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') - class FakeResponse: - status = 200 + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + return FetcherResponse([], 200, {}) - async def json(self): - return [] - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - class FakeSession: - def __init__(self, **_kwargs): - pass - - def post(self, *_args, **_kwargs): - return FakeResponse() - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', FakeSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', fake_post_fetch) search = sherlockeye.SearchSherlockeye('example.com') await search.process() @@ -292,24 +251,35 @@ def test_malformed_link_does_not_discard_later_valid_results(monkeypatch: pytest @pytest.mark.asyncio async def test_transport_failure_and_cancellation_are_distinct(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(sherlockeye.Core, 'sherlockeye_key', lambda: 'dummy-key') + session_exit_count = 0 - class FailedSession: - def __init__(self, **_kwargs): - raise RuntimeError('provider-secret') + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exit_count + try: + yield object() + finally: + session_exit_count += 1 - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', FailedSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'open_session', fake_open_session) + + async def failed_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + raise RuntimeError('provider-secret') + + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', failed_post_fetch) search = sherlockeye.SearchSherlockeye('example.com') await search.process() assert search.execution_status == 'failed' assert search.stop_reason == 'transport-error' + assert session_exit_count == 1 - class CancelledSession: - def __init__(self, **_kwargs): - raise asyncio.CancelledError + async def cancelled_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError - monkeypatch.setattr(sherlockeye.aiohttp, 'ClientSession', CancelledSession) + monkeypatch.setattr(sherlockeye.AsyncFetcher, 'post_fetch', cancelled_post_fetch) with pytest.raises(asyncio.CancelledError): await sherlockeye.SearchSherlockeye('example.com').process() + assert session_exit_count == 2 pytestmark = pytest.mark.provider_contract('sherlockeye') diff --git a/tests/discovery/test_subdomainfinderc99.py b/tests/discovery/test_subdomainfinderc99.py index c94581c6..42205179 100644 --- a/tests/discovery/test_subdomainfinderc99.py +++ b/tests/discovery/test_subdomainfinderc99.py @@ -12,11 +12,16 @@ from theHarvester.lib.core import FetcherResponse @pytest.mark.asyncio async def test_successful_response_returns_only_scoped_hostnames(monkeypatch) -> None: session = object() + session_exited = False calls: list[tuple[str, object]] = [] @contextlib.asynccontextmanager async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: - yield session + nonlocal session_exited + try: + yield session + finally: + session_exited = True async def fake_fetch(*_args, **kwargs): calls.append(('get', kwargs['session'])) @@ -41,6 +46,7 @@ async def test_successful_response_returns_only_scoped_hostnames(monkeypatch) -> assert search.execution_status == 'completed' assert search.stop_reason is None assert calls == [('get', session), ('post', session)] + assert session_exited is True @pytest.mark.asyncio @@ -63,8 +69,6 @@ async def test_empty_initial_response_is_transport_failure(monkeypatch) -> None: assert search.stop_reason == 'transport-error' - - @pytest.mark.asyncio async def test_malformed_scan_response_completes_without_evidence(monkeypatch) -> None: @contextlib.asynccontextmanager @@ -126,9 +130,15 @@ async def test_initial_failures_are_structured( @pytest.mark.asyncio async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: + session_exited = False + @contextlib.asynccontextmanager async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: - yield object() + nonlocal session_exited + try: + yield object() + finally: + session_exited = True async def fake_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: raise asyncio.CancelledError @@ -137,6 +147,7 @@ async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(subdomainfinderc99.AsyncFetcher, 'fetch', fake_fetch) with pytest.raises(asyncio.CancelledError): await subdomainfinderc99.SearchSubdomainfinderc99('example.test').process() + assert session_exited is True pytestmark = pytest.mark.provider_contract('subdomainfinderc99') diff --git a/tests/discovery/test_urlscan.py b/tests/discovery/test_urlscan.py index d2fd3700..74472881 100644 --- a/tests/discovery/test_urlscan.py +++ b/tests/discovery/test_urlscan.py @@ -10,13 +10,21 @@ from theHarvester.discovery import urlscan from theHarvester.lib.core import FetcherResponse +class ProviderSession: + def __init__(self) -> None: + self.exited = False + + @pytest.fixture(autouse=True) -def provider_session(monkeypatch: pytest.MonkeyPatch) -> object: - session = object() +def provider_session(monkeypatch: pytest.MonkeyPatch) -> ProviderSession: + session = ProviderSession() @contextlib.asynccontextmanager async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: - yield session + try: + yield session + finally: + session.exited = True monkeypatch.setattr(urlscan.AsyncFetcher, 'open_session', fake_open_session) return session @@ -25,7 +33,7 @@ def provider_session(monkeypatch: pytest.MonkeyPatch) -> object: @pytest.mark.asyncio async def test_process_collects_sequential_pages_and_preserves_all_routes( monkeypatch: pytest.MonkeyPatch, - provider_session: object, + provider_session: ProviderSession, ) -> None: responses = [ FetcherResponse( @@ -64,7 +72,6 @@ async def test_process_collects_sequential_pages_and_preserves_all_routes( status=200, headers={}, ), - FetcherResponse(body={'results': []}, status=200, headers={}), ] calls: list[dict[str, Any]] = [] @@ -73,7 +80,7 @@ async def test_process_collects_sequential_pages_and_preserves_all_routes( return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 2) await search.process(proxy=True) @@ -99,15 +106,15 @@ async def test_process_collects_sequential_pages_and_preserves_all_routes( ('AS64497', 'Example Transit Two', 'ip', '2001:db8::10'), } assert [call['params'] for call in calls] == [ - {'q': 'domain:example.com'}, - {'q': 'domain:example.com', 'search_after': '200,first'}, - {'q': 'domain:example.com', 'search_after': '100,second'}, + {'q': 'domain:example.com', 'size': 2}, + {'q': 'domain:example.com', 'size': 1, 'search_after': '200,first'}, ] assert all(call['url'] == 'https://urlscan.io/api/v1/search/' for call in calls) assert all(call['session'] is provider_session for call in calls) assert all(call['json'] is True for call in calls) assert all(call['include_metadata'] is True for call in calls) assert all('request_timeout' not in call for call in calls) + assert provider_session.exited is True assert search.execution_status == 'completed' assert search.stop_reason is None @@ -140,7 +147,7 @@ async def test_repeated_asn_relationship_is_retained_once_per_source_run(monkeyp monkeypatch.setattr(urlscan, 'datetime', TickingDateTime) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -153,7 +160,7 @@ async def test_valid_empty_response_is_completed(monkeypatch: pytest.MonkeyPatch return FetcherResponse(body={'results': []}, status=200, headers={}) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -180,7 +187,7 @@ async def test_missing_optional_fields_are_skipped(monkeypatch: pytest.MonkeyPat return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -209,7 +216,7 @@ async def test_malformed_nested_fields_preserve_valid_partial_results(monkeypatc return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -255,7 +262,7 @@ async def test_results_are_typed_and_scoped_before_insertion(monkeypatch: pytest return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -290,7 +297,7 @@ async def test_failed_first_page_is_attributed( return response monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -327,7 +334,7 @@ async def test_later_failure_preserves_partial_results( return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -342,7 +349,7 @@ async def test_fetch_exception_is_transport_failure(monkeypatch: pytest.MonkeyPa raise OSError('private transport details') monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -364,7 +371,7 @@ async def test_missing_cursor_stops_after_first_page(monkeypatch: pytest.MonkeyP ) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -396,7 +403,7 @@ async def test_repeated_cursor_stops_without_a_third_request(monkeypatch: pytest return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10) await search.process() @@ -408,46 +415,88 @@ async def test_repeated_cursor_stops_without_a_third_request(monkeypatch: pytest @pytest.mark.asyncio async def test_pagination_continues_beyond_the_removed_local_page_ceiling(monkeypatch: pytest.MonkeyPatch) -> None: - calls = 0 - - async def fake_fetch(**_kwargs: Any) -> FetcherResponse: - nonlocal calls - calls += 1 - if calls == 1002: - return FetcherResponse(body={'results': []}, status=200, headers={}) - return FetcherResponse( + calls: list[dict[str, Any]] = [] + first_page = [ + {'page': {'domain': f'page-{index}.example.com'}, 'sort': [10_001 - index, f'cursor-{index}']} + for index in range(1, 10_001) + ] + responses = [ + FetcherResponse(body={'results': first_page}, status=200, headers={}), + FetcherResponse( body={ 'results': [ { - 'page': {'domain': f'page-{calls}.example.com'}, - 'sort': [calls, f'cursor-{calls}'], + 'page': {'domain': 'page-10001.example.com'}, + 'sort': [0, 'cursor-10001'], } ] }, status=200, headers={}, - ) + ), + ] + + async def fake_fetch(**kwargs: Any) -> FetcherResponse: + calls.append(kwargs['params']) + return responses.pop(0) monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) - search = urlscan.SearchUrlscan('example.com') + search = urlscan.SearchUrlscan('example.com', 10_001) await search.process() - assert calls == 1002 - assert await search.get_hostnames() == {f'page-{page}.example.com' for page in range(1, 1002)} + assert calls == [ + {'q': 'domain:example.com', 'size': 10_000}, + {'q': 'domain:example.com', 'size': 1, 'search_after': '1,cursor-10000'}, + ] + assert await search.get_hostnames() == {f'page-{page}.example.com' for page in range(1, 10_002)} assert search.execution_status == 'completed' assert search.stop_reason is None @pytest.mark.asyncio -async def test_cancellation_propagates(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_operator_limit_sets_page_size_and_stops_without_an_extra_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[dict[str, Any]] = [] + results = [ + {'page': {'domain': f'result-{index}.example.com'}, 'sort': [10 - index, f'cursor-{index}']} for index in range(10) + ] + + async def fake_fetch(**kwargs: Any) -> FetcherResponse: + calls.append(kwargs['params']) + return FetcherResponse(body={'results': results}, status=200, headers={}) + + monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) + search = urlscan.SearchUrlscan('example.com', 10) + + await search.process() + + assert calls == [{'q': 'domain:example.com', 'size': 10}] + assert len(await search.get_hostnames()) == 10 + assert search.execution_status == 'completed' + assert search.stop_reason is None + + +@pytest.mark.parametrize('limit', [0, -1, True, 1.5]) +def test_limit_must_be_a_positive_integer(limit: Any) -> None: + with pytest.raises(ValueError, match='positive integer'): + urlscan.SearchUrlscan('example.com', limit) + + +@pytest.mark.asyncio +async def test_cancellation_propagates( + monkeypatch: pytest.MonkeyPatch, + provider_session: ProviderSession, +) -> None: async def fake_fetch(**_kwargs: Any) -> FetcherResponse: raise asyncio.CancelledError monkeypatch.setattr(urlscan.AsyncFetcher, 'fetch', fake_fetch) with pytest.raises(asyncio.CancelledError): - await urlscan.SearchUrlscan('example.com').process() + await urlscan.SearchUrlscan('example.com', 10).process() + assert provider_session.exited is True pytestmark = pytest.mark.provider_contract('urlscan') diff --git a/tests/discovery/test_virustotal.py b/tests/discovery/test_virustotal.py index a1c1ab6a..2aa39277 100644 --- a/tests/discovery/test_virustotal.py +++ b/tests/discovery/test_virustotal.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib from typing import TYPE_CHECKING, Any @@ -20,6 +21,7 @@ async def test_process_paginates_without_fixed_sleeps_and_keeps_scoped_evidence( ) -> None: monkeypatch.setattr(virustotal.Core, 'virustotal_key', staticmethod(lambda: 'test-key')) session = object() + session_exited = False calls: list[dict[str, Any]] = [] responses = [ FetcherResponse( @@ -49,6 +51,7 @@ async def test_process_paginates_without_fixed_sleeps_and_keeps_scoped_evidence( @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited assert kwargs == { 'headers': { 'Accept': 'application/json', @@ -56,7 +59,10 @@ async def test_process_paginates_without_fixed_sleeps_and_keeps_scoped_evidence( }, 'proxy': True, } - yield session + try: + yield session + finally: + session_exited = True async def fake_fetch(**kwargs: Any) -> FetcherResponse: calls.append(kwargs) @@ -77,6 +83,7 @@ async def test_process_paginates_without_fixed_sleeps_and_keeps_scoped_evidence( assert search.stop_reason is None assert [call['params'] for call in calls] == [{'limit': 10}, {'limit': 9, 'cursor': 'next-page'}] assert all(call['session'] is session for call in calls) + assert session_exited is True @pytest.mark.asyncio @@ -199,3 +206,27 @@ async def test_repeated_cursor_stops_without_spending_more_quota(monkeypatch: py assert search.execution_status == 'partial' assert search.stop_reason == 'repeated-cursor' assert responses == [] + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(virustotal.Core, 'virustotal_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_fetch(**_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(virustotal.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(virustotal.AsyncFetcher, 'fetch', fake_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await virustotal.SearchVirustotal('example.com', limit=10).process() + assert session_exited is True diff --git a/tests/discovery/test_whoisxml.py b/tests/discovery/test_whoisxml.py index 58b819ad..33d02fc2 100644 --- a/tests/discovery/test_whoisxml.py +++ b/tests/discovery/test_whoisxml.py @@ -1,8 +1,8 @@ from __future__ import annotations +import asyncio import contextlib -from collections.abc import AsyncIterator -from typing import Any +from typing import TYPE_CHECKING, Any import pytest @@ -10,14 +10,25 @@ from theHarvester.discovery import whoisxml from theHarvester.discovery.constants import MissingKey from theHarvester.lib.core import FetcherResponse +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +class ProviderSession: + def __init__(self) -> None: + self.exited = False + @pytest.fixture(autouse=True) -def provider_session(monkeypatch: pytest.MonkeyPatch) -> object: - session = object() +def provider_session(monkeypatch: pytest.MonkeyPatch) -> ProviderSession: + session = ProviderSession() @contextlib.asynccontextmanager async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: - yield session + try: + yield session + finally: + session.exited = True monkeypatch.setattr(whoisxml.AsyncFetcher, 'open_session', fake_open_session) return session @@ -27,7 +38,7 @@ def provider_session(monkeypatch: pytest.MonkeyPatch) -> object: @pytest.mark.asyncio async def test_response_body_is_not_logged_and_scoped_records_are_returned( monkeypatch: pytest.MonkeyPatch, - provider_session: object, + provider_session: ProviderSession, ) -> None: monkeypatch.setattr(whoisxml.Core, 'whoisxml_key', staticmethod(lambda: 'test-key')) calls: list[dict[str, Any]] = [] @@ -78,6 +89,7 @@ async def test_response_body_is_not_logged_and_scoped_records_are_returned( {'apiKey': 'test-key', 'domainName': 'example.com', 'searchAfter': 'www.example.com'}, ] assert all(call['json'] is True and call['include_metadata'] is True for call in calls) + assert provider_session.exited is True @pytest.mark.parametrize('key', [None, '', ' ']) @@ -179,3 +191,20 @@ async def test_later_page_failure_preserves_partial_results(monkeypatch: pytest. assert await search.get_hostnames() == {'api.example.com'} assert search.execution_status == 'partial' assert search.stop_reason == 'http-429' + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session( + monkeypatch: pytest.MonkeyPatch, + provider_session: ProviderSession, +) -> None: + monkeypatch.setattr(whoisxml.Core, 'whoisxml_key', staticmethod(lambda: 'test-key')) + + async def fake_fetch(**_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(whoisxml.AsyncFetcher, 'fetch', fake_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await whoisxml.SearchWhoisXML('example.com', 10).process() + assert provider_session.exited is True diff --git a/tests/discovery/test_zoomeyesearch.py b/tests/discovery/test_zoomeyesearch.py index 431c4b24..c655298e 100644 --- a/tests/discovery/test_zoomeyesearch.py +++ b/tests/discovery/test_zoomeyesearch.py @@ -1,3 +1,4 @@ +import asyncio import base64 import contextlib from collections.abc import AsyncIterator @@ -15,6 +16,7 @@ from theHarvester.lib.core import FetcherResponse async def test_process_reuses_session_and_collects_all_capabilities(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(zoomeyesearch.Core, 'zoomeye_key', staticmethod(lambda: 'test-key')) session = object() + session_exited = False calls: list[dict[str, Any]] = [] responses = [ FetcherResponse( @@ -39,11 +41,15 @@ async def test_process_reuses_session_and_collects_all_capabilities(monkeypatch: @contextlib.asynccontextmanager async def fake_open_session(**kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited assert kwargs == { 'headers': {'API-KEY': 'test-key', 'Content-Type': 'application/json'}, 'proxy': True, } - yield session + try: + yield session + finally: + session_exited = True async def fake_post_fetch(*args: Any, **kwargs: Any) -> FetcherResponse: calls.append({'url': args[0], **kwargs}) @@ -63,6 +69,7 @@ async def test_process_reuses_session_and_collects_all_capabilities(monkeypatch: assert search.stop_reason is None assert [call['url'] for call in calls] == [search.baseurl] assert all(call['session'] is session for call in calls) + assert session_exited is True assert base64.b64decode(calls[0]['json_body']['qbase64']).decode() == 'domain="example.com"' assert calls[0]['json_body']['page'] == 1 assert calls[0]['json_body']['pagesize'] == 2 @@ -219,3 +226,27 @@ async def test_banner_urls_are_absolute_http_and_scoped(monkeypatch: pytest.Monk assert urls == {'https://api.example.com/v1'} assert malformed is False + + +@pytest.mark.asyncio +async def test_cancellation_closes_provider_session(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(zoomeyesearch.Core, 'zoomeye_key', staticmethod(lambda: 'test-key')) + session_exited = False + + @contextlib.asynccontextmanager + async def fake_open_session(**_kwargs: Any) -> AsyncIterator[object]: + nonlocal session_exited + try: + yield object() + finally: + session_exited = True + + async def fake_post_fetch(*_args: Any, **_kwargs: Any) -> FetcherResponse: + raise asyncio.CancelledError('operator-stop') + + monkeypatch.setattr(zoomeyesearch.AsyncFetcher, 'open_session', fake_open_session) + monkeypatch.setattr(zoomeyesearch.AsyncFetcher, 'post_fetch', fake_post_fetch) + + with pytest.raises(asyncio.CancelledError, match='operator-stop'): + await zoomeyesearch.SearchZoomEye('example.com', 10).process() + assert session_exited is True diff --git a/tests/lib/test_core.py b/tests/lib/test_core.py index e9e81745..a7a0e6ba 100644 --- a/tests/lib/test_core.py +++ b/tests/lib/test_core.py @@ -1127,6 +1127,26 @@ async def test_fetch_all_propagates_metadata_opt_in(monkeypatch) -> None: assert [result.status for result in results] == [429, 429] +@pytest.mark.asyncio +async def test_fetch_all_reuses_a_caller_owned_session(monkeypatch: pytest.MonkeyPatch) -> None: + session = object() + seen_sessions: list[object] = [] + + async def fake_fetch(*_args: Any, session: object, **_kwargs: Any) -> str: + seen_sessions.append(session) + return 'ok' + + monkeypatch.setattr(AsyncFetcher, 'fetch', fake_fetch) + + results = await AsyncFetcher.fetch_all( + ['https://one.example', 'https://two.example'], + session=session, + ) + + assert results == ['ok', 'ok'] + assert seen_sessions == [session, session] + + @pytest.mark.asyncio async def test_fetch_uses_http_proxy_when_enabled(monkeypatch) -> None: reset_dummy_sessions() diff --git a/tests/lib/test_source_runner.py b/tests/lib/test_source_runner.py index 9068b40e..3e9c5ae2 100644 --- a/tests/lib/test_source_runner.py +++ b/tests/lib/test_source_runner.py @@ -162,7 +162,7 @@ def test_source_factories_match_the_catalog() -> None: ), ('thc', 'theHarvester.lib.source_runner.thc.SearchThc', ('example.test',), {}), ('tomba', 'theHarvester.lib.source_runner.tombasearch.SearchTomba', ('example.test', 25, 5), {}), - ('urlscan', 'theHarvester.lib.source_runner.urlscan.SearchUrlscan', ('example.test',), {}), + ('urlscan', 'theHarvester.lib.source_runner.urlscan.SearchUrlscan', ('example.test', 25), {}), ('virustotal', 'theHarvester.lib.source_runner.virustotal.SearchVirustotal', ('example.test', 25), {}), ( 'waybackarchive', diff --git a/tests/test_main.py b/tests/test_main.py index a1e6ec69..0dbb52f2 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -3831,8 +3831,8 @@ async def test_routeviews_pivots_from_attributed_ips_without_expanding_discovere calls: list[tuple[tuple[object, ...], tuple[str, ...]]] = [] class FakeUrlscan: - def __init__(self, _word: str) -> None: - pass + def __init__(self, _word: str, limit: int) -> None: + assert limit == 500 async def process(self, _proxy: bool) -> None: return None diff --git a/theHarvester/discovery/bevigil.py b/theHarvester/discovery/bevigil.py index 6fdc4bc6..48bf1e53 100644 --- a/theHarvester/discovery/bevigil.py +++ b/theHarvester/discovery/bevigil.py @@ -1,6 +1,7 @@ from urllib.parse import urlsplit, urlunsplit from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -49,48 +50,46 @@ class SearchBeVigil: (url_endpoint, 'urls'), ) - for endpoint, field in requests: - try: - responses = await AsyncFetcher.fetch_all( - [endpoint], - json=True, - proxy=self.proxy, - headers=headers, - include_metadata=True, - ) - except Exception: - self._stop('failed', 'transport-error') - return - response = responses[0] if responses else None - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - 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(field), list): - self._stop('failed', 'invalid-response') - return + try: + async with AsyncFetcher.open_session( + headers=headers, + proxy=self.proxy, + request_timeout=60, + ) as session: + for endpoint, field in requests: + responses = await AsyncFetcher.fetch_all( + [endpoint], + json=True, + proxy=self.proxy, + headers=headers, + include_metadata=True, + session=session, + ) + response = responses[0] if responses else None + if error := provider_http_error(response): + self._stop(*error) + return + assert isinstance(response, FetcherResponse) + if not isinstance(response.body, dict) or not isinstance(response.body.get(field), list): + self._stop('failed', 'invalid-response') + return - malformed = False - for value in response.body[field]: - if field == 'subdomains': - if not isinstance(value, str): - malformed = True - elif hostname := normalize_scoped_hostname(value, self.word): - self.totalhosts.add(hostname) - elif url := self._scoped_url(value): - self.urls.add(url) - elif not isinstance(value, str): - malformed = True - if malformed: - self._stop('failed', 'invalid-response') + malformed = False + for value in response.body[field]: + if field == 'subdomains': + if not isinstance(value, str): + malformed = True + elif hostname := normalize_scoped_hostname(value, self.word): + self.totalhosts.add(hostname) + elif url := self._scoped_url(value): + self.urls.add(url) + elif not isinstance(value, str): + malformed = True + if malformed: + self._stop('failed', 'invalid-response') + except Exception: + self._stop('failed', 'transport-error') + return if self.execution_status is not None and self._has_results(): self.execution_status = 'partial' diff --git a/theHarvester/discovery/dymosearch.py b/theHarvester/discovery/dymosearch.py index da2a5534..ec49923d 100644 --- a/theHarvester/discovery/dymosearch.py +++ b/theHarvester/discovery/dymosearch.py @@ -1,6 +1,7 @@ from typing import Any from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -55,18 +56,10 @@ class SearchDymo: proxy=self.proxy, include_metadata=True, ) - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - return - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return diff --git a/theHarvester/discovery/fofa.py b/theHarvester/discovery/fofa.py index fee02e7f..f8f29af7 100644 --- a/theHarvester/discovery/fofa.py +++ b/theHarvester/discovery/fofa.py @@ -4,6 +4,7 @@ from typing import Any from urllib.parse import urlparse from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -106,18 +107,10 @@ class SearchFofa: json=True, include_metadata=True, ) - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - return - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return diff --git a/theHarvester/discovery/fullhuntsearch.py b/theHarvester/discovery/fullhuntsearch.py index 948745b3..f9b65a31 100644 --- a/theHarvester/discovery/fullhuntsearch.py +++ b/theHarvester/discovery/fullhuntsearch.py @@ -4,6 +4,7 @@ from typing import Any, ClassVar from urllib.parse import quote from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -154,7 +155,7 @@ class SearchFullHunt: """Returns the headers needed for API requests""" return {'User-Agent': Core.get_user_agent(), 'X-API-KEY': self.key} - async def _fetch_data(self, endpoint: str) -> dict[str, Any]: + async def _fetch_data(self, endpoint: str, session: Any | None = None) -> dict[str, Any]: """Generic method to fetch data from a specific endpoint""" url = f'{self.BASE_URL}/{endpoint}' response = await AsyncFetcher.fetch_all( @@ -163,20 +164,13 @@ class SearchFullHunt: headers=self._get_headers(), proxy=self.proxy, include_metadata=True, + session=session, ) - metadata = response[0] if response and isinstance(response[0], FetcherResponse) else None - if metadata is None: - self._stop('failed', 'transport-error') - raise RuntimeError('FullHunt request failed') - if metadata.status in {401, 403}: - self._stop('failed', 'access-denied') - raise RuntimeError('FullHunt request was denied') - if metadata.status == 429: - self._stop('rate-limited', 'http-429') - raise RuntimeError('FullHunt request was rate limited') - if not 200 <= metadata.status < 300: - self._stop('failed', f'http-{metadata.status}') - raise RuntimeError(f'FullHunt request failed with HTTP {metadata.status}') + metadata = response[0] if response else None + if error := provider_http_error(metadata): + self._stop(*error) + raise RuntimeError(f'FullHunt request failed: {error[1]}') + assert isinstance(metadata, FetcherResponse) if not isinstance(metadata.body, dict): self._stop('failed', 'invalid-response') raise ValueError('FullHunt returned malformed data') @@ -231,7 +225,7 @@ class SearchFullHunt: return ' '.join(query_parts) - async def advanced_search(self) -> dict[str, Any]: + async def advanced_search(self, session: Any | None = None) -> dict[str, Any]: """Perform an advanced search using the configured filters This method uses the search endpoint with the filters configured via add_filter @@ -244,17 +238,17 @@ class SearchFullHunt: query = self._build_query_string() encoded_query = quote(query) endpoint = f'search?query={encoded_query}' - return await self._fetch_data(endpoint) + return await self._fetch_data(endpoint, session) - async def get_domain_details(self) -> dict[str, Any]: + async def get_domain_details(self, session: Any | None = None) -> dict[str, Any]: """Get comprehensive details about a domain""" endpoint = f'domain/{self.word}/details' - return await self._fetch_data(endpoint) + return await self._fetch_data(endpoint, session) - async def get_subdomains(self) -> dict[str, Any]: + async def get_subdomains(self, session: Any | None = None) -> dict[str, Any]: """Get subdomains for a domain""" endpoint = f'domain/{self.word}/subdomains' - return await self._fetch_data(endpoint) + return await self._fetch_data(endpoint, session) async def get_host_details(self, host: str) -> dict[str, Any]: """Get detailed information about a specific host""" @@ -438,30 +432,35 @@ class SearchFullHunt: async def do_search(self) -> None: """Main search method that calls the various endpoints""" try: - # First get domain details which includes most information - domain_details = await self.get_domain_details() - if not isinstance(domain_details.get('hosts'), list): - raise ValueError('FullHunt returned malformed domain details') - self.total_results['domain_details'] = domain_details - await self.extract_data_from_domain_details(domain_details) + async with AsyncFetcher.open_session( + headers=self._get_headers(), + proxy=self.proxy, + request_timeout=60, + ) as session: + # First get domain details which includes most information + domain_details = await self.get_domain_details(session) + if not isinstance(domain_details.get('hosts'), list): + raise ValueError('FullHunt returned malformed domain details') + self.total_results['domain_details'] = domain_details + await self.extract_data_from_domain_details(domain_details) - # If no hosts found in domain details, try the dedicated subdomains endpoint - if not self.total_results['hosts']: - subdomains_response = await self.get_subdomains() - hosts = subdomains_response.get('hosts') - if not isinstance(hosts, list): - raise ValueError('FullHunt returned malformed subdomains') - for host in hosts: - if normalized_host := normalize_scoped_hostname(host, self.word): - self.total_results['hosts'].append(normalized_host) - else: - self._stop('failed', 'invalid-response') - logger.info('FullHunt ignored a malformed subdomain item') + # If no hosts found in domain details, try the dedicated subdomains endpoint + if not self.total_results['hosts']: + subdomains_response = await self.get_subdomains(session) + hosts = subdomains_response.get('hosts') + if not isinstance(hosts, list): + raise ValueError('FullHunt returned malformed subdomains') + for host in hosts: + if normalized_host := normalize_scoped_hostname(host, self.word): + self.total_results['hosts'].append(normalized_host) + else: + self._stop('failed', 'invalid-response') + logger.info('FullHunt ignored a malformed subdomain item') - # If filters are set, perform an advanced search - if self.filters: - search_results = await self.advanced_search() - await self.extract_data_from_search_results(search_results) + # If filters are set, perform an advanced search + if self.filters: + search_results = await self.advanced_search(session) + await self.extract_data_from_search_results(search_results) except Exception as error: if self.execution_status is None: diff --git a/theHarvester/discovery/netlas.py b/theHarvester/discovery/netlas.py index 1958b809..8c8a5a29 100644 --- a/theHarvester/discovery/netlas.py +++ b/theHarvester/discovery/netlas.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Any from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -26,21 +27,13 @@ class SearchNetlas: self.stop_reason = reason def _response_body(self, response: Any) -> Any | None: - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return None - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return None - if response.status == 402: + if isinstance(response, FetcherResponse) and response.status == 402: self._stop('failed', 'quota-exhausted') return None - if response.status == 429: - self._stop('rate-limited', 'http-429') - return None - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return None + assert isinstance(response, FetcherResponse) if response.body is None: self._stop('failed', 'invalid-response') return None diff --git a/theHarvester/discovery/onyphe.py b/theHarvester/discovery/onyphe.py index 5cc8e9d0..dd47fa6c 100644 --- a/theHarvester/discovery/onyphe.py +++ b/theHarvester/discovery/onyphe.py @@ -4,6 +4,7 @@ from ipaddress import ip_address from urllib.parse import urlparse from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.asn_attribution import AsnAttributionObservation, SubjectKind from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -71,18 +72,10 @@ class SearchOnyphe: json=True, include_metadata=True, ) - if not isinstance(metadata, FetcherResponse): - self._stop('failed', 'transport-error') - return - if metadata.status == 429: - self._stop('rate-limited', 'http-429') - return - if metadata.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if not 200 <= metadata.status < 300: - self._stop('failed', f'http-{metadata.status}') + if error := provider_http_error(metadata): + self._stop(*error) return + assert isinstance(metadata, FetcherResponse) if not isinstance(metadata.body, dict): self._stop('failed', 'invalid-response') return diff --git a/theHarvester/discovery/provider_response.py b/theHarvester/discovery/provider_response.py new file mode 100644 index 00000000..4569f7de --- /dev/null +++ b/theHarvester/discovery/provider_response.py @@ -0,0 +1,14 @@ +from theHarvester.lib.core import FetcherResponse + + +def provider_http_error(response: object) -> tuple[str, str] | None: + """Classify transport and HTTP failures shared by provider adapters.""" + if not isinstance(response, FetcherResponse): + return 'failed', 'transport-error' + if response.status in {401, 403}: + return 'failed', 'access-denied' + if response.status == 429: + return 'rate-limited', 'http-429' + if not 200 <= response.status < 300: + return 'failed', f'http-{response.status}' + return None diff --git a/theHarvester/discovery/searchhunterhow.py b/theHarvester/discovery/searchhunterhow.py index dc493f0f..1850057f 100644 --- a/theHarvester/discovery/searchhunterhow.py +++ b/theHarvester/discovery/searchhunterhow.py @@ -6,6 +6,7 @@ from typing import Any from dateutil.relativedelta import relativedelta from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -69,18 +70,10 @@ class SearchHunterHow: params=request_params, include_metadata=True, ) - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - return - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return diff --git a/theHarvester/discovery/securityscorecard.py b/theHarvester/discovery/securityscorecard.py index 1fe8cb50..40deda28 100644 --- a/theHarvester/discovery/securityscorecard.py +++ b/theHarvester/discovery/securityscorecard.py @@ -4,6 +4,7 @@ from ipaddress import ip_address from typing import Any from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -40,18 +41,10 @@ class SearchSecurityScorecard: self.stop_reason = reason def _response_body(self, response: Any) -> dict[str, Any] | None: - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return None - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return None - if response.status == 429: - self._stop('rate-limited', 'http-429') - return None - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return None + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return None diff --git a/theHarvester/discovery/securitytrailssearch.py b/theHarvester/discovery/securitytrailssearch.py index 16d9de52..39d25089 100644 --- a/theHarvester/discovery/securitytrailssearch.py +++ b/theHarvester/discovery/securitytrailssearch.py @@ -4,6 +4,7 @@ from ipaddress import ip_address from typing import Any from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -27,18 +28,10 @@ class SearchSecuritytrail: self.stop_reason = reason def _body(self, response: Any) -> dict[str, Any] | None: - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return None - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return None - if response.status == 429: - self._stop('rate-limited', 'http-429') - return None - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return None + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return None diff --git a/theHarvester/discovery/sherlockeye.py b/theHarvester/discovery/sherlockeye.py index 1e718619..a264764a 100644 --- a/theHarvester/discovery/sherlockeye.py +++ b/theHarvester/discovery/sherlockeye.py @@ -1,13 +1,11 @@ import logging -import random from ipaddress import ip_address as normalize_ip_address from typing import Any from urllib.parse import urlparse -import aiohttp - from theHarvester.discovery.constants import MissingKey -from theHarvester.lib.core import Core +from theHarvester.discovery.provider_response import provider_http_error +from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname logger = logging.getLogger(__name__) @@ -52,19 +50,6 @@ class SearchSherlockeye: 'Content-Type': 'application/json', } - def _proxy_url(self) -> str | None: - if isinstance(self.proxy, str) and self.proxy: - return self.proxy - if isinstance(self.proxy, bool) and self.proxy: - try: - proxy_list = Core.proxy_list() - proxy_urls = [*proxy_list.get('http', []), *proxy_list.get('socks5', [])] - if proxy_urls: - return random.choice(proxy_urls) - except Exception: - return None - return None - def _add_hostname(self, hostname: str) -> None: if normalized := normalize_scoped_hostname(hostname, self.word): self.totalhosts.add(normalized) @@ -157,34 +142,33 @@ class SearchSherlockeye: 'value': self.word, 'timeoutSeconds': self.DEFAULT_TIMEOUT_SECONDS, } - timeout = aiohttp.ClientTimeout(total=self.DEFAULT_TIMEOUT_SECONDS + 30) - try: - async with aiohttp.ClientSession(headers=self._headers(), timeout=timeout) as session: - async with session.post( + async with AsyncFetcher.open_session( + headers=self._headers(), + proxy=self.proxy, + request_timeout=self.DEFAULT_TIMEOUT_SECONDS + 30, + ) as session: + response = await AsyncFetcher.post_fetch( self.SYNC_SEARCH_URL, - json=payload, - proxy=self._proxy_url(), - ) as response: - if response.status != 200: - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - elif response.status == 429: - self._stop('rate-limited', 'http-429') - else: - self._stop('failed', f'http-{response.status}') - logger.info('Sherlockeye API request failed with status %s', response.status) - return - - try: - response_data = await response.json() - except (aiohttp.ContentTypeError, ValueError): - self._stop('failed', 'invalid-response') - return - if isinstance(response_data, dict): - self._extract_response(response_data) - else: - self._stop('failed', 'invalid-response') + session=session, + json=True, + include_metadata=True, + json_body=payload, + ) + if error := provider_http_error(response): + self._stop(*error) + status = response.status if isinstance(response, FetcherResponse) else 'transport' + logger.info('Sherlockeye API request failed with status %s: %s', status, error[1]) + return + assert isinstance(response, FetcherResponse) + if response.status != 200: + self._stop('failed', f'http-{response.status}') + logger.info('Sherlockeye API request failed with status %s', response.status) + return + if isinstance(response.body, dict): + self._extract_response(response.body) + else: + self._stop('failed', 'invalid-response') except Exception as error: self._stop('failed', 'transport-error') logger.info('Sherlockeye API error: %s', type(error).__name__) diff --git a/theHarvester/discovery/urlscan.py b/theHarvester/discovery/urlscan.py index ff569688..29adb7ca 100644 --- a/theHarvester/discovery/urlscan.py +++ b/theHarvester/discovery/urlscan.py @@ -3,6 +3,7 @@ from datetime import UTC, datetime from ipaddress import ip_address from urllib.parse import urlsplit +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.asn_attribution import AsnAttributionObservation, SubjectKind from theHarvester.lib.core import AsyncFetcher, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -12,8 +13,13 @@ logger = logging.getLogger(__name__) class SearchUrlscan: - def __init__(self, word) -> None: + MAX_PAGE_SIZE = 10_000 + + def __init__(self, word: str, limit: int) -> None: + if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0: + raise ValueError('URLScan limit must be a positive integer') self.word = word + self.limit = limit self.totalhosts: set = set() self.totalips: set = set() self.urls: set = set() @@ -141,11 +147,16 @@ class SearchUrlscan: collected_at = datetime.now(UTC) cursor = None seen_cursors: set[str] = set() + records_seen = 0 malformed = False try: async with AsyncFetcher.open_session(proxy=self.proxy) as session: - while True: - params = {'q': f'domain:{self.word}'} + while records_seen < self.limit: + remaining = self.limit - records_seen + params: dict[str, str | int] = { + 'q': f'domain:{self.word}', + 'size': min(self.MAX_PAGE_SIZE, remaining), + } if cursor is not None: params['search_after'] = cursor response = await AsyncFetcher.fetch( @@ -156,18 +167,10 @@ class SearchUrlscan: include_metadata=True, ) - 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}') + if error := provider_http_error(response): + self._stop(*error) return + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict) or not isinstance(response.body.get('results'), list): self._stop('failed', 'invalid-response') return @@ -181,8 +184,12 @@ class SearchUrlscan: self.stop_reason = None if self._has_results() else 'no-results' return - malformed = self._parse_results(results, collected_at) or malformed - next_cursor = self._cursor(results[-1]) + page_results = results[:remaining] + records_seen += len(page_results) + malformed = self._parse_results(page_results, collected_at) or malformed + if records_seen >= self.limit: + break + next_cursor = self._cursor(page_results[-1]) if next_cursor is None: self._stop('failed', 'invalid-cursor') return @@ -194,6 +201,13 @@ class SearchUrlscan: except Exception as error: self._stop('failed', 'transport-error') logger.info('URLScan request failed: %s', type(error).__name__) + return + + if self.execution_status is not None and self._has_results(): + self.execution_status = 'partial' + elif self.execution_status is None: + self.execution_status = 'completed' + self.stop_reason = None if self._has_results() else 'no-results' async def get_hostnames(self) -> set: return self.totalhosts diff --git a/theHarvester/discovery/virustotal.py b/theHarvester/discovery/virustotal.py index ea78c208..e040f49b 100644 --- a/theHarvester/discovery/virustotal.py +++ b/theHarvester/discovery/virustotal.py @@ -3,6 +3,7 @@ from __future__ import annotations from typing import Any from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -44,18 +45,10 @@ class SearchVirustotal: json=True, include_metadata=True, ) - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - return - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return diff --git a/theHarvester/discovery/whoisxml.py b/theHarvester/discovery/whoisxml.py index 3c2216e4..0c244dda 100644 --- a/theHarvester/discovery/whoisxml.py +++ b/theHarvester/discovery/whoisxml.py @@ -1,6 +1,7 @@ from __future__ import annotations from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname @@ -39,18 +40,10 @@ class SearchWhoisXML: json=True, include_metadata=True, ) - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return - if response.status == 429: - self._stop('rate-limited', 'http-429') - return - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return diff --git a/theHarvester/discovery/zoomeyesearch.py b/theHarvester/discovery/zoomeyesearch.py index d9a91774..6ea42feb 100644 --- a/theHarvester/discovery/zoomeyesearch.py +++ b/theHarvester/discovery/zoomeyesearch.py @@ -8,6 +8,7 @@ from typing import Any from urllib.parse import urlsplit, urlunsplit from theHarvester.discovery.constants import MissingKey +from theHarvester.discovery.provider_response import provider_http_error from theHarvester.lib.core import AsyncFetcher, Core, FetcherResponse from theHarvester.lib.hostnames import normalize_scoped_hostname from theHarvester.parsers import myparser @@ -89,18 +90,10 @@ class SearchZoomEye: 'fields': self.RESPONSE_FIELDS, }, ) - if not isinstance(response, FetcherResponse): - self._stop('failed', 'transport-error') - return None - if response.status in {401, 403}: - self._stop('failed', 'access-denied') - return None - if response.status == 429: - self._stop('rate-limited', 'http-429') - return None - if not 200 <= response.status < 300: - self._stop('failed', f'http-{response.status}') + if error := provider_http_error(response): + self._stop(*error) return None + assert isinstance(response, FetcherResponse) if not isinstance(response.body, dict): self._stop('failed', 'invalid-response') return None diff --git a/theHarvester/lib/core.py b/theHarvester/lib/core.py index 28fec96d..1f14bd34 100644 --- a/theHarvester/lib/core.py +++ b/theHarvester/lib/core.py @@ -945,13 +945,30 @@ class AsyncFetcher: @classmethod async def fetch_all( cls, - urls, - headers=None, + urls: list[str], + headers: dict[str, str] | None = None, params: Sized = '', json: bool = False, - proxy: bool = False, + proxy: str | bool | None = False, include_metadata: bool = False, - ) -> list: + *, + session: aiohttp.ClientSession | None = None, + ) -> list[Any]: + if session is not None: + return list( + await asyncio.gather( + *[ + AsyncFetcher.fetch( + session=session, + url=url, + params=params, + json=json, + include_metadata=include_metadata, + ) + for url in urls + ] + ) + ) # By default, timeout is 5 minutes; 60 seconds should suffice headers = cls._default_headers(headers) timeout = cls._request_timeout(60) diff --git a/theHarvester/lib/source_runner.py b/theHarvester/lib/source_runner.py index 13573946..07adb3e3 100644 --- a/theHarvester/lib/source_runner.py +++ b/theHarvester/lib/source_runner.py @@ -176,7 +176,7 @@ SOURCE_FACTORIES: dict[str, SourceFactory] = { 'subdomainfinderc99': lambda request: subdomainfinderc99.SearchSubdomainfinderc99(request.target), 'thc': lambda request: thc.SearchThc(request.target), 'tomba': lambda request: tombasearch.SearchTomba(request.target, request.limit, request.start), - 'urlscan': lambda request: urlscan.SearchUrlscan(request.target), + 'urlscan': lambda request: urlscan.SearchUrlscan(request.target, request.limit), 'virustotal': lambda request: virustotal.SearchVirustotal(request.target, request.limit), 'waybackarchive': lambda request: waybackarchive.SearchWaybackarchive(request.target, request.limit), 'whoisxml': lambda request: whoisxml.SearchWhoisXML(request.target, request.limit),