arborist/tests/test_mesh_aead.py

"""Mesh wire — AEAD-encrypted gossip bodies.

Cleartext gossip stays the default; this module verifies the opt-in
encryption path. A sender flips ``encrypt=True`` on an announce call,
the body ships AEAD-protected under the per-epoch shared secret, and
a current member of the same epoch decrypts via their wrapped slot in
the secret envelope.

The signature still spans the full envelope (including the ciphertext),
so a MITM cannot swap payloads without invalidating sig verification.
Receivers without a slot in the epoch envelope (evicted / never
enrolled) get a clean 401 — confidentiality failure mirrors auth
failure at the wire layer.
"""

from __future__ import annotations

import base64
import time
from pathlib import Path

import pytest

from arborist.mesh import init_identity
from arborist.mesh.members import add_member
from arborist.mesh.state import recover_epoch_secret
from arborist.mesh.wire import (
    MeshWireServer,
    TYPE_ANNOUNCE_FALSIFICATION,
    TYPE_ANNOUNCE_ROOT,
    WireEnvelope,
    decrypt_envelope_body,
    encrypt_body_for_envelope,
)
from arborist.store import connect


def _bootstrap_two_peer_db(tmp_path: Path) -> tuple[Path, Path, dict, dict]:
    """alice + bob in the same group, sharing epoch 1's secret envelope.

    Real federation invariant: every peer at the same epoch sees the
    SAME `mesh_epochs.secret_envelope`. Without that, sender-side
    encryption under one DB's epoch secret cannot decrypt under
    another DB's epoch secret. The signing-only fixture in
    `test_mesh_wire.py` independently rotates each side, which is fine
    for sig checks but breaks AEAD by construction.

    Here we let alice rotate (genesis+add_member(bob)), then mirror
    her epoch 1 `mesh_epochs` row + `mesh_roster` rows + matching
    audit event into bob's DB. Same wrapped secret on both sides;
    bob's `recover_epoch_secret` resolves alice as actor and unwraps
    bob's slot using alice's dh_pub. The same path the wire receiver
    walks at runtime.
    """
    import json as _json
    import sqlite3 as _sqlite3

    from arborist.store import append_audit, latest_event_hash

    alice_db = tmp_path / "alice.db"
    bob_db = tmp_path / "bob.db"

    a_conn = connect(alice_db)
    try:
        alice_id = init_identity(a_conn, group_name="t", member_id="alice")
    finally:
        a_conn.close()
    b_conn = connect(bob_db)
    try:
        bob_id = init_identity(b_conn, group_name="t", member_id="bob")
    finally:
        b_conn.close()

    # Alice rotates to epoch 1, adding bob. Her DB now has epoch-1
    # secret_envelope wrapped to {alice, bob} via alice's dh_priv.
    a_conn = connect(alice_db)
    try:
        add_member(
            a_conn,
            member_id="bob",
            sign_pub=bob_id.sign_pub,
            dh_pub=bob_id.dh_pub,
            role="member",
        )
        epoch1_row = a_conn.execute(
            "SELECT epoch_id, started_at, started_event_hash, "
            "secret_envelope, reason FROM mesh_epochs WHERE epoch_id=1"
        ).fetchone()
        alice_roster_e1 = a_conn.execute(
            "SELECT member_id, sign_pub, dh_pub, role FROM mesh_roster "
            "WHERE epoch_id=1 ORDER BY member_id"
        ).fetchall()
        # Pull the rotation audit event so we can replay its body on bob's DB.
        rot_audit = a_conn.execute(
            "SELECT body, ts FROM audit_events WHERE event_hash=?",
            (epoch1_row["started_event_hash"],),
        ).fetchone()
    finally:
        a_conn.close()

    # Mirror epoch-1 onto bob's DB. Bob also needs alice's pubs in his
    # roster_at(0) so the roster lookup at epoch 1 (the rotator alice's
    # dh_pub) resolves. The rotation audit event must replay exactly so
    # the chain links cleanly; we use append_audit which recomputes the
    # event_hash deterministically — and since bob's chain is shorter
    # we just append a fresh rotate event with alice as actor, then
    # backfill epoch row's started_event_hash to that hash.
    b_conn = connect(bob_db)
    try:
        # Replay the rotate event on bob's chain. body is canonical-JSON.
        rot_body = _json.loads(rot_audit["body"])
        replayed_hash = append_audit(
            b_conn,
            event_type="mesh_epoch_rotate",
            body=rot_body,
            ts=int(rot_audit["ts"]),
        )
        # Insert the full epoch-1 roster (alice + bob) into bob's DB.
        for r in alice_roster_e1:
            b_conn.execute(
                "INSERT OR REPLACE INTO mesh_roster "
                "(epoch_id, member_id, sign_pub, dh_pub, role) "
                "VALUES (1, ?, ?, ?, ?)",
                (r["member_id"], r["sign_pub"], r["dh_pub"], r["role"]),
            )
        # Insert the SAME secret_envelope row so bob unwraps the same key.
        b_conn.execute(
            "INSERT OR REPLACE INTO mesh_epochs "
            "(epoch_id, started_at, started_event_hash, secret_envelope, reason) "
            "VALUES (1, ?, ?, ?, ?)",
            (
                int(epoch1_row["started_at"]),
                replayed_hash,
                epoch1_row["secret_envelope"],
                epoch1_row["reason"],
            ),
        )
        b_conn.commit()
    finally:
        b_conn.close()

    return alice_db, bob_db, alice_id, bob_id


