shellbrain 0.1.0__py3-none-any.whl

app/periphery/admin/instance_guard.py

@@ -0,0 +1,211 @@
+"""Guards that classify DB instances and refuse destructive actions on live targets."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+import hashlib
+from typing import Final
+from urllib.parse import urlparse
+
+import psycopg
+
+
+LIVE: Final[str] = "live"
+TEST: Final[str] = "test"
+SCRATCH: Final[str] = "scratch"
+ALLOWED_INSTANCE_MODES: Final[set[str]] = {LIVE, TEST, SCRATCH}
+PROTECTED_DB_NAMES: Final[set[str]] = {"shellbrain", "memory"}
+
+
+@dataclass(frozen=True)
+class InstanceMetadataRecord:
+    """Minimal typed view of one instance_metadata row."""
+
+    instance_id: str
+    instance_mode: str
+    created_at: str
+    created_by: str
+    notes: str | None = None
+
+
+def normalize_dsn(dsn: str) -> str:
+    """Normalize a DSN into a stable fingerprint source string."""
+
+    parsed = urlparse(dsn.replace("+psycopg", ""))
+    hostname = (parsed.hostname or "").lower()
+    port = parsed.port or 5432
+    db_name = parsed.path.lstrip("/")
+    return f"{hostname}:{port}/{db_name}"
+
+
+def dsn_fingerprint(dsn: str) -> str:
+    """Hash a normalized DSN to compare protected targets without exposing passwords."""
+
+    return hashlib.sha256(normalize_dsn(dsn).encode("utf-8")).hexdigest()
+
+
+def database_name_from_dsn(dsn: str) -> str:
+    """Extract the target database name from one DSN."""
+
+    return urlparse(dsn.replace("+psycopg", "")).path.lstrip("/")
+
+
+def assert_disposable_test_dsn(*, test_dsn: str, protected_dsn: str | None = None) -> None:
+    """Refuse to treat a protected or production-shaped database as disposable."""
+
+    if protected_dsn and dsn_fingerprint(test_dsn) == dsn_fingerprint(protected_dsn):
+        raise RuntimeError("Refusing destructive test setup against the protected live database DSN.")
+    db_name = database_name_from_dsn(test_dsn).lower()
+    if db_name in PROTECTED_DB_NAMES:
+        raise RuntimeError(
+            f"Refusing destructive test setup against protected database '{db_name}'. "
+            "Use an explicitly disposable scratch/test database name."
+        )
+
+
+def ensure_instance_metadata(
+    dsn: str,
+    *,
+    instance_mode: str,
+    created_by: str,
+    notes: str | None = None,
+) -> InstanceMetadataRecord:
+    """Create or preserve one instance_metadata row for the target database."""
+
+    if instance_mode not in ALLOWED_INSTANCE_MODES:
+        raise ValueError(f"Unsupported instance mode: {instance_mode}")
+    raw_dsn = dsn.replace("+psycopg", "")
+    instance_id = dsn_fingerprint(dsn)
+    created_at = datetime.now(timezone.utc).isoformat()
+    with psycopg.connect(raw_dsn) as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                CREATE TABLE IF NOT EXISTS instance_metadata (
+                  instance_id TEXT PRIMARY KEY,
+                  instance_mode TEXT NOT NULL,
+                  created_at TIMESTAMPTZ NOT NULL,
+                  created_by TEXT NOT NULL,
+                  notes TEXT NULL
+                )
+                """
+            )
+            cur.execute(
+                """
+                INSERT INTO instance_metadata (instance_id, instance_mode, created_at, created_by, notes)
+                VALUES (%s, %s, %s, %s, %s)
+                ON CONFLICT (instance_id) DO NOTHING
+                """,
+                (instance_id, instance_mode, created_at, created_by, notes),
+            )
+        conn.commit()
+    record = fetch_instance_metadata(dsn, include_legacy_lookup=False)
+    if record is None:
+        raise RuntimeError("Failed to persist instance metadata.")
+    return record
+
+
+def fetch_instance_metadata(dsn: str, *, include_legacy_lookup: bool = True) -> InstanceMetadataRecord | None:
+    """Read one instance_metadata row when the safety table exists."""
+
+    raw_dsn = dsn.replace("+psycopg", "")
+    instance_ids = [dsn_fingerprint(dsn)]
+    if include_legacy_lookup:
+        instance_ids.append(_legacy_dsn_fingerprint(dsn))
+    try:
+        with psycopg.connect(raw_dsn) as conn:
+            with conn.cursor() as cur:
+                cur.execute("SELECT to_regclass('public.instance_metadata');")
+                if cur.fetchone()[0] is None:
+                    return None
+                row = None
+                for instance_id in instance_ids:
+                    cur.execute(
+                        """
+                        SELECT instance_id, instance_mode, created_at::text, created_by, notes
+                        FROM instance_metadata
+                        WHERE instance_id = %s
+                        """,
+                        (instance_id,),
+                    )
+                    row = cur.fetchone()
+                    if row is not None:
+                        break
+    except psycopg.Error:
+        return None
+    if row is None:
+        return None
+    return InstanceMetadataRecord(*row)
+
+
+def assert_destructive_allowed(dsn: str, *, allowed_modes: tuple[str, ...] = (TEST, SCRATCH)) -> None:
+    """Fail closed unless the target database is explicitly classified as disposable."""
+
+    record = fetch_instance_metadata(dsn)
+    if record is None:
+        raise RuntimeError(
+            "Refusing destructive action because instance metadata is missing. "
+            "Stamp the database as test or scratch first."
+        )
+    if record.instance_mode not in allowed_modes:
+        raise RuntimeError(
+            f"Refusing destructive action against instance_mode={record.instance_mode!r}. "
+            f"Allowed modes: {', '.join(allowed_modes)}."
+        )
+
+
+def inspect_role_safety(dsn: str) -> list[str]:
+    """Return warnings when the provided DSN points at a dangerously privileged role."""
+
+    warnings: list[str] = []
+    raw_dsn = dsn.replace("+psycopg", "")
+    try:
+        with psycopg.connect(raw_dsn) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    SELECT r.rolsuper,
+                           pg_has_role(current_user, 'pg_database_owner', 'member'),
+                           has_schema_privilege(current_user, 'public', 'CREATE')
+                    FROM pg_roles r
+                    WHERE r.rolname = current_user
+                    """
+                )
+                row = cur.fetchone()
+    except psycopg.Error as exc:
+        return [f"Could not audit DB role safety: {exc}"]
+    if row is None:
+        return warnings
+    is_superuser, is_db_owner, can_create_in_public = row
+    if is_superuser:
+        warnings.append("Current DSN role is superuser-capable.")
+    if is_db_owner:
+        warnings.append("Current DSN role is a database owner.")
+    if can_create_in_public:
+        warnings.append("Current DSN role can CREATE in schema public.")
+    return warnings
+
+
+def fingerprint_summary(dsn: str) -> dict[str, str]:
+    """Return one printable DB identity summary for diagnostics and manifests."""
+
+    parsed = urlparse(dsn.replace("+psycopg", ""))
+    return {
+        "fingerprint": dsn_fingerprint(dsn),
+        "host": parsed.hostname or "",
+        "port": str(parsed.port or 5432),
+        "database": parsed.path.lstrip("/"),
+        "user": parsed.username or "",
+    }
+
+
+def _legacy_dsn_fingerprint(dsn: str) -> str:
+    """Compute the legacy fingerprint that also included the username."""
+
+    parsed = urlparse(dsn.replace("+psycopg", ""))
+    hostname = (parsed.hostname or "").lower()
+    port = parsed.port or 5432
+    db_name = parsed.path.lstrip("/")
+    user = parsed.username or ""
+    return hashlib.sha256(f"{hostname}:{port}/{db_name}?user={user}".encode("utf-8")).hexdigest()

app/periphery/admin/machine_state.py

@@ -0,0 +1,354 @@
+"""Machine-owned Shellbrain runtime config and bootstrap state."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from datetime import datetime, timezone
+import json
+import os
+from pathlib import Path
+import tomllib
+from typing import Any
+
+from app.boot.home import (
+    get_machine_backups_dir,
+    get_machine_config_path,
+    get_machine_models_dir,
+    get_shellbrain_home,
+)
+
+
+CONFIG_VERSION = 1
+BOOTSTRAP_VERSION = 1
+RUNTIME_MODE_MANAGED_LOCAL = "managed_local"
+BOOTSTRAP_STATE_PROVISIONING = "provisioning"
+BOOTSTRAP_STATE_READY = "ready"
+BOOTSTRAP_STATE_REPAIR_NEEDED = "repair_needed"
+
+
+@dataclass(frozen=True)
+class DatabaseState:
+    """Resolved application and admin DSNs for the active runtime."""
+
+    app_dsn: str
+    admin_dsn: str
+
+
+@dataclass(frozen=True)
+class ManagedInstanceState:
+    """Machine-owned managed Postgres runtime metadata."""
+
+    instance_id: str
+    container_name: str
+    image: str
+    host: str
+    port: int
+    db_name: str
+    data_dir: str
+    admin_user: str
+    admin_password: str
+    app_user: str
+    app_password: str
+
+
+@dataclass(frozen=True)
+class BackupState:
+    """Configured backup directory roots."""
+
+    root: str
+    mirror_root: str | None = None
+
+
+@dataclass(frozen=True)
+class EmbeddingRuntimeState:
+    """Pinned embedding runtime metadata."""
+
+    provider: str
+    model: str
+    model_revision: str | None
+    backend_version: str | None
+    cache_path: str
+    readiness_state: str
+    last_error: str | None = None
+
+
+@dataclass(frozen=True)
+class MachineConfig:
+    """Machine-owned runtime state for Shellbrain bootstrap and repair."""
+
+    config_version: int
+    bootstrap_version: int
+    runtime_mode: str
+    bootstrap_state: str
+    current_step: str | None
+    last_error: str | None
+    database: DatabaseState
+    managed: ManagedInstanceState
+    backups: BackupState
+    embeddings: EmbeddingRuntimeState
+
+    @property
+    def machine_instance_id(self) -> str:
+        """Return the active machine instance identifier."""
+
+        return self.managed.instance_id
+
+
+def default_paths() -> dict[str, str]:
+    """Return default machine-owned storage paths."""
+
+    return {
+        "backups_root": str(get_machine_backups_dir()),
+        "models_root": str(get_machine_models_dir()),
+    }
+
+
+def load_machine_config(path: Path | None = None) -> MachineConfig | None:
+    """Return the parsed machine config when it exists and is valid."""
+
+    target = path or get_machine_config_path()
+    try:
+        payload = tomllib.loads(target.read_text(encoding="utf-8"))
+    except FileNotFoundError:
+        return None
+    return _machine_config_from_payload(payload)
+
+
+def try_load_machine_config(path: Path | None = None) -> tuple[MachineConfig | None, str | None]:
+    """Best-effort machine-config load that reports corruption text instead of raising."""
+
+    target = path or get_machine_config_path()
+    try:
+        return load_machine_config(target), None
+    except (tomllib.TOMLDecodeError, ValueError) as exc:
+        return None, str(exc)
+
+
+def save_machine_config(config: MachineConfig, path: Path | None = None) -> Path:
+    """Persist one machine config with owner-only permissions."""
+
+    target = path or get_machine_config_path()
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(_render_machine_config(config), encoding="utf-8")
+    try:
+        os.chmod(target, 0o600)
+    except OSError:
+        pass
+    return target
+
+
+def backup_corrupt_machine_config(*, path: Path | None = None) -> Path | None:
+    """Rename a corrupt config aside and return the backup path when present."""
+
+    target = path or get_machine_config_path()
+    if not target.exists():
+        return None
+    backup = target.with_name(f"config.corrupt.{_timestamp_slug()}.toml")
+    target.replace(backup)
+    return backup
+
+
+def update_bootstrap_state(
+    config: MachineConfig,
+    *,
+    bootstrap_state: str,
+    current_step: str | None,
+    last_error: str | None,
+) -> MachineConfig:
+    """Return a copy with updated bootstrap status fields."""
+
+    return replace(
+        config,
+        bootstrap_state=bootstrap_state,
+        current_step=current_step,
+        last_error=last_error,
+    )
+
+
+def build_recovery_stub(*, current_step: str | None, last_error: str | None) -> dict[str, Any]:
+    """Return a minimal recovery payload for config corruption cases."""
+
+    return {
+        "config_version": CONFIG_VERSION,
+        "bootstrap_version": BOOTSTRAP_VERSION,
+        "runtime_mode": RUNTIME_MODE_MANAGED_LOCAL,
+        "bootstrap_state": BOOTSTRAP_STATE_REPAIR_NEEDED,
+        "current_step": current_step,
+        "last_error": last_error,
+        "paths": default_paths(),
+    }
+
+
+def save_recovery_stub(*, current_step: str | None, last_error: str | None, path: Path | None = None) -> Path:
+    """Persist a minimal repair-needed stub after config corruption."""
+
+    target = path or get_machine_config_path()
+    target.parent.mkdir(parents=True, exist_ok=True)
+    payload = build_recovery_stub(current_step=current_step, last_error=last_error)
+    target.write_text(_render_simple_payload(payload), encoding="utf-8")
+    try:
+        os.chmod(target, 0o600)
+    except OSError:
+        pass
+    return target
+
+
+def _machine_config_from_payload(payload: dict[str, Any]) -> MachineConfig:
+    """Validate and coerce a machine-config payload."""
+
+    if not isinstance(payload, dict):
+        raise ValueError("Machine config must be a TOML table.")
+    database = payload.get("database")
+    managed = payload.get("managed")
+    backups = payload.get("backups")
+    embeddings = payload.get("embeddings")
+    if not isinstance(database, dict) or not isinstance(managed, dict):
+        raise ValueError("Machine config is missing required database or managed sections.")
+    if not isinstance(backups, dict) or not isinstance(embeddings, dict):
+        raise ValueError("Machine config is missing required backups or embeddings sections.")
+    return MachineConfig(
+        config_version=int(payload.get("config_version") or 0),
+        bootstrap_version=int(payload.get("bootstrap_version") or 0),
+        runtime_mode=str(payload.get("runtime_mode") or ""),
+        bootstrap_state=str(payload.get("bootstrap_state") or ""),
+        current_step=_optional_str(payload.get("current_step")),
+        last_error=_optional_str(payload.get("last_error")),
+        database=DatabaseState(
+            app_dsn=_required_str(database, "app_dsn"),
+            admin_dsn=_required_str(database, "admin_dsn"),
+        ),
+        managed=ManagedInstanceState(
+            instance_id=_required_str(managed, "instance_id"),
+            container_name=_required_str(managed, "container_name"),
+            image=_required_str(managed, "image"),
+            host=_required_str(managed, "host"),
+            port=int(managed.get("port") or 0),
+            db_name=_required_str(managed, "db_name"),
+            data_dir=_required_str(managed, "data_dir"),
+            admin_user=_required_str(managed, "admin_user"),
+            admin_password=_required_str(managed, "admin_password"),
+            app_user=_required_str(managed, "app_user"),
+            app_password=_required_str(managed, "app_password"),
+        ),
+        backups=BackupState(
+            root=_required_str(backups, "root"),
+            mirror_root=_optional_str(backups.get("mirror_root")),
+        ),
+        embeddings=EmbeddingRuntimeState(
+            provider=_required_str(embeddings, "provider"),
+            model=_required_str(embeddings, "model"),
+            model_revision=_optional_str(embeddings.get("model_revision")),
+            backend_version=_optional_str(embeddings.get("backend_version")),
+            cache_path=_required_str(embeddings, "cache_path"),
+            readiness_state=_required_str(embeddings, "readiness_state"),
+            last_error=_optional_str(embeddings.get("last_error")),
+        ),
+    )
+
+
+def _required_str(payload: dict[str, Any], key: str) -> str:
+    """Return a required string field or raise."""
+
+    value = payload.get(key)
+    if not isinstance(value, str) or not value:
+        raise ValueError(f"Machine config field {key!r} must be a non-empty string.")
+    return value
+
+
+def _optional_str(value: Any) -> str | None:
+    """Return an optional string value."""
+
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise ValueError("Optional string fields must be strings when present.")
+    return value or None
+
+
+def _render_machine_config(config: MachineConfig) -> str:
+    """Render one machine config to TOML."""
+
+    payload = {
+        "config_version": config.config_version,
+        "bootstrap_version": config.bootstrap_version,
+        "runtime_mode": config.runtime_mode,
+        "bootstrap_state": config.bootstrap_state,
+        "current_step": config.current_step,
+        "last_error": config.last_error,
+        "database": {
+            "app_dsn": config.database.app_dsn,
+            "admin_dsn": config.database.admin_dsn,
+        },
+        "managed": {
+            "instance_id": config.managed.instance_id,
+            "container_name": config.managed.container_name,
+            "image": config.managed.image,
+            "host": config.managed.host,
+            "port": config.managed.port,
+            "db_name": config.managed.db_name,
+            "data_dir": config.managed.data_dir,
+            "admin_user": config.managed.admin_user,
+            "admin_password": config.managed.admin_password,
+            "app_user": config.managed.app_user,
+            "app_password": config.managed.app_password,
+        },
+        "backups": {
+            "root": config.backups.root,
+            "mirror_root": config.backups.mirror_root,
+        },
+        "embeddings": {
+            "provider": config.embeddings.provider,
+            "model": config.embeddings.model,
+            "model_revision": config.embeddings.model_revision,
+            "backend_version": config.embeddings.backend_version,
+            "cache_path": config.embeddings.cache_path,
+            "readiness_state": config.embeddings.readiness_state,
+            "last_error": config.embeddings.last_error,
+        },
+    }
+    return _render_simple_payload(payload)
+
+
+def _render_simple_payload(payload: dict[str, Any]) -> str:
+    """Render a limited nested mapping to TOML."""
+
+    root_lines: list[str] = []
+    section_lines: list[str] = []
+    for key, value in payload.items():
+        if isinstance(value, dict):
+            section_lines.extend(_render_table(key, value))
+        else:
+            root_lines.append(f"{key} = {_toml_value(value)}")
+    lines = [*root_lines]
+    if root_lines and section_lines:
+        lines.append("")
+    lines.extend(section_lines)
+    return "\n".join(lines) + "\n"
+
+
+def _render_table(name: str, payload: dict[str, Any]) -> list[str]:
+    """Render one TOML table."""
+
+    lines = [f"[{name}]"]
+    for key, value in payload.items():
+        lines.append(f"{key} = {_toml_value(value)}")
+    lines.append("")
+    return lines
+
+
+def _toml_value(value: Any) -> str:
+    """Render one limited TOML literal."""
+
+    if value is None:
+        return '""'
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, int):
+        return str(value)
+    return json.dumps(str(value))
+
+
+def _timestamp_slug() -> str:
+    """Return a filesystem-safe timestamp slug."""
+
+    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")

