vignette/scripts/check-postgres-rls-audit.py

#!/usr/bin/env python3
"""Local opt-in Postgres RLS/audit smoke checks for Vignette.

The script creates a small set of uniquely tagged fixture rows, verifies the
live RLS/audit behavior that unit tests mock, then removes only those rows.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import socket
import sys
import uuid
from dataclasses import dataclass, field
from typing import Any
from urllib.parse import urlparse


DSN_ENV_VARS = ("VIGNETTE_RLS_AUDIT_DSN", "POSTGRES_RLS_AUDIT_DSN")
LOCAL_HOSTS = {"", "localhost", "127.0.0.1", "::1", "db", "postgres"}


@dataclass
class CheckResult:
    name: str
    status: str
    detail: str


@dataclass
class SmokeState:
    run_id: str
    cohort_a: str
    cohort_b: str
    learner_a: str | None = None
    learner_b: str | None = None
    teacher_a: str | None = None
    admin: str | None = None
    session_a: str | None = None
    session_b: str | None = None
    created_user_ids: list[str] = field(default_factory=list)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Run a read/write smoke test against a local/dev Postgres database "
            "to prove session RLS, visible_to filtering, cohort scoping, and "
            "audit insert behavior."
        )
    )
    parser.add_argument(
        "--dsn",
        default=next((os.environ.get(name) for name in DSN_ENV_VARS if os.environ.get(name)), None),
        help=(
            "Postgres DSN. May also be supplied via VIGNETTE_RLS_AUDIT_DSN "
            "or POSTGRES_RLS_AUDIT_DSN. DATABASE_URL is intentionally not read."
        ),
    )
    parser.add_argument(
        "--write-fixtures",
        action="store_true",
        help="Required opt-in: insert uniquely tagged fixture rows and clean them up afterward.",
    )
    parser.add_argument(
        "--keep-fixtures",
        action="store_true",
        help="Leave the created fixture rows in place for manual inspection.",
    )
    parser.add_argument(
        "--run-id",
        default=None,
        help="Optional run id. Defaults to a random rls-audit-smoke-* id.",
    )
    parser.add_argument(
        "--allow-non-local",
        action="store_true",
        help="Allow a DSN host outside localhost/db/postgres. Intended only for disposable dev DBs.",
    )
    parser.add_argument(
        "--allow-skips",
        action="store_true",
        help="Return exit code 0 when checks are skipped. By default skipped proof is nonzero.",
    )
    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
    return parser


def dsn_host(dsn: str) -> str:
    parsed = urlparse(dsn)
    if parsed.hostname:
        return parsed.hostname
    if parsed.scheme and parsed.path and not parsed.netloc:
        return ""
    return parsed.hostname or ""


def is_local_dsn(dsn: str) -> bool:
    host = dsn_host(dsn)
    if host in LOCAL_HOSTS:
        return True
    try:
        return socket.gethostbyname(host).startswith("127.")
    except OSError:
        return False


def print_prerequisites(parser: argparse.ArgumentParser) -> None:
    parser.print_help(sys.stderr)
    print(
        "\nPrerequisites:\n"
        "  - Run against a local/dev Postgres DB initialized with infra/db/init/*.sql.\n"
        "  - Use a non-owner app role with NOBYPASSRLS, for example vignette_app.\n"
        "  - Supply a DSN explicitly via --dsn, VIGNETTE_RLS_AUDIT_DSN, or POSTGRES_RLS_AUDIT_DSN.\n"
        "  - Pass --write-fixtures to permit temporary fixture inserts.\n",
        file=sys.stderr,
    )


async def set_context(
    conn: Any,
    *,
    role: str | None = None,
    uid: str | None = None,
    cohort: str | None = None,
    ai_view: str | None = None,
) -> None:
    await conn.execute("SELECT set_config('app.ai_context', $1, true)", "1" if ai_view else "")
    await conn.execute("SELECT set_config('app.current_ai_view', $1, true)", ai_view or "")
    await conn.execute("SELECT set_config('app.current_role', $1, true)", role or "")
    await conn.execute("SELECT set_config('app.current_uid', $1, true)", uid or "")
    await conn.execute("SELECT set_config('app.current_cohort', $1, true)", cohort or "")
    await conn.execute(
        "SELECT set_config('app.current_sens_max', $1, true)",
        {"client": "1", "counselor": "0", "evaluator": "2"}.get(ai_view or "", ""),
    )


async def fetchval_as(
    conn: Any,
    query: str,
    *args: Any,
    role: str | None = None,
    uid: str | None = None,
    cohort: str | None = None,
    ai_view: str | None = None,
) -> Any:
    async with conn.transaction():
        await set_context(conn, role=role, uid=uid, cohort=cohort, ai_view=ai_view)
        return await conn.fetchval(query, *args)


async def execute_as(
    conn: Any,
    query: str,
    *args: Any,
    role: str | None = None,
    uid: str | None = None,
    cohort: str | None = None,
    ai_view: str | None = None,
) -> str:
    async with conn.transaction():
        await set_context(conn, role=role, uid=uid, cohort=cohort, ai_view=ai_view)
        return await conn.execute(query, *args)


async def preflight(conn: Any) -> list[CheckResult]:
    results: list[CheckResult] = []
    role_row = await conn.fetchrow(
        """
        SELECT r.rolname, r.rolsuper, r.rolbypassrls
        FROM pg_roles r
        WHERE r.rolname = current_user
        """
    )
    if role_row and (role_row["rolsuper"] or role_row["rolbypassrls"]):
        results.append(
            CheckResult(
                "preflight.rls_role",
                "fail",
                f"current_user={role_row['rolname']} can bypass RLS; use a NOBYPASSRLS app role",
            )
        )

    owner_tables = await conn.fetch(
        """
        SELECT n.nspname || '.' || c.relname AS table_name
        FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname IN ('app','audit')
          AND c.relname IN ('sessions','turns','app_user','audit_log')
          AND pg_get_userbyid(c.relowner) = current_user
        ORDER BY 1
        """
    )
    if owner_tables:
        table_list = ", ".join(row["table_name"] for row in owner_tables)
        results.append(
            CheckResult(
                "preflight.table_owner",
                "fail",
                f"current_user owns RLS target tables ({table_list}) and can bypass policies",
            )
        )

    required = await conn.fetchrow(
        """
        SELECT
          to_regclass('app.app_user') IS NOT NULL AS has_user,
          to_regclass('app.sessions') IS NOT NULL AS has_sessions,
          to_regclass('app.turns') IS NOT NULL AS has_turns,
          to_regclass('audit.audit_log') IS NOT NULL AS has_audit,
          to_regprocedure('app.current_uid()') IS NOT NULL AS has_current_uid,
          to_regprocedure('app.current_role_name()') IS NOT NULL AS has_current_role,
          to_regprocedure('app.is_ai_context()') IS NOT NULL AS has_ai_context
        """
    )
    missing = [key.removeprefix("has_") for key, value in dict(required or {}).items() if not value]
    if missing:
        results.append(
            CheckResult(
                "preflight.schema",
                "skip",
                "cannot prove RLS/audit behavior; missing schema objects: " + ", ".join(sorted(missing)),
            )
        )

    policies = await conn.fetch(
        """
        SELECT schemaname, tablename, policyname
        FROM pg_policies
        WHERE schemaname = 'app'
          AND (tablename, policyname) IN (
            ('sessions','p_sessions_select'),
            ('turns','p_turns_select')
          )
        """
    )
    found = {(row["tablename"], row["policyname"]) for row in policies}
    expected = {("sessions", "p_sessions_select"), ("turns", "p_turns_select")}
    missing_policies = sorted(f"{table}.{policy}" for table, policy in expected - found)
    if missing_policies:
        results.append(
            CheckResult(
                "preflight.policies",
                "skip",
                "cannot prove RLS behavior; missing policies: " + ", ".join(missing_policies),
            )
        )

    if not results:
        results.append(CheckResult("preflight", "pass", "schema, policies, and DB role are smoke-testable"))
    return results


async def create_fixtures(conn: Any, run_id: str) -> SmokeState:
    state = SmokeState(
        run_id=run_id,
        cohort_a=f"{run_id}-cohort-a",
        cohort_b=f"{run_id}-cohort-b",
    )
    prefix = f"rls-audit-smoke:{run_id}"
    users = [
        ("learner_a", "learner", state.cohort_a),
        ("learner_b", "learner", state.cohort_b),
        ("teacher_a", "instructor", state.cohort_a),
        ("admin", "admin", None),
    ]
    for key, role, cohort in users:
        row = await conn.fetchrow(
            """
            INSERT INTO app.app_user (external_id, email, display_name, role, cohort, consent_at)
            VALUES ($1, NULL, $2, $3, $4, now())
            RETURNING user_id
            """,
            f"{prefix}:{key}",
            f"RLS smoke {key} {run_id}",
            role,
            cohort,
        )
        user_id = str(row["user_id"])
        setattr(state, key, user_id)
        state.created_user_ids.append(user_id)

    async with conn.transaction():
        await set_context(conn, role="admin", uid=state.admin)
        state.session_a = str(
            await conn.fetchval(
                """
                INSERT INTO app.sessions (
                  learner_id, runtime_case_id, persona_code, persona_display_name,
                  persona_difficulty, session_no, theory_mode, stage_path
                )
                VALUES ($1::uuid, gen_random_uuid(), $2, $3, 'easy', 1, 'humanistic', '[]'::jsonb)
                RETURNING id
                """,
                state.learner_a,
                "RLS_SMOKE_A",
                f"RLS Smoke A {run_id}",
            )
        )
        state.session_b = str(
            await conn.fetchval(
                """
                INSERT INTO app.sessions (
                  learner_id, runtime_case_id, persona_code, persona_display_name,
                  persona_difficulty, session_no, theory_mode, stage_path
                )
                VALUES ($1::uuid, gen_random_uuid(), $2, $3, 'easy', 1, 'humanistic', '[]'::jsonb)
                RETURNING id
                """,
                state.learner_b,
                "RLS_SMOKE_B",
                f"RLS Smoke B {run_id}",
            )
        )
        await conn.execute(
            """
            INSERT INTO app.turns (session_id, seq, speaker, text, text_masked, actor_kind, visible_to)
            VALUES
              ($1::uuid, 1, 'counselor', $3, $3, 'human_learner', ARRAY['client','counselor','evaluator']),
              ($1::uuid, 2, 'client', $4, $4, 'client_ai', ARRAY['client','counselor','evaluator']),
              ($1::uuid, 3, 'counselor', $5, $5, 'evaluator_ai', ARRAY['evaluator']),
              ($2::uuid, 1, 'counselor', $6, $6, 'human_learner', ARRAY['client','counselor','evaluator'])
            """,
            state.session_a,
            state.session_b,
            f"{run_id} learner A visible counselor turn",
            f"{run_id} learner A visible client turn",
            f"{run_id} evaluator-only hidden turn",
            f"{run_id} learner B visible counselor turn",
        )
    return state


async def cleanup_fixtures(conn: Any, state: SmokeState) -> None:
    prefix = f"rls-audit-smoke:{state.run_id}:%"
    # Sessions/turns/audit are RLS-protected, so deletes must run inside the same
    # admin context the fixtures were created under; otherwise the NOBYPASSRLS app
    # role silently deletes 0 rows and the app_user delete fails on a FK reference.
    async with conn.transaction():
        await set_context(conn, role="admin", uid=state.admin)
        await conn.execute(
            """
            DELETE FROM audit.audit_log
            WHERE detail->>'smoke_run_id' = $1 OR target_id LIKE $2
            """,
            state.run_id,
            prefix,
        )
        await conn.execute(
            "DELETE FROM app.turns WHERE session_id = ANY($1::uuid[])",
            [sid for sid in (state.session_a, state.session_b) if sid],
        )
        await conn.execute(
            "DELETE FROM app.sessions WHERE learner_id = ANY($1::uuid[])",
            state.created_user_ids,
        )
    # app_user is not RLS-protected; delete it last once sessions no longer ref it.
    await conn.execute(
        "DELETE FROM app.app_user WHERE external_id LIKE $1",
        prefix,
    )


async def check_learner_isolation(conn: Any, state: SmokeState) -> CheckResult:
    own = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.sessions WHERE id = $1::uuid",
        state.session_a,
        role="learner",
        uid=state.learner_a,
    )
    other = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.sessions WHERE id = $1::uuid",
        state.session_b,
        role="learner",
        uid=state.learner_a,
    )
    other_turns = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.turns WHERE session_id = $1::uuid",
        state.session_b,
        role="learner",
        uid=state.learner_a,
    )
    if own == 1 and other == 0 and other_turns == 0:
        return CheckResult(
            "learner.session_turn_isolation",
            "pass",
            "learner A can read own session and cannot read learner B session or turns",
        )
    return CheckResult(
        "learner.session_turn_isolation",
        "fail",
        f"expected own=1 other=0 other_turns=0; got own={own} other={other} other_turns={other_turns}",
    )


async def check_teacher_cohort(conn: Any, state: SmokeState) -> CheckResult:
    visible_a = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.sessions WHERE id = $1::uuid",
        state.session_a,
        role="instructor",
        uid=state.teacher_a,
        cohort=state.cohort_a,
    )
    visible_b = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.sessions WHERE id = $1::uuid",
        state.session_b,
        role="instructor",
        uid=state.teacher_a,
        cohort=state.cohort_a,
    )
    if visible_a == 1 and visible_b == 0:
        return CheckResult(
            "teacher.cohort_scoped_sessions",
            "pass",
            "instructor scoped to cohort A sees learner A session and not cohort B session",
        )
    return CheckResult(
        "teacher.cohort_scoped_sessions",
        "fail",
        f"expected cohort A session=1 cohort B session=0; got {visible_a} and {visible_b}",
    )


async def check_visible_to(conn: Any, state: SmokeState) -> CheckResult:
    learner_count = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.turns WHERE session_id = $1::uuid AND seq = 3",
        state.session_a,
        role="learner",
        uid=state.learner_a,
    )
    client_ai_count = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.turns WHERE session_id = $1::uuid AND seq = 3",
        state.session_a,
        ai_view="client",
    )
    counselor_ai_count = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.turns WHERE session_id = $1::uuid AND seq = 3",
        state.session_a,
        ai_view="counselor",
    )
    evaluator_ai_count = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.turns WHERE session_id = $1::uuid AND seq = 3",
        state.session_a,
        ai_view="evaluator",
    )
    if (learner_count, client_ai_count, counselor_ai_count, evaluator_ai_count) == (0, 0, 0, 1):
        return CheckResult(
            "visible_to.evaluator_only_hidden",
            "pass",
            "evaluator-only turn is hidden from learner/client/counselor AI and visible to evaluator AI",
        )
    return CheckResult(
        "visible_to.evaluator_only_hidden",
        "fail",
        "expected learner/client/counselor/evaluator counts 0/0/0/1; got "
        f"{learner_count}/{client_ai_count}/{counselor_ai_count}/{evaluator_ai_count}",
    )


async def check_audit_insert(conn: Any, state: SmokeState) -> CheckResult:
    target_teacher = f"rls-audit-smoke:{state.run_id}:teacher-read:{state.session_a}"
    target_admin = f"rls-audit-smoke:{state.run_id}:admin-read:{state.session_a}"
    teacher_can_read = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.sessions WHERE id = $1::uuid",
        state.session_a,
        role="instructor",
        uid=state.teacher_a,
        cohort=state.cohort_a,
    )
    admin_can_read = await fetchval_as(
        conn,
        "SELECT count(*) FROM app.sessions WHERE id = $1::uuid",
        state.session_a,
        role="admin",
        uid=state.admin,
    )
    if teacher_can_read != 1 or admin_can_read != 1:
        return CheckResult(
            "audit.teacher_admin_read_insert",
            "fail",
            f"read path not visible before audit insert; instructor={teacher_can_read} admin={admin_can_read}",
        )
    await execute_as(
        conn,
        """
        INSERT INTO audit.audit_log (actor_uid, action, target_kind, target_id, detail)
        VALUES ($1::uuid, 'read_session', 'session', $2, $3::jsonb)
        """,
        state.teacher_a,
        target_teacher,
        json.dumps({"smoke_run_id": state.run_id, "role": "instructor"}),
        role="instructor",
        uid=state.teacher_a,
        cohort=state.cohort_a,
    )
    await execute_as(
        conn,
        """
        INSERT INTO audit.audit_log (actor_uid, action, target_kind, target_id, detail)
        VALUES ($1::uuid, 'read_session', 'session', $2, $3::jsonb)
        """,
        state.admin,
        target_admin,
        json.dumps({"smoke_run_id": state.run_id, "role": "admin"}),
        role="admin",
        uid=state.admin,
    )
    count = await conn.fetchval(
        """
        SELECT count(*)
        FROM audit.audit_log
        WHERE detail->>'smoke_run_id' = $1
          AND action = 'read_session'
          AND target_id IN ($2, $3)
        """,
        state.run_id,
        target_teacher,
        target_admin,
    )
    if count == 2:
        return CheckResult(
            "audit.teacher_admin_read_insert",
            "pass",
            "instructor and admin read_session audit rows inserted and verified",
        )
    return CheckResult(
        "audit.teacher_admin_read_insert",
        "fail",
        f"expected 2 read_session audit rows for run id; got {count}",
    )


async def run_smoke(args: argparse.Namespace) -> tuple[list[CheckResult], SmokeState | None]:
    try:
        import asyncpg
    except ImportError as exc:
        return [CheckResult("preflight.asyncpg", "fail", f"asyncpg is not installed: {exc}")], None

    conn = await asyncpg.connect(args.dsn)
    state: SmokeState | None = None
    try:
        results = await preflight(conn)
        if any(result.status != "pass" for result in results):
            return results, None
        run_id = args.run_id or f"rls-audit-smoke-{uuid.uuid4().hex[:12]}"
        state = await create_fixtures(conn, run_id)
        results.extend(
            [
                await check_learner_isolation(conn, state),
                await check_teacher_cohort(conn, state),
                await check_visible_to(conn, state),
                await check_audit_insert(conn, state),
            ]
        )
        return results, state
    finally:
        if state is not None and not args.keep_fixtures:
            await cleanup_fixtures(conn, state)
        await conn.close()


def emit(results: list[CheckResult], state: SmokeState | None, *, as_json: bool) -> None:
    if as_json:
        print(
            json.dumps(
                {
                    "run_id": state.run_id if state else None,
                    "status": (
                        "fail"
                        if any(result.status == "fail" for result in results)
                        else "skip"
                        if any(result.status == "skip" for result in results)
                        else "pass"
                    ),
                    "checks": [result.__dict__ for result in results],
                },
                indent=2,
            )
        )
        return
    if state:
        print(f"run_id: {state.run_id}")
    for result in results:
        print(f"{result.status.upper():4} {result.name}: {result.detail}")


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if not args.dsn or not args.write_fixtures:
        print_prerequisites(parser)
        return 2
    if not args.allow_non_local and not is_local_dsn(args.dsn):
        print(
            f"Refusing non-local DSN host '{dsn_host(args.dsn)}'. "
            "Use --allow-non-local only for disposable dev databases.",
            file=sys.stderr,
        )
        return 2

    results, state = asyncio.run(run_smoke(args))
    emit(results, state, as_json=args.json)
    has_fail = any(result.status == "fail" for result in results)
    has_skip = any(result.status == "skip" for result in results)
    if has_fail or (has_skip and not args.allow_skips):
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())