import asyncio import json import tomllib from datetime import UTC, datetime, timedelta from pathlib import Path from uuid import UUID import pytest from theHarvester import source_yields from theHarvester.lib import database as database_module from theHarvester.lib.active_evidence import ActionExecution, ActiveEvidence from theHarvester.lib.completed_result import CompletedResult, ResultObservation, SourceExecution from theHarvester.lib.database import ResultStore from theHarvester.lib.evidence_types import RESULT_KINDS RUN_ONE = UUID('11111111-1111-4111-8111-111111111111') RUN_TWO = UUID('22222222-2222-4222-8222-222222222222') def test_project_installs_harvest_yields_command() -> None: project = tomllib.loads(Path('pyproject.toml').read_text(encoding='utf-8')) assert project['project']['scripts']['harvest-yields'] == 'theHarvester.source_yields:main' def _completed_run( run_id: UUID, *, observations: tuple[ResultObservation, ...], resolved_hostnames: tuple[str, ...] = (), ) -> CompletedResult: started_at = datetime(2026, 8, 23, 12, tzinfo=UTC) + timedelta(minutes=int(str(run_id)[0])) sources = sorted({observation.source for observation in observations}) active_evidence = ( ActiveEvidence( executions=( ActionExecution.finish( action='dns-resolve', status='completed', duration_ms=1, groups={'hostname': resolved_hostnames}, ), ) ) if resolved_hostnames else ActiveEvidence() ) return CompletedResult.finish( run_id=run_id, target='example.test', started_at=started_at, completed_at=started_at + timedelta(seconds=1), groups={ observation.kind: [item.value for item in observations if item.kind == observation.kind] for observation in observations }, source_executions=tuple( SourceExecution( source=source, status='completed', duration_ms=1, result_count=sum(observation.source == source for observation in observations), ) for source in sources ), observations=observations, active_evidence=active_evidence, ) async def _create_database(database: Path) -> None: store = ResultStore(database) await store.initialize() await store.save_run( _completed_run( RUN_ONE, observations=( ResultObservation('alpha', 'hostname', 'shared.example.test'), ResultObservation('alpha', 'hostname', 'unique-alpha.example.test'), ResultObservation('beta', 'hostname', 'shared.example.test'), ResultObservation('beta', 'hostname', 'unique-beta.example.test'), ), resolved_hostnames=('shared.example.test', 'unique-alpha.example.test'), ) ) await store.save_run( _completed_run( RUN_TWO, observations=( ResultObservation('alpha', 'hostname', 'second-shared.example.test'), ResultObservation('beta', 'hostname', 'second-shared.example.test'), ResultObservation('gamma', 'hostname', 'second-gamma.example.test'), ResultObservation('alpha', 'ip', '192.0.2.1'), ResultObservation('beta', 'ip', '192.0.2.1'), ResultObservation('gamma', 'ip', '198.51.100.2'), ResultObservation('alpha', 'asn', 'AS64496'), ), resolved_hostnames=('second-shared.example.test',), ) ) await store.dispose() def test_missing_database_fails_without_creating_file( tmp_path: Path, capsys: pytest.CaptureFixture[str], ) -> None: database = tmp_path / 'missing.sqlite' with pytest.raises(SystemExit) as error: source_yields.main(['--database', str(database)]) assert error.value.code == 2 assert 'database does not exist' in capsys.readouterr().err assert not database.exists() def test_default_database_uses_the_standard_result_store( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: database = tmp_path / 'stash.sqlite' asyncio.run(_create_database(database)) monkeypatch.setattr(database_module, '_DEFAULT_DATABASE', database) assert source_yields.main([]) == 0 assert capsys.readouterr().out.startswith('Kind: hostname\nRun count: 2\n') def test_default_table_ranks_by_unique_per_run_and_aligns_columns( tmp_path: Path, capsys: pytest.CaptureFixture[str], ) -> None: database = tmp_path / 'runs.sqlite' asyncio.run(_create_database(database)) assert source_yields.main(['--database', str(database)]) == 0 assert capsys.readouterr().out.splitlines() == [ 'Kind: hostname', 'Run count: 2', 'SOURCE RUNS OBSERVED UNIQUE UNIQUE/RUN SHARED RESOLVED UNIQUE-RESOLVED UNIQUE-RESOLVED/RUN', 'gamma 1 1 1 1.00 0 0 0 0.00', 'alpha 2 3 1 0.50 2 3 1 0.50', 'beta 2 3 1 0.50 2 2 0 0.00', ] @pytest.mark.parametrize('kind', sorted(RESULT_KINDS)) def test_kind_accepts_every_result_kind_and_only_hostname_shows_resolution_columns( tmp_path: Path, capsys: pytest.CaptureFixture[str], kind: str, ) -> None: database = tmp_path / 'runs.sqlite' asyncio.run(_create_database(database)) assert source_yields.main(['--database', str(database), '--kind', kind]) == 0 lines = capsys.readouterr().out.splitlines() assert lines[0] == f'Kind: {kind}' assert ('RESOLVED' in lines[2]) is (kind == 'hostname') if kind == 'ip': assert [line.split() for line in lines[3:]] == [ ['gamma', '1', '1', '1', '1.00', '0'], ['alpha', '2', '1', '0', '0.00', '1'], ['beta', '2', '1', '0', '0.00', '1'], ] def test_run_id_selects_one_run( tmp_path: Path, capsys: pytest.CaptureFixture[str], ) -> None: database = tmp_path / 'runs.sqlite' asyncio.run(_create_database(database)) assert source_yields.main(['--database', str(database), '--run-id', str(RUN_ONE)]) == 0 lines = capsys.readouterr().out.splitlines() assert lines[1] == 'Run count: 1' assert [line.split() for line in lines[3:]] == [ ['alpha', '1', '2', '1', '1.00', '1', '2', '1', '1.00'], ['beta', '1', '2', '1', '1.00', '1', '1', '0', '0.00'], ] def test_unknown_run_id_fails_instead_of_reporting_an_empty_run( tmp_path: Path, capsys: pytest.CaptureFixture[str], ) -> None: database = tmp_path / 'runs.sqlite' asyncio.run(_create_database(database)) missing_run = UUID('33333333-3333-4333-8333-333333333333') with pytest.raises(SystemExit) as error: source_yields.main(['--database', str(database), '--run-id', str(missing_run)]) assert error.value.code == 2 assert 'completed result not found' in capsys.readouterr().err @pytest.mark.parametrize('kind', ['hostname', 'ip']) def test_json_format_is_machine_readable_and_uses_kind_specific_fields( tmp_path: Path, capsys: pytest.CaptureFixture[str], kind: str, ) -> None: database = tmp_path / 'runs.sqlite' asyncio.run(_create_database(database)) assert source_yields.main(['--database', str(database), '--kind', kind, '--format', 'json']) == 0 payload = json.loads(capsys.readouterr().out) assert payload['kind'] == kind assert payload['run_count'] == 2 assert [row['source'] for row in payload['source_yields']] == ['gamma', 'alpha', 'beta'] resolution_fields = { 'resolved_hostname_count', 'unique_resolved_hostname_count', 'unique_resolved_hostname_count_per_run', } assert all(resolution_fields <= row.keys() for row in payload['source_yields']) is (kind == 'hostname') assert {row['source']: row['run_count'] for row in payload['source_yields']} == {'alpha': 2, 'beta': 2, 'gamma': 1} assert {row['source']: row['unique_result_count_per_run'] for row in payload['source_yields']} == ( {'alpha': 0.5, 'beta': 0.5, 'gamma': 1.0} if kind == 'hostname' else {'alpha': 0.0, 'beta': 0.0, 'gamma': 1.0} )