def _ann_root_body() -> dict:
    return {
        "document_root": "ab" * 32,
        "source_uri": "https://en.wikipedia.org/wiki/Confidential",
        "chunking_version": "tok-512-v1",
        "canonicalization_version": "norm-v1",
        "schema_version": "v9.8.0",
    }


def _build_encrypted_announce(
    db_path: Path,
    sign_priv: bytes,
    *,
    sender_id: str,
    epoch_id: int,
    ts: int | None = None,
    body: dict | None = None,
) -> tuple[WireEnvelope, dict]:
    """Construct + sign an encrypted ANNOUNCE_ROOT envelope from a peer.

    Returns (envelope, signed_payload). The envelope dataclass instance
    is handy for tests that want to mutate ciphertext before re-signing.
    """
    if body is None:
        body = _ann_root_body()
    if ts is None:
        ts = int(time.time())
    conn = connect(db_path)
    try:
        secret = recover_epoch_secret(conn, epoch_id=epoch_id)
    finally:
        conn.close()
    encrypted = encrypt_body_for_envelope(
        type_=TYPE_ANNOUNCE_ROOT,
        sender_id=sender_id,
        epoch_id=epoch_id,
        ts=ts,
        body=body,
        epoch_secret=secret,
    )
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id=sender_id,
        epoch_id=epoch_id,
        ts=ts,
        body={},
        encrypted_body=encrypted,
    )
    return env, env.sign_with(sign_priv)


# ---------------------------------------------------------------------------
# 1. Round-trip: alice encrypts, bob (current member) decrypts + accepts.
# ---------------------------------------------------------------------------


