Nexum-Expanded/runtime/src/nexum_runtime/tooling/security.py

"""Least-privilege policy, action approvals, and secret-safe tool output."""

from __future__ import annotations

import fcntl
import hashlib
import hmac
import json
import os
import re
import secrets
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Literal, Mapping, cast

from .contracts import ToolCall, ToolRisk, ToolSpec
from .sandbox import control_root


PolicyDecision = Literal["allow", "approve", "deny"]
_SECRET_ENV_MARKERS = (
    "API_KEY",
    "PASSWORD",
    "PRIVATE_KEY",
    "SECRET",
    "TOKEN",
)
_BEARER_RE = re.compile(r"(?i)(authorization\s*:\s*bearer\s+)[^\s,;]+")
_URL_CREDENTIAL_RE = re.compile(r"(https?://)[^/@\s]+@", re.IGNORECASE)
_SENSITIVE_KEY_RE = re.compile(
    r"(?:^|_)(?:access_token|api_key|authorization|bearer|client_secret|"
    r"credential|password|passphrase|private_key|secret|token)(?:$|_)",
    re.IGNORECASE,
)


def _atomic_json(path: Path, payload: Mapping[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    os.chmod(path.parent, 0o700)
    temporary = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
    temporary.write_text(
        json.dumps(payload, sort_keys=True, separators=(",", ":")),
        encoding="utf-8",
    )
    temporary.chmod(0o600)
    with temporary.open("rb") as handle:
        os.fsync(handle.fileno())
    os.replace(temporary, path)
    directory_fd = os.open(path.parent, os.O_RDONLY)
    try:
        os.fsync(directory_fd)
    finally:
        os.close(directory_fd)


class SecretRedactor:
    """Redact execution-only credentials before output reaches model or logs."""

    def __init__(
        self,
        environment: Mapping[str, str] | None = None,
        extra_values: tuple[str, ...] = (),
    ) -> None:
        source = os.environ if environment is None else environment
        values = {
            str(value)
            for name, value in source.items()
            if value
            and len(str(value)) >= 8
            and any(marker in name.upper() for marker in _SECRET_ENV_MARKERS)
        }
        values.update(value for value in extra_values if len(value) >= 8)
        self._values = tuple(sorted(values, key=len, reverse=True))

    def redact(self, text: str) -> str:
        redacted = text
        for value in self._values:
            redacted = redacted.replace(value, "[REDACTED]")
        redacted = _BEARER_RE.sub(r"\1[REDACTED]", redacted)
        return _URL_CREDENTIAL_RE.sub(r"\1[REDACTED]@", redacted)


def redact_sensitive_value(
    value: Any,
    *,
    redactor: SecretRedactor | None = None,
    sensitive: bool = False,
) -> Any:
    """Recursively redact structured execution data before durable reuse."""

    active = redactor if redactor is not None else SecretRedactor()
    if sensitive:
        return "[REDACTED]"
    if isinstance(value, str):
        return active.redact(value)
    if isinstance(value, list):
        return [
            redact_sensitive_value(item, redactor=active)
            for item in value
        ]
    if isinstance(value, tuple):
        return tuple(
            redact_sensitive_value(item, redactor=active)
            for item in value
        )
    if isinstance(value, Mapping):
        return {
            str(key): redact_sensitive_value(
                item,
                redactor=active,
                sensitive=bool(_SENSITIVE_KEY_RE.search(str(key))),
            )
            for key, item in value.items()
        }
    return value


@dataclass(frozen=True)
class ToolPolicy:
    read: PolicyDecision = "allow"
    workspace_write: PolicyDecision = "allow"
    external_effect: PolicyDecision = "approve"
    destructive: PolicyDecision = "approve"
    tool_overrides: tuple[tuple[str, PolicyDecision], ...] = ()

    @classmethod
    def load(cls, workspace: str | Path) -> "ToolPolicy":
        path = control_root(workspace) / "policy.json"
        if not path.is_file():
            return cls()
        payload = json.loads(path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            raise ValueError("tool policy must contain an object")
        allowed = {"allow", "approve", "deny"}

        def decision(name: str, default: PolicyDecision) -> PolicyDecision:
            value = str(payload.get(name, default))
            if value not in allowed:
                raise ValueError(f"invalid policy decision for {name}")
            return cast(PolicyDecision, value)

        raw_overrides = payload.get("tools")
        overrides: list[tuple[str, PolicyDecision]] = []
        if isinstance(raw_overrides, dict):
            for name, value in sorted(raw_overrides.items()):
                rendered = str(value)
                if rendered not in allowed:
                    raise ValueError(f"invalid tool policy decision for {name}")
                overrides.append((str(name), cast(PolicyDecision, rendered)))
        return cls(
            read=decision("read", "allow"),
            workspace_write=decision("workspace_write", "allow"),
            external_effect=decision("external_effect", "approve"),
            destructive=decision("destructive", "approve"),
            tool_overrides=tuple(overrides),
        )

    def decision(self, spec: ToolSpec) -> PolicyDecision:
        overrides = dict(self.tool_overrides)
        if spec.name in overrides:
            return cast(PolicyDecision, overrides[spec.name])
        return cast(PolicyDecision, getattr(self, spec.risk))


@dataclass(frozen=True)
class ApprovalRecord:
    approval_id: str
    session_sha256: str
    action_sha256: str
    tool: str
    risk: ToolRisk
    status: str
    created_unix: int
    decided_unix: int = 0
    consumed_unix: int = 0


class ApprovalStore:
    """One-shot approval records bound to one exact model-selected action."""

    def __init__(self, workspace: str | Path) -> None:
        self.root = control_root(workspace) / "approvals"
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        os.chmod(self.root, 0o700)

    @staticmethod
    def action_sha256(session_id: str, call: ToolCall) -> str:
        canonical = json.dumps(
            {
                "session_id": session_id,
                "call_id": call.call_id,
                "name": call.name,
                "args": call.args,
            },
            sort_keys=True,
            separators=(",", ":"),
        ).encode("utf-8")
        return hashlib.sha256(canonical).hexdigest()

    def _path(self, approval_id: str) -> Path:
        if not re.fullmatch(r"[a-f0-9]{32}", approval_id):
            raise ValueError("approval id is invalid")
        return self.root / f"{approval_id}.json"

    def _load_locked(self, path: Path) -> ApprovalRecord:
        payload = json.loads(path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            raise RuntimeError("approval record is invalid")
        return ApprovalRecord(**payload)

    def request(
        self,
        session_id: str,
        call: ToolCall,
        spec: ToolSpec,
    ) -> ApprovalRecord:
        action_sha256 = self.action_sha256(session_id, call)
        for path in self.root.glob("*.json"):
            try:
                existing = self._load_locked(path)
            except (OSError, TypeError, ValueError, RuntimeError):
                continue
            if (
                hmac.compare_digest(existing.action_sha256, action_sha256)
                and existing.status in {"pending", "approved"}
            ):
                return existing
        record = ApprovalRecord(
            approval_id=secrets.token_hex(16),
            session_sha256=hashlib.sha256(session_id.encode("utf-8")).hexdigest(),
            action_sha256=action_sha256,
            tool=call.name,
            risk=spec.risk,
            status="pending",
            created_unix=int(time.time()),
        )
        _atomic_json(self._path(record.approval_id), asdict(record))
        return record

    def decide(
        self,
        approval_id: str,
        *,
        approved: bool,
        session_id: str = "",
    ) -> ApprovalRecord:
        path = self._path(approval_id)
        lock_path = path.with_suffix(".lock")
        with lock_path.open("a+b") as lock_handle:
            lock_path.chmod(0o600)
            fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX)
            try:
                record = self._load_locked(path)
                if session_id:
                    expected_session = hashlib.sha256(
                        session_id.encode("utf-8")
                    ).hexdigest()
                    if not hmac.compare_digest(
                        record.session_sha256, expected_session
                    ):
                        raise PermissionError(
                            "approval does not belong to this session"
                        )
                if record.status not in {"pending", "approved"}:
                    raise RuntimeError("approval is no longer pending")
                updated = ApprovalRecord(
                    **{
                        **asdict(record),
                        "status": "approved" if approved else "denied",
                        "decided_unix": int(time.time()),
                    }
                )
                _atomic_json(path, asdict(updated))
                return updated
            finally:
                fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN)

    def consume(
        self,
        approval_id: str,
        *,
        session_id: str,
        call: ToolCall,
    ) -> ApprovalRecord:
        path = self._path(approval_id)
        lock_path = path.with_suffix(".lock")
        with lock_path.open("a+b") as lock_handle:
            lock_path.chmod(0o600)
            fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX)
            try:
                record = self._load_locked(path)
                expected = self.action_sha256(session_id, call)
                if not hmac.compare_digest(record.action_sha256, expected):
                    raise PermissionError("approval does not match the selected action")
                if record.status != "approved":
                    raise PermissionError(f"approval status is {record.status}")
                updated = ApprovalRecord(
                    **{
                        **asdict(record),
                        "status": "consumed",
                        "consumed_unix": int(time.time()),
                    }
                )
                _atomic_json(path, asdict(updated))
                return updated
            finally:
                fcntl.flock(lock_handle.fileno(), fcntl.LOCK_UN)

    @staticmethod
    def _session_sha256(session_id: str) -> str:
        return hashlib.sha256(session_id.encode("utf-8")).hexdigest()

    @classmethod
    def _require_session(
        cls, record: ApprovalRecord, session_id: str
    ) -> ApprovalRecord:
        if session_id and not hmac.compare_digest(
            record.session_sha256, cls._session_sha256(session_id)
        ):
            raise PermissionError("approval does not belong to this session")
        return record

    def get(self, approval_id: str, *, session_id: str = "") -> ApprovalRecord:
        return self._require_session(
            self._load_locked(self._path(approval_id)), session_id
        )

    def list(self, *, session_id: str = "") -> tuple[ApprovalRecord, ...]:
        records: list[ApprovalRecord] = []
        for path in sorted(self.root.glob("*.json")):
            try:
                record = self._load_locked(path)
            except (OSError, TypeError, ValueError, RuntimeError):
                continue
            if session_id and not hmac.compare_digest(
                record.session_sha256, self._session_sha256(session_id)
            ):
                continue
            records.append(record)
        return tuple(records)


__all__ = [
    "ApprovalRecord",
    "ApprovalStore",
    "PolicyDecision",
    "SecretRedactor",
    "ToolPolicy",
    "redact_sensitive_value",
]