workspaceguard-cli 0.1.1__py3-none-any.whl

workspaceguard/isolation_guard.py

@@ -0,0 +1,158 @@
+"""
+Core isolation + metering logic. Ported from src/core/isolation-guard.ts.
+Never imports a specific backend directly -- all backend-specific behavior
+goes through the BackendAdapter interface, a fixed architectural boundary.
+The CLI and the library export are both thin wrappers around this class;
+there is exactly one implementation of the isolation/metering logic.
+"""
+from __future__ import annotations
+
+import asyncio
+import math
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import DefaultDict, List, Optional
+
+from .circuit_breaker import CircuitBreaker
+from .config import load_config, resolve_workspace_id, save_config, set_workspace_cap, upsert_workspace
+from .log import ConsoleLogger, FailClosedEvent, Logger
+from .namespace import ensure_workspace_dirs
+from .types import BackendAdapter, IdentityNotFoundError, WorkspaceGuardConfig
+from .usage import UsageMeter, WorkspaceUsage
+from .vault import Vault
+
+
+@dataclass
+class WorkspaceUsageReport:
+    workspace_id: str
+    identity: str
+    monthly_message_cap: Optional[int]
+    percent_used: Optional[int]
+    period: str
+    message_count: int
+    estimated_bytes: int
+
+
+class IsolationGuard:
+    def __init__(
+        self,
+        data_dir: str,
+        backend: BackendAdapter,
+        logger: Optional[Logger] = None,
+        circuit_open_cooldown_ms: Optional[int] = None,
+    ) -> None:
+        self._data_dir = data_dir
+        self._backend = backend
+        self._logger = logger or ConsoleLogger()
+        self._vault = Vault(os.path.join(self._data_dir, ".workspaceguard", "master.key"))
+        kwargs = {} if circuit_open_cooldown_ms is None else {"open_cooldown_ms": circuit_open_cooldown_ms}
+        self._circuit = CircuitBreaker(backend.name, self._logger, **kwargs)
+        self._usage = UsageMeter(self._data_dir)
+        self._config = WorkspaceGuardConfig(backend="odysseus", identity_header="", workspaces=[])
+        # Per-workspace lock so a concurrent quota check-then-record can't
+        # race across two calls to chat() for the same workspace (in-process
+        # only, matching this project's documented single-sidecar-process
+        # deployment model -- see usage.py).
+        self._workspace_locks: DefaultDict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+
+    @property
+    def _config_path(self) -> str:
+        return os.path.join(self._data_dir, "workspaceguard.config.yaml")
+
+    async def init(self) -> None:
+        await self._vault.init()
+        self._config = load_config(self._config_path)
+
+    async def add_workspace(self, workspace_id: str, identity: str) -> None:
+        self._config = upsert_workspace(self._config, workspace_id, identity)
+        save_config(self._config_path, self._config)
+        await ensure_workspace_dirs(self._data_dir, workspace_id)
+
+    async def set_secret(self, workspace_id: str, secret: str) -> None:
+        await self._vault.write_secret(self._data_dir, workspace_id, secret)
+
+    async def rotate_key(self, workspace_id: str) -> None:
+        await self._vault.rotate(self._data_dir, workspace_id)
+
+    def resolve_workspace(self, identity_header_value: Optional[str]) -> str:
+        """
+        Resolves a workspace from an untrusted identity header value. Fails
+        closed on any miss -- never falls back to a default workspace. This
+        is the single choke point every request-handling path must go
+        through.
+        """
+        workspace_id = resolve_workspace_id(self._config, identity_header_value)
+        if not workspace_id:
+            self._logger.log(
+                FailClosedEvent(
+                    reason="identity_not_found",
+                    detail={"identityHeaderValue": identity_header_value},
+                )
+            )
+            raise IdentityNotFoundError(identity_header_value)
+        return workspace_id
+
+    async def chat(self, identity_header_value: Optional[str], message: str) -> str:
+        """
+        The quota check, backend call, and usage record are wrapped in a
+        single per-workspace lock so two concurrent chat() calls for the
+        same workspace can never both pass the same pre-cap quota check
+        (regression: previously unlocked, letting N concurrent requests at
+        cap-1 all read the same count and all proceed, and letting
+        concurrent record() calls lose updates to each other). Different
+        workspaces never contend with each other.
+        """
+        workspace_id = self.resolve_workspace(identity_header_value)
+        entry = next((w for w in self._config.workspaces if w.workspace_id == workspace_id), None)
+        cap = entry.monthly_message_cap if entry is not None else None
+
+        async with self._workspace_locks[workspace_id]:
+            await self._usage.check_quota(workspace_id, cap)
+
+            async def _call() -> str:
+                return await self._backend.forward_chat(workspace_id, message)
+
+            response = await self._circuit.call(_call)
+            await self._usage.record(workspace_id, message)
+            return response
+
+    async def status(self) -> List[dict]:
+        return [{"workspaceId": w.workspace_id, "identity": w.identity} for w in self._config.workspaces]
+
+    async def set_cap(self, workspace_id: str, cap: Optional[int]) -> None:
+        self._config = set_workspace_cap(self._config, workspace_id, cap)
+        save_config(self._config_path, self._config)
+
+    async def usage_report(self) -> List[WorkspaceUsageReport]:
+        """
+        The admin-visibility surface this project actually ships (the OSS
+        free tier) -- a hosted billing dashboard on top of this data is a
+        separate, closed-source product, never built in this repo.
+        """
+        ids = [w.workspace_id for w in self._config.workspaces]
+        usage_by_id = await self._usage.get_all(ids)
+        reports: List[WorkspaceUsageReport] = []
+        for w in self._config.workspaces:
+            usage = usage_by_id.get(w.workspace_id) or WorkspaceUsage(period="", message_count=0, estimated_bytes=0)
+            percent_used: Optional[int] = None
+            if w.monthly_message_cap is not None and w.monthly_message_cap > 0:
+                # math.floor(x + 0.5), not the builtin round(): Python's
+                # round() uses banker's rounding (round-half-to-even) while
+                # the TS original's Math.round() rounds half-up -- this
+                # kept percentUsed diverging between the two --json outputs
+                # for exact-half cases (e.g. 12.5% -> 12 vs 13) despite both
+                # READMEs documenting the outputs as identical.
+                percent_used = math.floor((usage.message_count / w.monthly_message_cap) * 100 + 0.5)
+            reports.append(
+                WorkspaceUsageReport(
+                    workspace_id=w.workspace_id,
+                    identity=w.identity,
+                    monthly_message_cap=w.monthly_message_cap,
+                    percent_used=percent_used,
+                    period=usage.period,
+                    message_count=usage.message_count,
+                    estimated_bytes=usage.estimated_bytes,
+                )
+            )
+        return reports

