hydrogen/core/runner.py

from __future__ import annotations

import logging
import sys
from collections import Counter
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
from pathlib import Path
from time import perf_counter

from core.evaluation import evaluate_exit_code
from core.module_loader import load_module_descriptors, load_prevalidated_workers
from core.runtime import resolve_config, validate_reporting_plugins
from core.schemas.config import Config
from core.schemas.results import AuditResults
from core.schemas.status import AuditStatus
from reporting.bootstrap import bootstrap_reporting
from reporting.models import ModuleAuditResult, ModuleExecutionStats, PluginRuntimeError
from reporting.service import ReportService

logger = logging.getLogger(__name__)


def run_audit(modules_path: Path, config: Config, package_prefix: str = "modules") -> int:
    _configure_logging(config)

    loaded_modules, module_discovery_errors = load_module_descriptors(
        modules_path, config, package_prefix
    )
    reporting_bootstrap = bootstrap_reporting(config.plugin_packages)
    reporting_validation_errors = validate_reporting_plugins(
        reporting_bootstrap.renderer_registry,
        reporting_bootstrap.transport_registry,
    )
    resolved_config, config_errors = resolve_config(
        config,
        loaded_modules,
        reporting_bootstrap.renderer_registry,
        reporting_bootstrap.transport_registry,
        reporting_bootstrap.errors,
    )

    worker_load_result = load_prevalidated_workers(loaded_modules, resolved_config)
    plugin_errors = [
        *module_discovery_errors,
        *reporting_validation_errors,
        *config_errors,
        *worker_load_result.errors,
    ]

    results = _run_workers(worker_load_result.loaded_workers, resolved_config)
    exit_code = evaluate_exit_code([r.result for r in results], resolved_config)
    report_service = ReportService(
        renderer_registry=reporting_bootstrap.renderer_registry,
        transport_registry=reporting_bootstrap.transport_registry,
    )
    report = report_service.build_report(results, exit_code, plugin_errors)

    if resolved_config.dry_run:
        logger.info("dry_run enabled, skipping report publish")
    else:
        publish_errors = _publish_reports(report_service, report, resolved_config.reports)
        if publish_errors:
            report.plugin_errors.extend(publish_errors)

    return report.exit_code


def _run_workers(workers: list, config) -> list[ModuleAuditResult]:
    if not workers:
        return []

    indexed_workers = list(enumerate(workers))
    results_by_index: dict[int, ModuleAuditResult] = {}

    with ThreadPoolExecutor(max_workers=config.max_concurrency) as executor:
        pending: dict[Future[ModuleAuditResult], int] = {}
        next_index = 0

        while next_index < len(indexed_workers) and len(pending) < config.max_concurrency:
            index, worker = indexed_workers[next_index]
            pending[executor.submit(_run_single_worker, worker)] = index
            next_index += 1

        stop_submitting = False
        while pending:
            done, _ = wait(tuple(pending), return_when=FIRST_COMPLETED)
            for future in done:
                index = pending.pop(future)
                result = future.result()
                results_by_index[index] = result

                if config.fail_fast and result.result.status is AuditStatus.FAIL:
                    stop_submitting = True

            while (
                not stop_submitting
                and next_index < len(indexed_workers)
                and len(pending) < config.max_concurrency
            ):
                index, worker = indexed_workers[next_index]
                pending[executor.submit(_run_single_worker, worker)] = index
                next_index += 1

    return [results_by_index[index] for index in sorted(results_by_index)]


def _run_single_worker(worker) -> ModuleAuditResult:
    started_at = perf_counter()
    try:
        result = worker.worker.run()
        duration_seconds = perf_counter() - started_at
        counts = Counter(finding.severity for finding in result.findings)
        return ModuleAuditResult(
            module=worker.manifest,
            result=result,
            stats=ModuleExecutionStats(
                duration_seconds=duration_seconds, finding_counts=dict(counts)
            ),
        )
    except Exception as exc:
        logger.exception("worker %s failed during execution", worker.manifest.identifier)
        duration_seconds = perf_counter() - started_at
        failed_result = AuditResults(status=AuditStatus.FAIL, findings=[], risk_level=1.0)
        return ModuleAuditResult(
            module=worker.manifest,
            result=failed_result,
            stats=ModuleExecutionStats(duration_seconds=duration_seconds, finding_counts={}),
            error=PluginRuntimeError(
                plugin_kind="module",
                plugin_name=worker.manifest.identifier,
                stage="run",
                message=str(exc),
            ),
        )


def _configure_logging(config: Config) -> None:
    level_name = config.logging.level.upper()
    level = getattr(logging, level_name, logging.INFO)
    root_logger = logging.getLogger()
    stream = sys.stdout if config.logging.output.value == "stdout" else sys.stderr

    handler = logging.StreamHandler(stream)
    handler.setLevel(level)
    handler.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s"))

    root_logger.handlers.clear()
    root_logger.addHandler(handler)
    root_logger.setLevel(level)


def _publish_reports(report_service, report, outputs) -> list[PluginRuntimeError]:
    errors: list[PluginRuntimeError] = []
    for output in outputs:
        try:
            report_service.publish(report, output)
        except Exception as exc:
            logger.exception(
                "failed to publish output %s via %s",
                output.renderer.type,
                output.transport.type,
            )
            errors.append(
                PluginRuntimeError(
                    plugin_kind="report_output",
                    plugin_name=f"{output.renderer.type}->{output.transport.type}",
                    stage="publish",
                    message=str(exc),
                )
            )
    return errors