ForcePilot/backend/package/yuxi/channel/secrets/audit.py

from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Any

from yuxi.channel.secrets.models import (
    SecretRef,
    SecretsAuditFinding,
    SecretsAuditReport,
    SecretSource,
    is_expected_resolved_secret_value,
    is_non_empty_string,
    is_secret_field_name,
    is_secret_ref,
    parse_secret_ref,
    secret_ref_key,
)
from yuxi.channel.secrets.providers import (
    SecretsProvider,
)
from yuxi.channel.secrets.resolver import SecretResolver

logger = logging.getLogger(__name__)


class SecretsAuditResult:
    def __init__(
        self,
        report: SecretsAuditReport,
        exit_code: int = 0,
    ):
        self.report = report
        self.exit_code = exit_code


def _resolve_json_path(path: str, key: str) -> str:
    if not path:
        return key
    return f"{path}.{key}"


def _find_plaintext_secrets(
    config: dict,
    path: str = "",
    file_label: str = "config",
) -> list[SecretsAuditFinding]:
    findings: list[SecretsAuditFinding] = []

    def _walk(current: Any, current_path: str):
        if not isinstance(current, dict):
            return

        for key, value in current.items():
            entry_path = _resolve_json_path(current_path, key)

            if is_secret_ref(value):
                continue

            if isinstance(value, dict):
                _walk(value, entry_path)
            elif isinstance(value, list):
                for i, item in enumerate(value):
                    if isinstance(item, dict) and not is_secret_ref(item):
                        _walk(item, f"{entry_path}[{i}]")
            elif is_secret_field_name(key) and is_non_empty_string(value) and value != "***":
                findings.append(
                    SecretsAuditFinding(
                        code="PLAINTEXT_FOUND",
                        severity="warn",
                        file=file_label,
                        json_path=entry_path,
                        message=f"{entry_path} is stored as plaintext.",
                    )
                )

    _walk(config, path)
    return findings


async def _check_ref_resolvability(
    findings: list[SecretsAuditFinding],
    config: dict,
    providers: dict[SecretSource, SecretsProvider] | None = None,
    allow_exec: bool = True,
) -> dict[str, Any]:
    resolver = SecretResolver(providers)
    refs_collected: list[tuple[SecretRef, str, str]] = []

    def _collect_refs(current: Any, current_path: str, file_label: str):
        if not isinstance(current, dict):
            return

        for key, value in current.items():
            entry_path = _resolve_json_path(current_path, key)

            if is_secret_ref(value):
                ref = parse_secret_ref(value)
                if ref:
                    refs_collected.append((ref, entry_path, file_label))
            elif isinstance(value, dict):
                _collect_refs(value, entry_path, file_label)
            elif isinstance(value, list):
                for i, item in enumerate(value):
                    if isinstance(item, dict):
                        _collect_refs(item, f"{entry_path}[{i}]", file_label)

    _collect_refs(config, "", "config")

    if not allow_exec:
        exec_refs = [r for r, _, _ in refs_collected if r.source == SecretSource.EXEC]
        skipped = len(exec_refs)
        refs_collected = [(r, p, f) for r, p, f in refs_collected if r.source != SecretSource.EXEC]
    else:
        skipped = 0

    refs = [r for r, _, _ in refs_collected]
    if not refs:
        return {"refs_checked": 0, "skipped_exec_refs": skipped, "resolvability_complete": True}

    try:
        resolved = await resolver.resolve_secret_ref_values(refs)
    except Exception as e:
        for ref, entry_path, file_label in refs_collected:
            findings.append(
                SecretsAuditFinding(
                    code="REF_UNRESOLVED",
                    severity="error",
                    file=file_label,
                    json_path=entry_path,
                    message=f"Failed to resolve {secret_ref_key(ref)} ({e}).",
                )
            )
        return {
            "refs_checked": len(refs),
            "skipped_exec_refs": skipped,
            "resolvability_complete": skipped == 0,
        }

    for ref, entry_path, file_label in refs_collected:
        key = secret_ref_key(ref)
        if key not in resolved or resolved[key] is None:
            findings.append(
                SecretsAuditFinding(
                    code="REF_UNRESOLVED",
                    severity="error",
                    file=file_label,
                    json_path=entry_path,
                    message=f"Failed to resolve {key} (resolved value is missing).",
                )
            )
        elif not is_expected_resolved_secret_value(resolved[key], "string"):
            findings.append(
                SecretsAuditFinding(
                    code="REF_UNRESOLVED",
                    severity="error",
                    file=file_label,
                    json_path=entry_path,
                    message=f"Failed to resolve {key} (not a non-empty string).",
                )
            )

    return {
        "refs_checked": len(refs),
        "skipped_exec_refs": skipped,
        "resolvability_complete": skipped == 0,
    }