workspaceguard/log.py

@@ -0,0 +1,43 @@
+"""
+Structured event logging. Ported from src/core/log.ts. Every fail-closed
+decision and every circuit-breaker state change gets one JSON line -- the
+only way the "zero cross-workspace leaks" claim is verifiable outside of
+tests.
+"""
+from __future__ import annotations
+
+import json
+import sys
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, Dict, Literal, Optional
+
+
+@dataclass
+class FailClosedEvent:
+    reason: str
+    detail: Dict[str, Any]
+    type: Literal["fail_closed"] = "fail_closed"
+
+
+@dataclass
+class CircuitStateChangeEvent:
+    backend: str
+    state: Literal["open", "closed"]
+    type: Literal["circuit_state_change"] = "circuit_state_change"
+
+
+LogEvent = Any  # FailClosedEvent | CircuitStateChangeEvent
+
+
+class Logger(ABC):
+    @abstractmethod
+    def log(self, event: LogEvent) -> None: ...
+
+
+class ConsoleLogger(Logger):
+    def log(self, event: LogEvent) -> None:
+        line: Dict[str, Any] = {"ts": datetime.now(timezone.utc).isoformat()}
+        line.update(vars(event))
+        print(json.dumps(line), file=sys.stdout)

workspaceguard/namespace.py

@@ -0,0 +1,21 @@
+"""
+Chat history and memory get a directory boundary per workspace, not a
+shared table/file with a workspace column -- easier to verify by inspection
+than a query filter. Ported from src/core/namespace.ts.
+"""
+from __future__ import annotations
+
+import os
+
+
+def chat_history_dir(data_dir: str, workspace_id: str) -> str:
+    return os.path.join(data_dir, "workspaces", workspace_id, "chat")
+
+
+def memory_dir(data_dir: str, workspace_id: str) -> str:
+    return os.path.join(data_dir, "workspaces", workspace_id, "memory")
+
+
+async def ensure_workspace_dirs(data_dir: str, workspace_id: str) -> None:
+    os.makedirs(chat_history_dir(data_dir, workspace_id), exist_ok=True)
+    os.makedirs(memory_dir(data_dir, workspace_id), exist_ok=True)

workspaceguard/types.py