app/periphery/admin/privileges.py

@@ -0,0 +1,42 @@
+"""Admin helpers that reconcile safe runtime privileges for the app role."""
+
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+import psycopg
+
+
+def reconcile_app_role_privileges(*, admin_dsn: str, app_dsn: str) -> None:
+    """Grant only the DML/runtime privileges the app role needs after admin migrations."""
+
+    app_user = _username_from_dsn(app_dsn)
+    if not app_user:
+        raise RuntimeError("Could not determine the app-role username from SHELLBRAIN_DB_DSN.")
+
+    raw_admin_dsn = admin_dsn.replace("+psycopg", "")
+    with psycopg.connect(raw_admin_dsn, autocommit=True) as conn:
+        with conn.cursor() as cur:
+            cur.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC")
+            cur.execute(f'REVOKE CREATE ON SCHEMA public FROM "{app_user}"')
+            cur.execute(f'GRANT USAGE ON SCHEMA public TO "{app_user}"')
+            cur.execute(
+                f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO "{app_user}"'
+            )
+            cur.execute(
+                f'GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA public TO "{app_user}"'
+            )
+            cur.execute(
+                f'ALTER DEFAULT PRIVILEGES FOR ROLE CURRENT_USER IN SCHEMA public '
+                f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO "{app_user}"'
+            )
+            cur.execute(
+                f'ALTER DEFAULT PRIVILEGES FOR ROLE CURRENT_USER IN SCHEMA public '
+                f'GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO "{app_user}"'
+            )
+
+
+def _username_from_dsn(dsn: str) -> str:
+    """Extract the username from one DSN."""
+
+    return urlparse(dsn.replace("+psycopg", "")).username or ""