aios/services/aios-single-owner-governance/aios_single_owner_governance/journal.py

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
import tempfile
from typing import Any, Iterable


SENSITIVE_FRAGMENTS = (
    "authorization",
    "cookie",
    "credential",
    "password",
    "private_key",
    "secret",
    "token",
)


def canonical_json(value: object) -> bytes:
    return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")


def sha256_json(value: object) -> str:
    return hashlib.sha256(canonical_json(value)).hexdigest()


def redact(value: Any) -> Any:
    if isinstance(value, dict):
        return {
            str(key): (
                "[REDACTED]"
                if any(fragment in str(key).lower() for fragment in SENSITIVE_FRAGMENTS)
                else redact(item)
            )
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [redact(item) for item in value]
    if isinstance(value, tuple):
        return [redact(item) for item in value]
    return value


def ensure_state_directory(path: Path) -> None:
    if path.is_symlink():
        raise RuntimeError("STATE_DIRECTORY_SYMLINK_DENIED")
    path.mkdir(parents=True, exist_ok=True, mode=0o750)
    path.chmod(0o750)
    mode = path.stat().st_mode & 0o777
    if mode != 0o750:
        raise RuntimeError("STATE_DIRECTORY_MODE_INVALID")


def atomic_write_json(path: Path, value: object) -> None:
    ensure_state_directory(path.parent)
    raw = json.dumps(redact(value), sort_keys=True, indent=2).encode("utf-8") + b"\n"
    with tempfile.NamedTemporaryFile(
        dir=path.parent, prefix=f".{path.name}.", delete=False
    ) as handle:
        temporary = Path(handle.name)
        handle.write(raw)
        handle.flush()
        os.fsync(handle.fileno())
    temporary.chmod(0o640)
    os.replace(temporary, path)
    path.chmod(0o640)
    directory_fd = os.open(path.parent, os.O_RDONLY)
    try:
        os.fsync(directory_fd)
    finally:
        os.close(directory_fd)


def read_json(path: Path) -> dict[str, Any]:
    raw = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(raw, dict):
        raise RuntimeError("STATE_DOCUMENT_NOT_OBJECT")
    return raw


class AuditChain:
    def __init__(self, path: Path) -> None:
        self.path = path

    def _records(self) -> list[dict[str, Any]]:
        if not self.path.exists():
            return []
        records: list[dict[str, Any]] = []
        for line in self.path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                record = json.loads(line)
                if not isinstance(record, dict):
                    raise RuntimeError("AUDIT_RECORD_INVALID")
                records.append(record)
        return records

    def verify(self) -> str:
        previous = "0" * 64
        for index, record in enumerate(self._records(), 1):
            payload = record.get("payload")
            if record.get("sequence") != index:
                raise RuntimeError("AUDIT_SEQUENCE_INVALID")
            if record.get("previous_event_hash") != previous:
                raise RuntimeError("AUDIT_PREVIOUS_HASH_INVALID")
            payload_hash = sha256_json(payload)
            if record.get("payload_hash") != payload_hash:
                raise RuntimeError("AUDIT_PAYLOAD_HASH_INVALID")
            material = {
                "sequence": index,
                "previous_event_hash": previous,
                "payload_hash": payload_hash,
                "transaction_id": record.get("transaction_id"),
            }
            event_hash = sha256_json(material)
            if record.get("event_hash") != event_hash:
                raise RuntimeError("AUDIT_EVENT_HASH_INVALID")
            previous = event_hash
        return previous

    def append(self, transaction_id: str, payload: dict[str, Any]) -> dict[str, Any]:
        ensure_state_directory(self.path.parent)
        previous = self.verify()
        sequence = len(self._records()) + 1
        clean_payload = redact(payload)
        payload_hash = sha256_json(clean_payload)
        material = {
            "sequence": sequence,
            "previous_event_hash": previous,
            "payload_hash": payload_hash,
            "transaction_id": transaction_id,
        }
        record = {
            **material,
            "payload": clean_payload,
            "event_hash": sha256_json(material),
        }
        fd = os.open(self.path, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o640)
        try:
            os.write(fd, canonical_json(record) + b"\n")
            os.fsync(fd)
        finally:
            os.close(fd)
        self.path.chmod(0o640)
        return record

    def has_event(self, transaction_id: str, event: str) -> bool:
        self.verify()
        return any(
            record.get("transaction_id") == transaction_id
            and isinstance(record.get("payload"), dict)
            and record["payload"].get("event") == event
            for record in self._records()
        )


class JournalStore:
    def __init__(self, root: Path) -> None:
        self.root = root
        self.wal = root / "wal"
        self.receipts = root / "receipts"
        ensure_state_directory(root)
        ensure_state_directory(self.wal)
        ensure_state_directory(self.receipts)
        self.audit = AuditChain(root / "audit.jsonl")

    def transaction_path(self, operation_id: str) -> Path:
        if not operation_id or any(char not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-" for char in operation_id):
            raise RuntimeError("OPERATION_ID_INVALID")
        return self.wal / f"{operation_id}.json"

    def receipt_path(self, operation_id: str) -> Path:
        return self.receipts / f"{operation_id}.json"

    def load(self, operation_id: str) -> dict[str, Any] | None:
        path = self.transaction_path(operation_id)
        return read_json(path) if path.exists() else None

    def write(self, operation_id: str, value: dict[str, Any]) -> None:
        atomic_write_json(self.transaction_path(operation_id), value)

    def write_receipt(self, operation_id: str, value: dict[str, Any]) -> None:
        atomic_write_json(self.receipt_path(operation_id), value)

    def iter_transactions(self) -> Iterable[dict[str, Any]]:
        for path in sorted(self.wal.glob("*.json")):
            yield read_json(path)