def test_encrypted_announce_round_trip(tmp_path):
    """Alice sends an encrypted ANNOUNCE_ROOT; bob's server decrypts and
    accepts. The audit_event reflects the cleartext body and marks
    `encrypted=True` for operator visibility."""
    alice_db, bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    env, signed = _build_encrypted_announce(
        alice_db, alice_id.sign_priv, sender_id="alice", epoch_id=1
    )
    # Ciphertext must NOT contain the plaintext doc_root.
    assert env.encrypted_body is not None
    ct_b64 = env.encrypted_body["ct_b64"]
    assert "document_root" not in base64.b64decode(ct_b64).decode(
        "utf-8", errors="ignore"
    )

    srv = MeshWireServer(bob_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 200, body
    assert body["wire_type"] == TYPE_ANNOUNCE_ROOT

    b_conn = connect(bob_db)
    try:
        last = b_conn.execute(
            "SELECT body, subject_root FROM audit_events ORDER BY seq DESC LIMIT 1"
        ).fetchone()
    finally:
        b_conn.close()
    import json
    body_json = json.loads(last["body"])
    assert body_json["encrypted"] is True
    # Decrypted plaintext threaded into the audit body.
    assert body_json["envelope_body"]["document_root"] == "ab" * 32
    assert last["subject_root"] == "ab" * 32


# ---------------------------------------------------------------------------
# 2. Receiver missing the epoch secret — evicted member returns 401.
# ---------------------------------------------------------------------------


def test_encrypted_announce_evicted_receiver_401(tmp_path):
    """Carol has alice in her roster (so signature verifies) but no slot
    in the epoch envelope (she manufactured her own genesis). The
    decrypt path therefore fails and `handle_announce` returns 401.
    """
    alice_db, _bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    carol_db = tmp_path / "carol.db"
    c_conn = connect(carol_db)
    try:
        carol_id = init_identity(c_conn, group_name="t", member_id="carol")
        # Splice alice into carol's epoch-1 roster directly so the SIG
        # check passes — but carol's secret_envelope at epoch 1 has
        # carol's slot only (she was its founder), wrapped to carol's
        # OWN dh keys, not alice's. So decrypt under that slot yields
        # the wrong key for alice's ciphertext → AEAD tag mismatch.
        # First bump carol's epoch to 1 by adding alice via add_member,
        # which creates a fresh epoch 1 wrapped to carol+alice.
        add_member(
            c_conn,
            member_id="alice",
            sign_pub=alice_id.sign_pub,
            dh_pub=alice_id.dh_pub,
            role="member",
        )
    finally:
        c_conn.close()

    # Build an encrypted ANNOUNCE on alice's DB (alice's epoch-1 secret).
    _env, signed = _build_encrypted_announce(
        alice_db, alice_id.sign_priv, sender_id="alice", epoch_id=1
    )
    # Send to carol's server. Sig verifies (alice is in carol's roster),
    # but carol's epoch-1 secret differs from alice's → AEAD fails.
    srv = MeshWireServer(carol_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 401
    assert "AEAD decrypt" in body["error"] or "decrypt" in body["error"].lower()

    # No audit event was written for the rejected message.
    c_conn = connect(carol_db)
    try:
        last = c_conn.execute(
            "SELECT event_type FROM audit_events ORDER BY seq DESC LIMIT 1"
        ).fetchone()
    finally:
        c_conn.close()
    # The most recent event should NOT be a `mesh_received` for the rejected msg.
    assert last is None or last["event_type"] != "mesh_received"


# ---------------------------------------------------------------------------
# 3. Tampered ciphertext: sig stays valid (we re-sign) but AEAD fails.
# ---------------------------------------------------------------------------


def test_encrypted_announce_tampered_ciphertext_401(tmp_path):
    """Flip a byte of the ciphertext and re-sign with a valid key
    (alice's). The Ed25519 sig still verifies — the protection here
    comes from AEAD's authentication tag, which detects the tamper.
    """
    alice_db, bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    env, _signed = _build_encrypted_announce(
        alice_db, alice_id.sign_priv, sender_id="alice", epoch_id=1
    )
    # Flip a byte inside ct_b64.
    assert env.encrypted_body is not None
    raw_ct = bytearray(base64.b64decode(env.encrypted_body["ct_b64"]))
    raw_ct[0] ^= 0x01
    tampered_ct_b64 = base64.b64encode(bytes(raw_ct)).decode("ascii")
    tampered_env = WireEnvelope(
        type=env.type,
        sender_id=env.sender_id,
        epoch_id=env.epoch_id,
        ts=env.ts,
        body=env.body,
        encrypted_body={
            "nonce_b64": env.encrypted_body["nonce_b64"],
            "ct_b64": tampered_ct_b64,
        },
    )
    signed = tampered_env.sign_with(alice_id.sign_priv)

    srv = MeshWireServer(bob_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 401
    assert "decrypt" in body["error"].lower()


# ---------------------------------------------------------------------------
# 4. Wrong nonce / wrong AAD — both yield AEAD tag mismatch.
# ---------------------------------------------------------------------------


def test_encrypted_announce_swapped_nonce_401(tmp_path):
    """Re-sign with a freshly randomized nonce (key + AAD unchanged).
    AEAD won't authenticate.
    """
    alice_db, bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    env, _signed = _build_encrypted_announce(
        alice_db, alice_id.sign_priv, sender_id="alice", epoch_id=1
    )
    assert env.encrypted_body is not None
    import os as _os

    bad_nonce_b64 = base64.b64encode(_os.urandom(12)).decode("ascii")
    tampered_env = WireEnvelope(
        type=env.type,
        sender_id=env.sender_id,
        epoch_id=env.epoch_id,
        ts=env.ts,
        body=env.body,
        encrypted_body={
            "nonce_b64": bad_nonce_b64,
            "ct_b64": env.encrypted_body["ct_b64"],
        },
    )
    signed = tampered_env.sign_with(alice_id.sign_priv)
    srv = MeshWireServer(bob_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 401


def test_encrypted_announce_aad_drift_401(tmp_path):
    """Build a ciphertext with one ts, then sign an envelope with a
    different ts. AAD is recomputed receiver-side from the envelope
    metadata, so it won't match — AEAD fails.
    """
    alice_db, bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    env, _signed = _build_encrypted_announce(
        alice_db, alice_id.sign_priv, sender_id="alice", epoch_id=1, ts=1000
    )
    assert env.encrypted_body is not None
    drifted_env = WireEnvelope(
        type=env.type,
        sender_id=env.sender_id,
        epoch_id=env.epoch_id,
        ts=2000,  # different ts → AAD changes
        body=env.body,
        encrypted_body=env.encrypted_body,
    )
    signed = drifted_env.sign_with(alice_id.sign_priv)
    srv = MeshWireServer(bob_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
    finally:
        srv.stop()
    assert status == 401


# ---------------------------------------------------------------------------
# 5. Mixed mode: cleartext + encrypted from the same client both flow.
# ---------------------------------------------------------------------------


def test_mixed_cleartext_and_encrypted_round_trip(tmp_path):
    """Alice sends one cleartext ANNOUNCE_ROOT, then one encrypted
    ANNOUNCE_FALSIFICATION. Bob accepts both; audit chain logs the
    `encrypted` flag accurately on each."""
    from arborist.mesh.wire import MeshWireClient

    alice_db, bob_db, _alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    bob_srv = MeshWireServer(bob_db, host="127.0.0.1", port=0)
    bob_srv.start_in_thread()
    try:
        with MeshWireClient(alice_db, bob_srv.url) as c:
            r1 = c.announce_root(
                document_root="cd" * 32,
                source_uri="https://x/cleartext",
                chunking_version="tok-512-v1",
                canonicalization_version="norm-v1",
                schema_version="v9.8.0",
            )
            r2 = c.announce_falsification(
                cache_key="ef" * 32,
                reason="encrypted falsification",
                encrypt=True,
            )
        assert r1["wire_type"] == TYPE_ANNOUNCE_ROOT
        assert r2["wire_type"] == TYPE_ANNOUNCE_FALSIFICATION
    finally:
        bob_srv.stop()

    b_conn = connect(bob_db)
    try:
        rows = b_conn.execute(
            "SELECT body FROM audit_events WHERE event_type='mesh_received' "
            "ORDER BY seq DESC LIMIT 2"
        ).fetchall()
    finally:
        b_conn.close()
    import json
    bodies = [json.loads(r["body"]) for r in rows]
    # rows are descending; second message (encrypted falsification) is newest
    enc, plain = bodies[0], bodies[1]
    assert enc["wire_type"] == TYPE_ANNOUNCE_FALSIFICATION
    assert enc["encrypted"] is True
    assert enc["envelope_body"]["cache_key"] == "ef" * 32
    assert plain["wire_type"] == TYPE_ANNOUNCE_ROOT
    assert plain["encrypted"] is False
    assert plain["envelope_body"]["document_root"] == "cd" * 32


# ---------------------------------------------------------------------------
# 6. Cleartext-only mode (current default) — wire format unchanged.
# ---------------------------------------------------------------------------


def test_cleartext_default_envelope_unchanged(tmp_path):
    """An envelope with `encrypt=False` (default) MUST canonicalize
    identically to the v1 wire format — no `encrypted_body` field on
    the canonical bytes, no signature change, no audit-body change.
    """
    from arborist.mesh.wire import MeshWireClient

    alice_db, bob_db, _alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    bob_srv = MeshWireServer(bob_db, host="127.0.0.1", port=0)
    bob_srv.start_in_thread()
    try:
        with MeshWireClient(alice_db, bob_srv.url) as c:
            r = c.announce_root(
                document_root="01" * 32,
                source_uri="https://x/plain",
                chunking_version="tok-512-v1",
                canonicalization_version="norm-v1",
                schema_version="v9.8.0",
            )
        assert r["ok"] is True
    finally:
        bob_srv.stop()

    b_conn = connect(bob_db)
    try:
        last = b_conn.execute(
            "SELECT body FROM audit_events WHERE event_type='mesh_received' "
            "ORDER BY seq DESC LIMIT 1"
        ).fetchone()
    finally:
        b_conn.close()
    import json
    body_json = json.loads(last["body"])
    assert body_json["encrypted"] is False
    assert body_json["envelope_body"]["document_root"] == "01" * 32


# ---------------------------------------------------------------------------
# Hardening: round-trip the AEAD primitive directly (no HTTP).
# ---------------------------------------------------------------------------


def test_decrypt_envelope_body_direct(tmp_path):
    """Sanity check: encrypt_body_for_envelope + decrypt_envelope_body
    are exact inverses given the same key, AAD, and metadata."""
    alice_db, _bob_db, _alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    a_conn = connect(alice_db)
    try:
        secret = recover_epoch_secret(a_conn, epoch_id=1)
    finally:
        a_conn.close()
    body = {"document_root": "11" * 32, "source_uri": "https://x"}
    enc = encrypt_body_for_envelope(
        type_=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=42,
        body=body,
        epoch_secret=secret,
    )
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=42,
        body={},
        encrypted_body=enc,
    )
    decoded = decrypt_envelope_body(env, secret)
    assert decoded == body


def test_decrypt_envelope_body_rejects_no_encrypted_body():
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=0,
        body={"document_root": "00" * 32},
    )
    with pytest.raises(ValueError, match="no encrypted_body"):
        decrypt_envelope_body(env, b"\x00" * 32)