rustpad/tests/share_session_sql_test.py

#!/usr/bin/env python3
from __future__ import annotations

import hashlib
import re
import sqlite3
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
MIGRATIONS = ROOT / "migrations" / "sqlite"
SQLITE_QUERIES = ROOT / "src" / "queries" / "sqlite.rs"


def query(name: str) -> str:
    source = SQLITE_QUERIES.read_text()
    pattern = rf'Query::{re.escape(name)}\s*=>\s*\{{\s*r#"(.*?)"#\s*\}}'
    match = re.search(pattern, source, re.S)
    if not match:
        raise AssertionError(f"Query::{name} not found")
    return match.group(1)


def parse_rfc3339(value: str) -> datetime:
    parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
    if parsed.tzinfo is None:
        raise ValueError("timestamp must include timezone")
    return parsed.astimezone(timezone.utc)


def active_permission(row: tuple[str, str, str | None] | None, now: datetime) -> str | None:
    if row is None:
        return None
    permission, session_expires_at, link_expires_at = row
    try:
        if parse_rfc3339(session_expires_at) <= now:
            return None
        if link_expires_at is not None and parse_rfc3339(link_expires_at) <= now:
            return None
    except ValueError:
        return None
    return permission if permission in {"ro", "rw"} else None


def link_active(expires_at: str | None, now: datetime) -> bool:
    if expires_at is None:
        return True
    try:
        return parse_rfc3339(expires_at) > now
    except ValueError:
        return False


def apply(connection: sqlite3.Connection, paths: list[Path]) -> None:
    for path in paths:
        connection.executescript(path.read_text())


