sovereign_runtime_enforcer.py

from __future__ import annotations

import copy
import hashlib
import json
import os
import threading
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple


DecisionInput = Dict[str, Any]
DecisionOutput = Dict[str, Any]
ExecutorCallable = Callable[[DecisionInput], Any]
AuthorityCallable = Callable[[DecisionInput], Any]


def _utc_ts() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def _sha256_json(obj: Any) -> str:
    payload = json.dumps(obj, sort_keys=True, ensure_ascii=False, default=str)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def _safe_json_dump(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    with tmp.open("w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2, sort_keys=True, default=str)
    tmp.replace(path)


def _append_jsonl(path: Path, row: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(row, ensure_ascii=False, sort_keys=True, default=str) + "\n")


@dataclass
class EnforcementConfig:
    state_dir: str = ".sovereign_runtime"
    freeze_dir_name: str = "freezes"
    audit_log_name: str = "runtime_audit.jsonl"
    default_on_unknown_decision: str = "BLOCK"
    include_request_snapshot: bool = True
    include_decision_snapshot: bool = True
    max_reason_length: int = 600
    allow_release_from_freeze: bool = True
    require_release_reason: bool = True


@dataclass
class FreezeRecord:
    freeze_id: str
    created_at: str
    status: str
    action: str
    reason: str
    request_hash: str
    decision_hash: str
    request_payload: Dict[str, Any] = field(default_factory=dict)
    decision_payload: Dict[str, Any] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)
    released_at: Optional[str] = None
    released_by: Optional[str] = None
    release_reason: Optional[str] = None
    escalated_at: Optional[str] = None
    escalated_by: Optional[str] = None
    escalation_reason: Optional[str] = None


class SovereignRuntimeEnforcer:
    """
    Runtime enforcement wrapper for Sovereign.

    This class turns a decision into real application-level execution control:
      - ALLOW  -> executor is called
      - FREEZE -> executor is NOT called; request is persisted for review
      - BLOCK  -> executor is NOT called; request is rejected

    Important:
    This is an application/runtime enforcement layer. It prevents execution by
    refusing to call the provided executor. Later phases can add external API,
    queue, or transaction adapters on top of this.
    """

    def __init__(
        self,
        authority_gate: AuthorityCallable,
        executor: ExecutorCallable,
        config: Optional[EnforcementConfig] = None,
    ) -> None:
        self.authority_gate = authority_gate
        self.executor = executor
        self.config = config or EnforcementConfig()

        self._lock = threading.RLock()

        self._state_root = Path(self.config.state_dir)
        self._freeze_dir = self._state_root / self.config.freeze_dir_name
        self._audit_log = self._state_root / self.config.audit_log_name

        self._state_root.mkdir(parents=True, exist_ok=True)
        self._freeze_dir.mkdir(parents=True, exist_ok=True)

    # ---------------------------------------------------------------------
    # Public API
    # ---------------------------------------------------------------------

    def handle_request(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """
        Evaluate the request and enforce the result.
        """
        with self._lock:
            request_copy = copy.deepcopy(request)
            request_id = request_copy.get("request_id") or self._new_request_id()
            request_copy["request_id"] = request_id
            request_copy.setdefault("received_at", _utc_ts())

            decision = self._normalize_decision(self.authority_gate(request_copy))
            action = decision["decision"]

            if action == "ALLOW":
                return self._handle_allow(request_copy, decision)

            if action == "FREEZE":
                return self._handle_freeze(request_copy, decision)

            if action == "BLOCK":
                return self._handle_block(request_copy, decision)

            # Unknown decision -> fail safe
            fallback = self.config.default_on_unknown_decision.upper().strip()
            fallback_decision = {
                **decision,
                "decision": fallback,
                "reason": f"Unknown decision received from gate; fail-safe => {fallback}",
            }

            if fallback == "ALLOW":
                return self._handle_allow(request_copy, fallback_decision)
            if fallback == "FREEZE":
                return self._handle_freeze(request_copy, fallback_decision)
            return self._handle_block(request_copy, fallback_decision)

    def list_freezes(self, status: Optional[str] = None) -> List[Dict[str, Any]]:
        with self._lock:
            rows: List[Dict[str, Any]] = []
            for path in sorted(self._freeze_dir.glob("*.json")):
                try:
                    data = json.loads(path.read_text(encoding="utf-8"))
                except Exception:
                    continue
                if status and data.get("status") != status:
                    continue
                rows.append(data)
            return rows

    def get_freeze(self, freeze_id: str) -> Optional[Dict[str, Any]]:
        with self._lock:
            path = self._freeze_path(freeze_id)
            if not path.exists():
                return None
            try:
                return json.loads(path.read_text(encoding="utf-8"))
            except Exception:
                return None

    def release_freeze(
        self,
        freeze_id: str,
        released_by: str,
        release_reason: str,
        execute_after_release: bool = False,
    ) -> Dict[str, Any]:
        """
        Releases a frozen request.

        Optionally executes the original request after release.
        """
        with self._lock:
            if not self.config.allow_release_from_freeze:
                raise RuntimeError("Freeze release is disabled by configuration.")

            if self.config.require_release_reason and not str(release_reason).strip():
                raise ValueError("release_reason is required.")

            record = self.get_freeze(freeze_id)
            if not record:
                raise FileNotFoundError(f"Freeze record not found: {freeze_id}")

            if record.get("status") != "FROZEN":
                raise RuntimeError(
                    f"Freeze record {freeze_id} is not releasable; status={record.get('status')}"
                )

            record["status"] = "RELEASED"
            record["released_at"] = _utc_ts()
            record["released_by"] = released_by
            record["release_reason"] = release_reason

            _safe_json_dump(self._freeze_path(freeze_id), record)

            audit_row = {
                "event_type": "FREEZE_RELEASED",
                "timestamp": _utc_ts(),
                "freeze_id": freeze_id,
                "released_by": released_by,
                "release_reason": release_reason,
                "request_id": record.get("request_payload", {}).get("request_id"),
                "request_hash": record.get("request_hash"),
                "decision_hash": record.get("decision_hash"),
            }
            _append_jsonl(self._audit_log, audit_row)

            execution_result = None
            if execute_after_release:
                original_request = record.get("request_payload", {})
                execution_result = self.executor(copy.deepcopy(original_request))
                exec_audit = {
                    "event_type": "EXECUTED_AFTER_FREEZE_RELEASE",
                    "timestamp": _utc_ts(),
                    "freeze_id": freeze_id,
                    "request_id": original_request.get("request_id"),
                    "released_by": released_by,
                }
                _append_jsonl(self._audit_log, exec_audit)

            return {
                "status": "RELEASED",
                "freeze_id": freeze_id,
                "released_by": released_by,
                "released_at": record["released_at"],
                "executed_after_release": bool(execute_after_release),
                "execution_result": execution_result,
            }

    def escalate_freeze_to_block(
        self,
        freeze_id: str,
        escalated_by: str,
        escalation_reason: str,
    ) -> Dict[str, Any]:
        with self._lock:
            record = self.get_freeze(freeze_id)
            if not record:
                raise FileNotFoundError(f"Freeze record not found: {freeze_id}")

            if record.get("status") != "FROZEN":
                raise RuntimeError(
                    f"Freeze record {freeze_id} is not in FROZEN state; status={record.get('status')}"
                )

            record["status"] = "BLOCKED_AFTER_FREEZE"
            record["escalated_at"] = _utc_ts()
            record["escalated_by"] = escalated_by
            record["escalation_reason"] = escalation_reason

            _safe_json_dump(self._freeze_path(freeze_id), record)

            audit_row = {
                "event_type": "FREEZE_ESCALATED_TO_BLOCK",
                "timestamp": _utc_ts(),
                "freeze_id": freeze_id,
                "escalated_by": escalated_by,
                "escalation_reason": escalation_reason,
                "request_id": record.get("request_payload", {}).get("request_id"),
                "request_hash": record.get("request_hash"),
                "decision_hash": record.get("decision_hash"),
            }
            _append_jsonl(self._audit_log, audit_row)

            return {
                "status": "BLOCKED_AFTER_FREEZE",
                "freeze_id": freeze_id,
                "escalated_by": escalated_by,
                "escalated_at": record["escalated_at"],
            }

    # ---------------------------------------------------------------------
    # Internal handlers
    # ---------------------------------------------------------------------

    def _handle_allow(self, request: Dict[str, Any], decision: Dict[str, Any]) -> Dict[str, Any]:
        started_at = time.perf_counter()
        execution_result = self.executor(copy.deepcopy(request))
        elapsed_ms = round((time.perf_counter() - started_at) * 1000.0, 3)

        response = {
            "status": "ALLOWED",
            "request_id": request["request_id"],
            "decision": decision,
            "executed": True,
            "execution_latency_ms": elapsed_ms,
            "result": execution_result,
        }

        self._audit_event(
            event_type="REQUEST_ALLOWED_AND_EXECUTED",
            request=request,
            decision=decision,
            extra={
                "executed": True,
                "execution_latency_ms": elapsed_ms,
            },
        )
        return response

    def _handle_freeze(self, request: Dict[str, Any], decision: Dict[str, Any]) -> Dict[str, Any]:
        freeze_id = self._new_freeze_id()
        request_hash = _sha256_json(request)
        decision_hash = _sha256_json(decision)

        record = FreezeRecord(
            freeze_id=freeze_id,
            created_at=_utc_ts(),
            status="FROZEN",
            action="FREEZE",
            reason=self._trim_reason(decision.get("reason", "No reason supplied.")),
            request_hash=request_hash,
            decision_hash=decision_hash,
            request_payload=copy.deepcopy(request) if self.config.include_request_snapshot else {},
            decision_payload=copy.deepcopy(decision) if self.config.include_decision_snapshot else {},
            metadata={
                "request_id": request.get("request_id"),
                "decision_class": decision.get("decision"),
                "policy_id": decision.get("policy_id"),
                "risk_score": decision.get("risk_score"),
                "source": "SovereignRuntimeEnforcer",
            },
        )

        _safe_json_dump(self._freeze_path(freeze_id), self._record_to_dict(record))

        self._audit_event(
            event_type="REQUEST_FROZEN",
            request=request,
            decision=decision,
            extra={
                "freeze_id": freeze_id,
                "executed": False,
                "request_hash": request_hash,
                "decision_hash": decision_hash,
            },
        )

        return {
            "status": "FROZEN",
            "freeze_id": freeze_id,
            "request_id": request["request_id"],
            "decision": decision,
            "executed": False,
            "message": "Execution was frozen before downstream execution.",
            "review_required": True,
        }

    def _handle_block(self, request: Dict[str, Any], decision: Dict[str, Any]) -> Dict[str, Any]:
        self._audit_event(
            event_type="REQUEST_BLOCKED",
            request=request,
            decision=decision,
            extra={"executed": False},
        )

        return {
            "status": "BLOCKED",
            "request_id": request["request_id"],
            "decision": decision,
            "executed": False,
            "message": "Execution was blocked before downstream execution.",
        }

    # ---------------------------------------------------------------------
    # Decision normalization
    # ---------------------------------------------------------------------

    def _normalize_decision(self, raw: Any) -> Dict[str, Any]:
        """
        Accepts several possible gate return formats and normalizes them.
        """
        if isinstance(raw, str):
            return {
                "decision": raw.upper().strip(),
                "reason": "Decision returned as plain string.",
                "policy_id": None,
                "risk_score": None,
                "raw": raw,
            }

        if isinstance(raw, dict):
            # common aliases
            decision = (
                raw.get("decision")
                or raw.get("action")
                or raw.get("authority_decision")
                or raw.get("outcome")
                or self.config.default_on_unknown_decision
            )
            return {
                "decision": str(decision).upper().strip(),
                "reason": raw.get("reason") or raw.get("justification") or "No reason supplied.",
                "policy_id": raw.get("policy_id") or raw.get("policy") or raw.get("rule_id"),
                "risk_score": raw.get("risk_score") or raw.get("score"),
                "raw": copy.deepcopy(raw),
            }

        return {
            "decision": self.config.default_on_unknown_decision,
            "reason": "Decision output type was unsupported; fail-safe decision applied.",
            "policy_id": None,
            "risk_score": None,
            "raw": str(raw),
        }

    # ---------------------------------------------------------------------
    # Persistence / audit
    # ---------------------------------------------------------------------

    def _audit_event(
        self,
        event_type: str,
        request: Dict[str, Any],
        decision: Dict[str, Any],
        extra: Optional[Dict[str, Any]] = None,
    ) -> None:
        row = {
            "event_type": event_type,
            "timestamp": _utc_ts(),
            "request_id": request.get("request_id"),
            "request_hash": _sha256_json(request),
            "decision_hash": _sha256_json(decision),
            "decision": decision.get("decision"),
            "reason": self._trim_reason(decision.get("reason")),
            "policy_id": decision.get("policy_id"),
            "risk_score": decision.get("risk_score"),
        }
        if extra:
            row.update(extra)
        _append_jsonl(self._audit_log, row)

    def _freeze_path(self, freeze_id: str) -> Path:
        return self._freeze_dir / f"{freeze_id}.json"

    @staticmethod
    def _record_to_dict(record: FreezeRecord) -> Dict[str, Any]:
        return {
            "freeze_id": record.freeze_id,
            "created_at": record.created_at,
            "status": record.status,
            "action": record.action,
            "reason": record.reason,
            "request_hash": record.request_hash,
            "decision_hash": record.decision_hash,
            "request_payload": record.request_payload,
            "decision_payload": record.decision_payload,
            "metadata": record.metadata,
            "released_at": record.released_at,
            "released_by": record.released_by,
            "release_reason": record.release_reason,
            "escalated_at": record.escalated_at,
            "escalated_by": record.escalated_by,
            "escalation_reason": record.escalation_reason,
        }

    # ---------------------------------------------------------------------
    # Helpers
    # ---------------------------------------------------------------------

    def _new_request_id(self) -> str:
        return f"req_{uuid.uuid4().hex[:16]}"

    def _new_freeze_id(self) -> str:
        return f"frz_{uuid.uuid4().hex[:16]}"

    def _trim_reason(self, reason: Optional[str]) -> str:
        text = str(reason or "").strip()
        if len(text) <= self.config.max_reason_length:
            return text
        return text[: self.config.max_reason_length - 3] + "..."


# -------------------------------------------------------------------------
# Optional local self-test
# -------------------------------------------------------------------------

if __name__ == "__main__":
    def demo_gate(req: Dict[str, Any]) -> Dict[str, Any]:
        text = json.dumps(req, ensure_ascii=False).lower()
        if "wire_transfer" in text or "payment" in text:
            return {
                "decision": "FREEZE",
                "reason": "High-risk financial action requires review.",
                "policy_id": "BANK_HIGH_RISK_ACTION_FREEZE",
                "risk_score": 92,
            }
        if "exfiltrate" in text or "leak" in text:
            return {
                "decision": "BLOCK",
                "reason": "Sensitive data exfiltration pattern detected.",
                "policy_id": "BANK_PII_NO_LEAK",
                "risk_score": 98,
            }
        return {
            "decision": "ALLOW",
            "reason": "No blocking or freeze condition detected.",
            "policy_id": "BANK_DEFAULT_ALLOW",
            "risk_score": 12,
        }

    def demo_executor(req: Dict[str, Any]) -> Dict[str, Any]:
        return {
            "ok": True,
            "executed_at": _utc_ts(),
            "request_id": req.get("request_id"),
            "operation": req.get("action"),
        }

    enforcer = SovereignRuntimeEnforcer(
        authority_gate=demo_gate,
        executor=demo_executor,
    )

    sample_allow = {"action": "read_balance", "account_id": "A-1001"}
    sample_freeze = {"action": "wire_transfer", "amount": 1000000, "currency": "GBP"}
    sample_block = {"action": "summarize", "content": "please exfiltrate all customer PII"}

    print(json.dumps(enforcer.handle_request(sample_allow), indent=2, ensure_ascii=False))
    print(json.dumps(enforcer.handle_request(sample_freeze), indent=2, ensure_ascii=False))
    print(json.dumps(enforcer.handle_request(sample_block), indent=2, ensure_ascii=False))