def _unquote_value(value: str) -> str:
    """Remove matching outer quotes and handle escaped quotes."""
    value = value.strip()
    if len(value) >= 2:
        if value[0] == '"' and value[-1] == '"':
            return value[1:-1].replace('\\"', '"')
        if value[0] == "'" and value[-1] == "'":
            return value[1:-1].replace("\\'", "'")
    return value


def _scan_env_file(env_path: str | Path, known_env_vars: set[str]) -> list[SecretsAuditFinding]:
    findings: list[SecretsAuditFinding] = []
    env_path = Path(env_path)
    if not env_path.exists():
        return findings

    try:
        content = env_path.read_text(encoding="utf-8")
    except OSError:
        return findings

    pattern = re.compile(r"^\s*(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(.+)$")
    for line in content.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        match = pattern.match(line)
        if not match:
            continue
        key = match.group(1)
        if key not in known_env_vars:
            continue
        value = _unquote_value(match.group(2))
        if value:
            findings.append(
                SecretsAuditFinding(
                    code="PLAINTEXT_FOUND",
                    severity="warn",
                    file=str(env_path),
                    json_path=f"$env.{key}",
                    message=f"Potential secret found in .env ({key}).",
                )
            )
    return findings


async def run_secrets_audit(
    config: dict[str, Any],
    providers: dict[SecretSource, SecretsProvider] | None = None,
    env_path: str | Path | None = None,
    known_env_vars: set[str] | None = None,
    allow_exec: bool = True,
) -> SecretsAuditResult:
    files_scanned: list[str] = ["config"]

    findings = _find_plaintext_secrets(config)

    resolution = await _check_ref_resolvability(
        findings,
        config,
        providers=providers,
        allow_exec=allow_exec,
    )

    if env_path and known_env_vars:
        env_findings = _scan_env_file(env_path, known_env_vars)
        findings.extend(env_findings)
        files_scanned.append(str(env_path))

    plaintext_count = sum(1 for f in findings if f.code == "PLAINTEXT_FOUND")
    unresolved_count = sum(1 for f in findings if f.code == "REF_UNRESOLVED")

    if unresolved_count > 0:
        status = "unresolved"
    elif findings:
        status = "findings"
    else:
        status = "clean"

    report = SecretsAuditReport(
        version=1,
        status=status,
        resolution=resolution,
        files_scanned=files_scanned,
        summary={
            "plaintextCount": plaintext_count,
            "unresolvedRefCount": unresolved_count,
        },
        findings=findings,
    )

    exit_code = 2 if unresolved_count > 0 else 0

    return SecretsAuditResult(report=report, exit_code=exit_code)


async def run_secrets_audit_strict(
    config: dict[str, Any],
    providers: dict[SecretSource, SecretsProvider] | None = None,
    env_path: str | Path | None = None,
    known_env_vars: set[str] | None = None,
    allow_exec: bool = True,
) -> SecretsAuditResult:
    result = await run_secrets_audit(
        config=config,
        providers=providers,
        env_path=env_path,
        known_env_vars=known_env_vars,
        allow_exec=allow_exec,
    )
    if result.report.status == "findings" and result.exit_code == 0:
        result.exit_code = 1
    return result