def main() -> None:
    migrations = sorted(MIGRATIONS.glob("*.sql"))
    before_sessions = [path for path in migrations if path.name < "0026_share_link_sessions.sql"]
    session_migration = MIGRATIONS / "0026_share_link_sessions.sql"
    label_migration = MIGRATIONS / "0027_share_link_labels.sql"

    db = sqlite3.connect(":memory:")
    db.execute("PRAGMA foreign_keys = ON")
    apply(db, before_sessions)
    db.execute(
        "INSERT INTO users (nickname, nickname_key, email, email_key, password_hash) "
        "VALUES (?, ?, ?, ?, ?)",
        ("Owner", "owner", "owner@example.test", "owner@example.test", "hash"),
    )

    raw_link = "legacy-share-token"
    share_hash = hashlib.sha256(raw_link.encode()).hexdigest()
    db.execute(
        "INSERT INTO resource_share_links "
        "(token_hash, resource_kind, resource_slug, permission, expires_at, created_by) "
        "VALUES (?, 'workspace', 'private-space', 'ro', NULL, 1)",
        (share_hash,),
    )
    db.executescript(session_migration.read_text())
    db.executescript(label_migration.read_text())
    columns = {row[1] for row in db.execute("PRAGMA table_info(resource_share_links)")}
    assert "token" not in columns, "fresh schema must not contain a plaintext token column"
    assert "label" in columns, "share links must support labels"

    source_sql = query("SHARE_LINK_SESSION_SOURCE")
    insert_sql = query("SHARE_SESSION_INSERT")
    permission_sql = query("SHARE_SESSION_PERMISSION")
    revoke_sessions_sql = query("SHARE_SESSIONS_DELETE_BY_LINK")
    delete_expired_sql = query("SHARE_SESSIONS_DELETE_EXPIRED")
    link_insert_sql = query("SHARE_LINK_INSERT")
    sharing_list_sql = query("RESOURCE_SHARING_LINKS")

    now = datetime(2026, 8, 2, 23, 0, tzinfo=timezone.utc)
    source = db.execute(source_sql, (share_hash, "workspace", "private-space")).fetchone()
    assert source == (share_hash, "ro", None)

    session_raw = "opaque-browser-session"
    session_hash = hashlib.sha256(session_raw.encode()).hexdigest()
    db.execute(
        insert_sql,
        (
            session_hash,
            share_hash,
            "workspace",
            "private-space",
            "2026-08-03T00:00:00+00:00",
        ),
    )

    def permission(slug: str = "private-space") -> str | None:
        row = db.execute(permission_sql, (session_hash, "workspace", slug)).fetchone()
        return active_permission(row, now)

    assert permission() == "ro", "ro share session must remain read-only"
    assert permission("another-space") is None, "session must be scoped to one resource"

    db.execute(
        "UPDATE resource_share_links SET permission = 'rw' WHERE token_hash = ?", (share_hash,)
    )
    assert permission() == "rw", "ro -> rw change must affect active sessions immediately"
    db.execute(
        "UPDATE resource_share_links SET permission = 'ro' WHERE token_hash = ?", (share_hash,)
    )
    assert permission() == "ro", "rw -> ro change must remove write access immediately"

    db.execute(
        "UPDATE resource_share_sessions SET expires_at = ? WHERE session_token_hash = ?",
        ("2026-08-02T22:59:59+00:00", session_hash),
    )
    assert permission() is None, "expired browser session must be denied"
    db.execute(
        "UPDATE resource_share_sessions SET expires_at = ? WHERE session_token_hash = ?",
        ("2026-08-03T00:00:00+00:00", session_hash),
    )

    # Offset timestamps are parsed chronologically rather than compared as text.
    db.execute(
        "UPDATE resource_share_links SET expires_at = ? WHERE token_hash = ?",
        ("2026-08-03T01:30:00+02:00", share_hash),
    )
    row = db.execute(permission_sql, (session_hash, "workspace", "private-space")).fetchone()
    assert active_permission(row, now) == "ro"
    later = datetime(2026, 8, 2, 23, 31, tzinfo=timezone.utc)
    assert active_permission(row, later) is None, "expired link must invalidate its sessions"

    db.execute(
        "UPDATE resource_share_links SET expires_at = 'not-a-date' WHERE token_hash = ?",
        (share_hash,),
    )
    row = db.execute(permission_sql, (session_hash, "workspace", "private-space")).fetchone()
    assert active_permission(row, now) is None, "malformed expiration must fail closed"
    assert not link_active("not-a-date", now)

    db.execute(
        "UPDATE resource_share_links SET expires_at = NULL, revoked_at = ? WHERE token_hash = ?",
        ("2026-08-02T23:05:00+00:00", share_hash),
    )
    assert db.execute(permission_sql, (session_hash, "workspace", "private-space")).fetchone() is None
    deleted = db.execute(
        revoke_sessions_sql, (share_hash, "workspace", "private-space")
    ).rowcount
    assert deleted == 1, "revoking a link must remove all derived sessions"

    # Newly created links persist only the hash and an optional identifying label.
    second_raw = "one-time-returned-token"
    second_hash = hashlib.sha256(second_raw.encode()).hexdigest()
    db.execute(
        link_insert_sql,
        (second_hash, "QA link", "workspace", "second-space", "rw", None, 1),
    )
    assert db.execute(
        "SELECT label, permission FROM resource_share_links WHERE token_hash = ?", (second_hash,)
    ).fetchone() == ("QA link", "rw")
    listed_link = db.execute(sharing_list_sql, ("workspace", "second-space")).fetchone()
    assert listed_link == (
        second_hash,
        "QA link",
        "rw",
        None,
        listed_link[4],
    ), "management listing must return the label but never a plaintext share token"

    db.execute(
        insert_sql,
        (
            hashlib.sha256(b"expired-row").hexdigest(),
            second_hash,
            "workspace",
            "second-space",
            "2026-08-02T22:00:00+00:00",
        ),
    )
    db.execute(delete_expired_sql, ("2026-08-02T23:00:00+00:00",))
    assert db.execute(
        "SELECT COUNT(*) FROM resource_share_sessions WHERE share_token_hash = ?", (second_hash,)
    ).fetchone() == (0,), "expired session cleanup must remove stale rows"

    db.execute("DELETE FROM resource_share_links WHERE token_hash = ?", (second_hash,))
    assert db.execute(
        "SELECT COUNT(*) FROM resource_share_sessions WHERE share_token_hash = ?", (second_hash,)
    ).fetchone() == (0,), "link deletion must cascade to sessions"

    print("share session SQL regression tests: 17 passed")


if __name__ == "__main__":
    main()