@@ -0,0 +1,78 @@
+"""
+Shared types for WorkspaceGuard's isolation and metering pipeline. Used by
+both the CLI (workspaceguard/cli.py) and the programmatic library API
+(workspaceguard/__init__.py).
+
+This is a Python port of the TypeScript original's src/core/types.ts. Field
+names follow Python (snake_case) convention; the config/usage persistence
+layer (workspaceguard/config.py, workspaceguard/usage.py) re-serializes to
+the same camelCase keys the npm package's YAML/JSON files use, so a
+`workspaceguard.config.yaml` or `usage.json` written by one distribution
+stays readable by the other.
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+
+@dataclass
+class WorkspaceEntry:
+    workspace_id: str
+    identity: str
+    monthly_message_cap: Optional[int] = None
+    """Monthly message cap for usage metering. None = unlimited (free-tier default)."""
+
+
+@dataclass
+class WorkspaceGuardConfig:
+    backend: str = "odysseus"
+    identity_header: str = ""
+    workspaces: List[WorkspaceEntry] = field(default_factory=list)
+
+
+class BackendAdapter(ABC):
+    """
+    Fixed architectural boundary: all backend-specific behavior goes through
+    this interface. Core isolation/metering logic never imports a specific
+    backend directly.
+    """
+
+    name: str
+
+    @abstractmethod
+    async def health_check(self) -> bool: ...
+
+    @abstractmethod
+    async def forward_chat(self, workspace_id: str, message: str) -> str: ...
+
+
+class IdentityNotFoundError(Exception):
+    def __init__(self, header_value: Optional[str]) -> None:
+        super().__init__("no workspace configured for this identity")
+        self.header_value = header_value
+
+
+class VaultDecryptionError(Exception):
+    def __init__(self, workspace_id: str) -> None:
+        super().__init__(f"could not decrypt vault for workspace {workspace_id}")
+        self.workspace_id = workspace_id
+
+
+class BackendUnreachableError(Exception):
+    def __init__(self, backend: str) -> None:
+        super().__init__(f"backend {backend} did not respond")
+        self.backend = backend
+
+
+class BackendCircuitOpenError(Exception):
+    def __init__(self, backend: str) -> None:
+        super().__init__(f"backend {backend} circuit is open, refusing calls")
+        self.backend = backend
+
+
+class WorkspaceNotFoundError(Exception):
+    def __init__(self, workspace_id: str) -> None:
+        super().__init__(f'no workspace configured with id "{workspace_id}"')
+        self.workspace_id = workspace_id

workspaceguard/usage.py

@@ -0,0 +1,149 @@
+"""
+Per-workspace usage metering -- the wedge this project ships instead of
+per-user isolation (already native in the target platform, see README).
+Ported from src/core/usage.ts. Reads/writes are not lock-guarded against
+concurrent processes, the same accepted tradeoff config.py already makes for
+`workspaceguard.config.yaml` (single-sidecar-process deployment model).
+
+The usage.json file uses camelCase keys (period, messageCount,
+estimatedBytes) to stay wire-compatible with the npm package's own store.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Dict, List, Optional
+
+UsageStore = Dict[str, "WorkspaceUsage"]
+
+
+@dataclass
+class WorkspaceUsage:
+    period: str
+    message_count: int
+    estimated_bytes: int
+
+
+class QuotaExceededError(Exception):
+    def __init__(self, workspace_id: str, cap: int) -> None:
+        super().__init__(f"workspace {workspace_id} has reached its monthly cap of {cap} messages")
+        self.workspace_id = workspace_id
+        self.cap = cap
+
+
+def _current_period() -> str:
+    now = datetime.now(timezone.utc)
+    return f"{now.year}-{now.month:02d}"
+
+
+def usage_path(data_dir: str) -> str:
+    return os.path.join(data_dir, ".workspaceguard", "usage.json")
+
+
+def _load_usage_raw(data_dir: str) -> Dict[str, dict]:
+    """
+    Only a missing file (no workspace has recorded usage yet) is a
+    legitimate empty store. Every other failure (permission error, disk
+    error, a corrupted/truncated JSON file) must propagate so a caller like
+    check_quota() fails closed instead of silently reading every
+    workspace's usage as zero -- a catch-everything here used to let any
+    I/O hiccup or a race-corrupted usage.json lift every quota in the
+    store.
+    """
+    try:
+        with open(usage_path(data_dir), "r", encoding="utf-8") as fh:
+            raw = fh.read()
+    except FileNotFoundError:
+        return {}
+
+    try:
+        parsed = json.loads(raw)
+    except json.JSONDecodeError as err:
+        raise ValueError(f"usage store at {usage_path(data_dir)} is corrupted and could not be parsed: {err}") from err
+    if not isinstance(parsed, dict):
+        raise ValueError(f"usage store at {usage_path(data_dir)} has an unexpected shape (expected a JSON object)")
+    return parsed
+
+
+def load_usage(data_dir: str) -> UsageStore:
+    raw = _load_usage_raw(data_dir)
+    store: UsageStore = {}
+    for workspace_id, entry in raw.items():
+        try:
+            store[workspace_id] = WorkspaceUsage(
+                period=entry["period"],
+                message_count=entry["messageCount"],
+                estimated_bytes=entry["estimatedBytes"],
+            )
+        except (KeyError, TypeError):
+            continue
+    return store
+
+
+def save_usage(data_dir: str, store: UsageStore) -> None:
+    path = usage_path(data_dir)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    payload = {
+        workspace_id: {
+            "period": usage.period,
+            "messageCount": usage.message_count,
+            "estimatedBytes": usage.estimated_bytes,
+        }
+        for workspace_id, usage in store.items()
+    }
+    with open(path, "w", encoding="utf-8") as fh:
+        json.dump(payload, fh, indent=2)
+
+
+def current_usage_for(store: UsageStore, workspace_id: str) -> WorkspaceUsage:
+    """
+    Rolls a workspace's counter over to a fresh period on read -- a stale
+    count from last month must never inflate this month's total or trip a
+    cap that shouldn't apply yet (no cron/scheduler needed, since every real
+    read/write already knows what period it is).
+    """
+    existing = store.get(workspace_id)
+    period = _current_period()
+    if existing is None or existing.period != period:
+        return WorkspaceUsage(period=period, message_count=0, estimated_bytes=0)
+    return existing
+
+
+class UsageMeter:
+    def __init__(self, data_dir: str) -> None:
+        self.data_dir = data_dir
+
+    async def record(self, workspace_id: str, message: str) -> WorkspaceUsage:
+        store = load_usage(self.data_dir)
+        usage = current_usage_for(store, workspace_id)
+        usage = WorkspaceUsage(
+            period=usage.period,
+            message_count=usage.message_count + 1,
+            estimated_bytes=usage.estimated_bytes + len(message.encode("utf-8")),
+        )
+        store[workspace_id] = usage
+        save_usage(self.data_dir, store)
+        return usage
+
+    async def get(self, workspace_id: str) -> WorkspaceUsage:
+        store = load_usage(self.data_dir)
+        return current_usage_for(store, workspace_id)
+
+    async def get_all(self, workspace_ids: List[str]) -> Dict[str, WorkspaceUsage]:
+        store = load_usage(self.data_dir)
+        return {workspace_id: current_usage_for(store, workspace_id) for workspace_id in workspace_ids}
+
+    async def check_quota(self, workspace_id: str, cap: Optional[int]) -> None:
+        """
+        Raises QuotaExceededError once a workspace has reached its
+        configured cap. cap=None means unlimited -- never blocks a
+        workspace that has no cap set, so this stays free/OSS-safe by
+        default.
+        """
+        if cap is None:
+            return
+        usage = await self.get(workspace_id)
+        if usage.message_count >= cap:
+            raise QuotaExceededError(workspace_id, cap)

