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

from __future__ import annotations

import asyncio
import json
import logging
import os
import time
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any

from yuxi.channel.secrets.models import (
    SINGLE_VALUE_FILE_REF_ID,
    SecretRef,
    SecretSource,
    is_record,
)

logger = logging.getLogger(__name__)

DEFAULT_FILE_MAX_BYTES = 1024 * 1024
DEFAULT_FILE_TIMEOUT_MS = 5_000
DEFAULT_EXEC_TIMEOUT_MS = 5_000
DEFAULT_EXEC_MAX_OUTPUT_BYTES = 1024 * 1024
DEFAULT_EXEC_MAX_BATCH_BYTES = 256 * 1024


class SecretsProvider(ABC):
    """Abstract base for secret resolution providers."""

    @abstractmethod
    async def resolve(self, ref: SecretRef) -> str | None: ...

    @abstractmethod
    async def resolve_batch(self, refs: list[SecretRef]) -> dict[str, str | None]: ...

    @property
    @abstractmethod
    def source_type(self) -> SecretSource: ...


class EnvProvider(SecretsProvider):
    def __init__(self, allowlist: set[str] | None = None):
        self._allowlist = allowlist

    @property
    def source_type(self) -> SecretSource:
        return SecretSource.ENV

    async def resolve(self, ref: SecretRef) -> str | None:
        resolved = await self.resolve_batch([ref])
        return resolved.get(ref.ref)

    async def resolve_batch(self, refs: list[SecretRef]) -> dict[str, str | None]:
        result: dict[str, str | None] = {}
        for ref in refs:
            if self._allowlist is not None and ref.ref not in self._allowlist:
                logger.warning("EnvProvider: ref '%s' not in allowlist", ref.ref)
                result[ref.ref] = None
                continue
            value = os.getenv(ref.ref)
            if value is None:
                logger.warning("EnvProvider: env var '%s' not set", ref.ref)
            result[ref.ref] = value
        return result


