Vestibule/scripts/test-native-messaging.py

#!/usr/bin/env python3
"""Smoke test for usher native messaging — production protocol.

Tests the full Argon2id unlock round-trip without LibreWolf. The
protocol steps are a table: each step declares its label, request,
timeout, and pass criterion. The driver runs the table in one pass and
collects failure messages — adding a protocol step is one table entry.

Steps:
  1. hello handshake
  2. ping/pong
  3. set-unlock "test-password" -> unlock-set ok=true (hash stored to disk)
  4. unlock "wrong-password"    -> granted=false, reason="invalid"
  5. unlock "test-password"     -> granted=true
  6. simulate-wake              -> Wake event

Cleans up the hash file before and after so the test is idempotent.

Usage:
    python3 scripts/test-native-messaging.py [path/to/usher]
"""
import contextlib
import json
import os
import select
import struct
import subprocess
import sys

DEFAULT_USHER = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "helper", "target", "release", "usher",
)
USHER = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_USHER

# Hash file path — must match helper/src/storage.rs
HASH_FILE = os.path.join(
    os.environ.get("XDG_CONFIG_HOME", os.path.expanduser("~/.config")),
    "vestibule", "unlock.hash"
)

STANDARD_TIMEOUT_S = 10
ARGON2_TIMEOUT_S = 15   # Argon2id at 64 MiB takes ~1-3 s
WAKE_TIMEOUT_S = 5
EXIT_TIMEOUT_S = 5


def cleanup_hash():
    with contextlib.suppress(FileNotFoundError):
        os.remove(HASH_FILE)


def send(proc, obj):
    data = json.dumps(obj).encode("utf-8")
    proc.stdin.write(struct.pack("<I", len(data)))
    proc.stdin.write(data)
    proc.stdin.flush()


def recv(proc, timeout):
    ready, _, _ = select.select([proc.stdout], [], [], timeout)
    if not ready:
        return None
    header = proc.stdout.read(4)
    if len(header) < 4:
        return None
    (n,) = struct.unpack("<I", header)
    body = proc.stdout.read(n)
    if len(body) < n:
        return None
    return json.loads(body.decode("utf-8"))


def exchange(proc, request, timeout):
    send(proc, request)
    return recv(proc, timeout)


# ─── Step criteria ─────────────────────────────────────────────────────
#
# A criterion maps a decoded response (or None on timeout/truncation) to
# a failure message, or None when the step passes.

def matches(**expected):
    """Criterion factory: every expected key/value pair must be present."""
    def criterion(response):
        if response is None:
            return "no response within timeout"
        mismatched = [
            f"{key}: expected {value!r}, got {response.get(key)!r}"
            for key, value in expected.items()
            if response.get(key) != value
        ]
        return "; ".join(mismatched) or None
    return criterion


def stores_argon2_hash(response):
    """set-unlock must succeed AND create the hash file on disk."""
    failure = matches(type="unlock-set", ok=True)(response)
    if failure:
        return failure
    return None if os.path.exists(HASH_FILE) else f"hash file not created at {HASH_FILE}"


STEPS = [
    ("hello", {"type": "hello", "client": "smoke-test", "version": "0.0.0"},
     STANDARD_TIMEOUT_S, matches(type="hello", server="usher")),
    ("ping", {"type": "ping", "echo": "production-123"},
     STANDARD_TIMEOUT_S, matches(type="pong", echo="production-123")),
    ("set-unlock", {"type": "set-unlock", "password": "test-password-123"},
     ARGON2_TIMEOUT_S, stores_argon2_hash),
    ("unlock(wrong)", {"type": "unlock", "password": "wrong-password"},
     ARGON2_TIMEOUT_S, matches(type="unlock-result", granted=False)),
    ("unlock(correct)", {"type": "unlock", "password": "test-password-123"},
     ARGON2_TIMEOUT_S, matches(type="unlock-result", granted=True)),
    ("wake", {"type": "simulate-wake"},
     WAKE_TIMEOUT_S, matches(type="wake")),
]


def run_steps(proc):
    """Run every table step; return the failure messages, in order."""
    def run(step):
        label, request, timeout, criterion = step
        response = exchange(proc, request, timeout)
        print(f"  {label:<15} ->  {response}")
        return criterion(response)
    return [message for message in map(run, STEPS) if message]


def shutdown(proc):
    """Close stdin; usher must exit promptly. Returns failure messages."""
    proc.stdin.close()
    try:
        proc.wait(timeout=EXIT_TIMEOUT_S)
    except subprocess.TimeoutExpired:
        proc.kill()
        return [f"usher did not exit within {EXIT_TIMEOUT_S}s of stdin close"]
    return []


def main():
    if not os.path.exists(USHER):
        print(f"error: usher binary not found at {USHER}", file=sys.stderr)
        print("  build it first:  (cd helper && cargo build --release)", file=sys.stderr)
        sys.exit(1)

    # Remove any hash left over from a previous run.
    cleanup_hash()

    print(f"launching {USHER}")
    proc = subprocess.Popen(
        [USHER],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        bufsize=0,
    )

    failures = run_steps(proc)
    failures.extend(shutdown(proc))

    stderr = proc.stderr.read().decode("utf-8", errors="replace").strip()
    if stderr:
        print("\n--- stderr ---")
        print(stderr)
        print("--- end stderr ---\n")

    # Leave no hash behind.
    cleanup_hash()

    if failures:
        print("FAIL:")
        print("\n".join(f"  - {failure}" for failure in failures))
        sys.exit(1)

    print("\nOK — usher production protocol works: Argon2id unlock + wake events.")


if __name__ == "__main__":
    main()