workspaceguard/vault.py

@@ -0,0 +1,130 @@
+"""
+One derived AES-256-GCM key per workspace: chat history and memory get a
+directory boundary (namespace.py), the API-key vault gets its own encrypted
+file per workspace so a leaked file alone (without the master key) reveals
+nothing. Ported from src/core/vault.ts, using the `cryptography` package's
+AESGCM primitive (the Python ecosystem's standard, actively maintained
+choice for authenticated symmetric encryption) in place of Node's
+node:crypto.
+
+On-disk layout is kept identical to the TypeScript original (12-byte IV +
+16-byte GCM auth tag + ciphertext, concatenated) so a vault file's format is
+conceptually portable between the two distributions even though a real
+deployment would run one or the other against a given data directory, not
+both against the same one simultaneously.
+"""
+from __future__ import annotations
+
+import hashlib
+import hmac
+import os
+import secrets
+from typing import Optional
+
+from cryptography.exceptions import InvalidTag
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+from .types import VaultDecryptionError
+
+KEY_BYTES = 32
+IV_BYTES = 12
+TAG_BYTES = 16
+
+
+class Vault:
+    def __init__(self, master_key_path: str) -> None:
+        self._master_key_path = master_key_path
+        self._master_key: Optional[bytes] = None
+
+    async def init(self) -> None:
+        try:
+            with open(self._master_key_path, "r", encoding="utf-8") as fh:
+                raw = fh.read().strip()
+            import base64
+
+            self._master_key = base64.b64decode(raw)
+        except FileNotFoundError:
+            import base64
+
+            self._master_key = secrets.token_bytes(KEY_BYTES)
+            os.makedirs(os.path.dirname(self._master_key_path), exist_ok=True)
+            with open(self._master_key_path, "w", encoding="utf-8") as fh:
+                fh.write(base64.b64encode(self._master_key).decode("ascii"))
+            os.chmod(self._master_key_path, 0o600)
+
+    def _require_master_key(self) -> bytes:
+        if self._master_key is None:
+            raise RuntimeError("Vault.init() must run before any vault operation")
+        return self._master_key
+
+    def _derive_workspace_key(self, workspace_id: str, generation: int) -> bytes:
+        master = self._require_master_key()
+        message = f"{workspace_id}:{generation}".encode("utf-8")
+        return hmac.new(master, message, hashlib.sha256).digest()
+
+    def vault_path(self, data_dir: str, workspace_id: str) -> str:
+        return os.path.join(data_dir, "workspaces", workspace_id, "vault.bin")
+
+    def _generation_path(self, data_dir: str, workspace_id: str) -> str:
+        return os.path.join(data_dir, "workspaces", workspace_id, "key-generation")
+
+    def _read_generation(self, data_dir: str, workspace_id: str) -> int:
+        try:
+            with open(self._generation_path(data_dir, workspace_id), "r", encoding="utf-8") as fh:
+                return int(fh.read().strip() or "0")
+        except (FileNotFoundError, ValueError):
+            return 0
+
+    async def write_secret(self, data_dir: str, workspace_id: str, plaintext: str) -> None:
+        generation = self._read_generation(data_dir, workspace_id)
+        key = self._derive_workspace_key(workspace_id, generation)
+        iv = secrets.token_bytes(IV_BYTES)
+        aesgcm = AESGCM(key)
+        sealed = aesgcm.encrypt(iv, plaintext.encode("utf-8"), None)
+        ciphertext, auth_tag = sealed[:-TAG_BYTES], sealed[-TAG_BYTES:]
+        path = self.vault_path(data_dir, workspace_id)
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        with open(path, "wb") as fh:
+            fh.write(iv + auth_tag + ciphertext)
+        os.chmod(path, 0o600)
+
+    async def read_secret(self, data_dir: str, workspace_id: str) -> str:
+        generation = self._read_generation(data_dir, workspace_id)
+        key = self._derive_workspace_key(workspace_id, generation)
+        path = self.vault_path(data_dir, workspace_id)
+        try:
+            with open(path, "rb") as fh:
+                payload = fh.read()
+        except FileNotFoundError:
+            raise VaultDecryptionError(workspace_id)
+        iv = payload[:IV_BYTES]
+        auth_tag = payload[IV_BYTES : IV_BYTES + TAG_BYTES]
+        ciphertext = payload[IV_BYTES + TAG_BYTES :]
+        try:
+            aesgcm = AESGCM(key)
+            plaintext = aesgcm.decrypt(iv, ciphertext + auth_tag, None)
+            return plaintext.decode("utf-8")
+        except (InvalidTag, ValueError):
+            raise VaultDecryptionError(workspace_id)
+
+    async def rotate(self, data_dir: str, workspace_id: str) -> None:
+        """
+        Rotation increments the per-workspace key generation BEFORE
+        re-encrypting, so the new ciphertext is under a genuinely different
+        derived key -- anything encrypted under the old generation can no
+        longer be decrypted once the generation file is bumped. Behavior for
+        an in-flight stream during rotation is a known open question, not
+        resolved here (matches the TypeScript original).
+        """
+        try:
+            existing: Optional[str] = await self.read_secret(data_dir, workspace_id)
+        except VaultDecryptionError:
+            existing = None
+        current_generation = self._read_generation(data_dir, workspace_id)
+        next_generation = current_generation + 1
+        generation_path = self._generation_path(data_dir, workspace_id)
+        os.makedirs(os.path.dirname(generation_path), exist_ok=True)
+        with open(generation_path, "w", encoding="utf-8") as fh:
+            fh.write(str(next_generation))
+        if existing is not None:
+            await self.write_secret(data_dir, workspace_id, existing)