class FileProvider(SecretsProvider):
    def __init__(
        self,
        base_dir: str | Path | None = None,
        mode: str = "json",
        max_bytes: int = DEFAULT_FILE_MAX_BYTES,
        timeout_ms: int = DEFAULT_FILE_TIMEOUT_MS,
        allow_insecure_path: bool = False,
        filename: str = "secrets.json",
    ):
        self._base_dir = Path(base_dir) if base_dir else Path.cwd()
        self._mode = mode
        self._max_bytes = max_bytes
        self._timeout_ms = timeout_ms
        self._allow_insecure_path = allow_insecure_path
        self._filename = filename
        self._payload_cache: Any = None
        self._payload_cache_mtime: float | None = None

    @property
    def source_type(self) -> SecretSource:
        return SecretSource.FILE

    async def resolve(self, ref: SecretRef) -> str | None:
        resolved = await self.resolve_batch([ref])
        return resolved.get(ref.ref)

    async def resolve_batch(self, refs: list[SecretRef]) -> dict[str, str | None]:
        payload = await self._read_payload()
        if payload is None:
            return {ref.ref: None for ref in refs}

        if self._mode == "singleValue":
            result: dict[str, str | None] = {}
            payload_str = payload if isinstance(payload, str) else json.dumps(payload)
            for ref in refs:
                if ref.ref != SINGLE_VALUE_FILE_REF_ID:
                    logger.warning(
                        "FileProvider: singleValue mode expects ref id '%s', got '%s'",
                        SINGLE_VALUE_FILE_REF_ID,
                        ref.ref,
                    )
                    result[ref.ref] = None
                else:
                    result[ref.ref] = payload_str
            return result

        result = {}
        for ref in refs:
            try:
                result[ref.ref] = self._resolve_json_pointer(payload, ref.ref)
            except (KeyError, TypeError) as e:
                logger.warning("FileProvider: JSON pointer '%s' not found: %s", ref.ref, e)
                result[ref.ref] = None
        return result

    def _resolve_json_pointer(self, payload: Any, pointer: str) -> str:
        if not pointer.startswith("/"):
            raise KeyError(f"JSON pointer must start with '/': {pointer}")
        tokens = pointer[1:].split("/")
        current = payload
        for token in tokens:
            token = token.replace("~1", "/").replace("~0", "~")
            if isinstance(current, dict):
                current = current[token]
            elif isinstance(current, list):
                try:
                    idx = int(token)
                except ValueError:
                    raise KeyError(f"Invalid array index '{token}' in JSON pointer")
                current = current[idx]
            else:
                raise KeyError(f"Cannot index into non-container at token '{token}'")
        if isinstance(current, str):
            return current
        return json.dumps(current)

    def _validate_path(self, file_path: Path) -> None:
        if self._allow_insecure_path:
            return
        resolved = file_path.resolve()
        base_resolved = self._base_dir.resolve()
        try:
            resolved.relative_to(base_resolved)
        except ValueError:
            raise PermissionError(
                f"FileProvider: path '{resolved}' is outside allowed base directory '{base_resolved}'"
            )

    async def _read_payload(self) -> Any:
        file_path = self._base_dir / self._filename
        self._validate_path(file_path)

        try:
            current_mtime = os.path.getmtime(file_path)
        except OSError:
            current_mtime = None

        if self._payload_cache_mtime is not None and current_mtime == self._payload_cache_mtime:
            return self._payload_cache

        try:
            content = await asyncio.wait_for(
                asyncio.to_thread(file_path.read_text, encoding="utf-8"),
                timeout=self._timeout_ms / 1000,
            )
            content = content.strip()
        except FileNotFoundError:
            logger.warning("FileProvider: file not found: %s", file_path)
            self._payload_cache_mtime = current_mtime
            self._payload_cache = None
            return None
        except OSError as e:
            logger.warning("FileProvider: read error for %s: %s", file_path, e)
            self._payload_cache_mtime = current_mtime
            self._payload_cache = None
            return None
        except TimeoutError:
            logger.warning("FileProvider: read timeout for %s", file_path)
            self._payload_cache = None
            return None

        self._payload_cache_mtime = current_mtime
        try:
            self._payload_cache = json.loads(content)
        except json.JSONDecodeError:
            if self._mode == "singleValue":
                self._payload_cache = content
            else:
                logger.warning("FileProvider: invalid JSON in %s", file_path)
                self._payload_cache = None
        return self._payload_cache

    def invalidate_cache(self) -> None:
        self._payload_cache_mtime = None
        self._payload_cache = None


class ExecProvider(SecretsProvider):
    EXEC_PROTOCOL_VERSION = 1

    def __init__(
        self,
        command: str | None = None,
        args: list[str] | None = None,
        trusted_dirs: set[str] | None = None,
        timeout_ms: int = DEFAULT_EXEC_TIMEOUT_MS,
        no_output_timeout_ms: int | None = None,
        max_output_bytes: int = DEFAULT_EXEC_MAX_OUTPUT_BYTES,
        max_batch_bytes: int = DEFAULT_EXEC_MAX_BATCH_BYTES,
        stderr_max_bytes: int = 4096,
        json_only: bool = True,
        pass_env: list[str] | None = None,
        extra_env: dict[str, str] | None = None,
        allow_insecure_path: bool = False,
        allow_symlink_command: bool = False,
    ):
        self._command = command
        self._args = args or []
        self._trusted_dirs = trusted_dirs or set()
        self._timeout_ms = timeout_ms
        self._no_output_timeout_ms = no_output_timeout_ms or timeout_ms
        self._max_output_bytes = max_output_bytes
        self._max_batch_bytes = max_batch_bytes
        self._stderr_max_bytes = stderr_max_bytes
        self._json_only = json_only
        self._pass_env = pass_env or []
        self._extra_env = extra_env or {}
        self._allow_insecure_path = allow_insecure_path
        self._allow_symlink_command = allow_symlink_command

    @property
    def source_type(self) -> SecretSource:
        return SecretSource.EXEC

    async def resolve(self, ref: SecretRef) -> str | None:
        resolved = await self.resolve_batch([ref])
        return resolved.get(ref.ref)

    def _validate_command(self, cmd: str) -> str:
        if not cmd:
            raise ValueError("ExecProvider: command is empty")

        cmd_path = Path(cmd)
        if not self._allow_insecure_path and not cmd_path.is_absolute():
            raise ValueError(f"ExecProvider: command must be an absolute path, got '{cmd}'")

        resolved_cmd = cmd_path.resolve()

        if not self._allow_symlink_command:
            real_cmd = resolved_cmd.resolve()
            if real_cmd != resolved_cmd:
                raise ValueError(f"ExecProvider: command '{cmd}' is a symlink and allow_symlink_command is False")

        if self._trusted_dirs:
            cmd_dir = resolved_cmd.parent
            in_trusted = any(
                cmd_dir == Path(td).resolve() or cmd_dir.is_relative_to(Path(td).resolve()) for td in self._trusted_dirs
            )
            if not in_trusted:
                raise ValueError(f"ExecProvider: command '{resolved_cmd}' is not in trusted directories")

        return str(resolved_cmd)

    async def resolve_batch(self, refs: list[SecretRef]) -> dict[str, str | None]:
        ids = list(dict.fromkeys(r.ref for r in refs))

        if not self._command:
            if not refs:
                return {}
            raise ValueError("ExecProvider: no command configured")

        cmd = self._command
        cmd = self._validate_command(cmd)

        request = {
            "protocolVersion": self.EXEC_PROTOCOL_VERSION,
            "provider": "exec",
            "ids": ids,
        }
        input_data = json.dumps(request)

        if len(input_data.encode("utf-8")) > self._max_batch_bytes:
            logger.warning(
                "ExecProvider: request exceeds max batch bytes (%d > %d)",
                len(input_data.encode("utf-8")),
                self._max_batch_bytes,
            )
            return {ref.ref: None for ref in refs}

        child_env: dict[str, str] = {}
        for key in self._pass_env:
            value = os.getenv(key)
            if value is not None:
                child_env[key] = value
        child_env.update(self._extra_env)

        proc = None
        try:
            proc = await asyncio.create_subprocess_exec(
                cmd,
                *self._args,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=child_env,
            )

            if proc.stdin is None or proc.stdout is None:
                logger.warning("ExecProvider: failed to open process streams")
                await self._kill_process(proc)
                return {ref.ref: None for ref in refs}

            proc.stdin.write(input_data.encode())
            await proc.stdin.drain()
            proc.stdin.close()
            await proc.stdin.wait_closed()

            stdout_bytes, stderr_bytes = await self._read_with_limits(proc)
        except asyncio.CancelledError:
            if proc is not None:
                await self._kill_process(proc)
            raise
        except TimeoutError:
            if proc is not None:
                await self._kill_process(proc)
            logger.warning("ExecProvider: command timed out: %s", cmd[:80])
            return {ref.ref: None for ref in refs}
        except Exception as e:
            if proc is not None:
                await self._kill_process(proc)
            logger.warning("ExecProvider: command failed: %s: %s", cmd[:80], e)
            return {ref.ref: None for ref in refs}

        if proc.returncode != 0:
            stderr_text = stderr_bytes.decode("utf-8", errors="replace")[:200] if stderr_bytes else ""
            logger.warning(
                "ExecProvider: command exit code %d: %s",
                proc.returncode,
                stderr_text,
            )
            return {ref.ref: None for ref in refs}

        stdout = stdout_bytes.decode("utf-8").strip() if stdout_bytes else ""
        return self._parse_exec_values(ids, stdout)

    async def _read_with_limits(self, proc) -> tuple[bytes, bytes]:
        stdout_chunks: list[bytes] = []
        stderr_chunks: list[bytes] = []
        global_deadline = time.monotonic() + self._timeout_ms / 1000

        async def _read_stream(
            stream: asyncio.StreamReader | None,
            chunks: list[bytes],
            byte_counter: list[int],
            label: str,
            max_bytes: int,
        ) -> None:
            if stream is None:
                return
            while True:
                remaining = global_deadline - time.monotonic()
                if remaining <= 0:
                    raise TimeoutError(f"ExecProvider: global timeout ({label})")

                chunk_timeout = min(remaining, self._no_output_timeout_ms / 1000)
                try:
                    chunk = await asyncio.wait_for(stream.read(8192), timeout=chunk_timeout)
                except TimeoutError:
                    raise TimeoutError(f"ExecProvider: no output timeout ({label})")

                if not chunk:
                    return

                byte_counter[0] += len(chunk)
                if byte_counter[0] > max_bytes:
                    raise ValueError(
                        f"ExecProvider: {label} output exceeded max bytes ({byte_counter[0]} > {max_bytes})"
                    )

                chunks.append(chunk)

        try:
            stdout_counter = [0]
            stderr_counter = [0]
            stderr_limit = min(self._max_output_bytes, self._stderr_max_bytes)
            await asyncio.gather(
                _read_stream(proc.stdout, stdout_chunks, stdout_counter, "stdout", self._max_output_bytes),
                _read_stream(proc.stderr, stderr_chunks, stderr_counter, "stderr", stderr_limit),
            )
        except Exception:
            await self._kill_process(proc)
            raise

        await proc.wait()
        return b"".join(stdout_chunks), b"".join(stderr_chunks)

    async def _kill_process(self, proc) -> None:
        try:
            proc.kill()
            await asyncio.wait_for(proc.wait(), timeout=5.0)
        except TimeoutError:
            logger.warning("ExecProvider: process did not exit after kill")
        except Exception as e:
            logger.warning("ExecProvider: failed to kill process: %s", e)

    def _parse_exec_values(self, ids: list[str], stdout: str) -> dict[str, str | None]:
        trimmed = stdout.strip()
        if not trimmed:
            logger.warning("ExecProvider: empty stdout")
            return {ref_id: None for ref_id in ids}

        if not self._json_only and len(ids) == 1:
            try:
                parsed = json.loads(trimmed)
            except json.JSONDecodeError:
                return {ids[0]: trimmed}
        else:
            try:
                parsed = json.loads(trimmed)
            except json.JSONDecodeError:
                logger.warning("ExecProvider: invalid JSON response")
                return {ref_id: None for ref_id in ids}

        if not is_record(parsed):
            if not self._json_only and len(ids) == 1 and isinstance(parsed, str):
                return {ids[0]: parsed}
            logger.warning("ExecProvider: response must be an object")
            return {ref_id: None for ref_id in ids}

        protocol_version = parsed.get("protocolVersion")
        if protocol_version != self.EXEC_PROTOCOL_VERSION:
            logger.warning("ExecProvider: protocolVersion must be %d", self.EXEC_PROTOCOL_VERSION)
            return {ref_id: None for ref_id in ids}

        response_values = parsed.get("values")
        if not is_record(response_values):
            logger.warning('ExecProvider: response missing "values"')
            return {ref_id: None for ref_id in ids}

        response_errors = parsed.get("errors") if is_record(parsed.get("errors")) else None

        result: dict[str, str | None] = {}
        for ref_id in ids:
            if response_errors and ref_id in response_errors:
                error_entry = response_errors[ref_id]
                if is_record(error_entry) and isinstance(error_entry.get("message"), str):
                    msg = error_entry["message"].strip()
                    logger.warning("ExecProvider: error for id '%s': %s", ref_id, msg)
                else:
                    logger.warning("ExecProvider: error for id '%s'", ref_id)
                result[ref_id] = None
                continue

            if ref_id not in response_values:
                logger.warning("ExecProvider: response missing id '%s'", ref_id)
                result[ref_id] = None
                continue

            value = response_values[ref_id]
            result[ref_id] = str(value) if not isinstance(value, str) else value

        return result