arborist/tests/test_mesh_wire.py

"""Mesh wire — unit tests for envelopes + signature verification.

These tests don't bind a network port. End-to-end HTTP exchange is
exercised in tests/test_mesh_wire_e2e.py.
"""

from __future__ import annotations

import time
from pathlib import Path

import pytest

from arborist.mesh import init_identity
from arborist.mesh.members import add_member
from arborist.mesh.crypto import (
    generate_dh_keypair,
    generate_signing_keypair,
)
from arborist.mesh.wire import (
    ALL_TYPES,
    PATH_ANNOUNCE,
    PATH_INFO,
    PATH_REQUEST,
    TYPE_ANNOUNCE_FALSIFICATION,
    TYPE_ANNOUNCE_ROOT,
    TYPE_DELIVER_BODY,
    TYPE_REQUEST_BODY,
    WIRE_VERSION,
    WireEnvelope,
    verify_envelope_sig,
)
from arborist.store import connect


def _bootstrap_two_peer_db(tmp_path: Path) -> tuple[Path, Path, dict, dict]:
    """Initialise two peers (alice admin + bob member) on separate DBs in
    the same group with bob enrolled at epoch 1 on BOTH peers.

    Returns (alice_db, bob_db, alice_identity, bob_identity).
    """
    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()

    # Each peer's roster needs the other peer at the same epoch_id so
    # signature verification can find sender_id.sign_pub. Add bob at
    # alice's DB; add alice at bob's DB. Same epoch (1) on both sides.
    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",
        )
    finally:
        a_conn.close()

    b_conn = connect(bob_db)
    try:
        add_member(
            b_conn,
            member_id="alice",
            sign_pub=alice_id.sign_pub,
            dh_pub=alice_id.dh_pub,
            role="member",
        )
    finally:
        b_conn.close()

    return alice_db, bob_db, alice_id, bob_id


# ---------------------------------------------------------------------------
# Envelope canonicalization + sign / verify
# ---------------------------------------------------------------------------


def test_canonical_bytes_stable_across_field_order():
    """Same data, different in-memory dict order → same canonical bytes."""
    a = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=100,
        body={"document_root": "ab" * 32, "source_uri": "https://x"},
    )
    b = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=100,
        body={"source_uri": "https://x", "document_root": "ab" * 32},
    )
    assert a.canonical_bytes() == b.canonical_bytes()


def test_signed_round_trip():
    """sign_with -> from_signed reconstructs the same envelope and sig."""
    sign_priv, _sign_pub = generate_signing_keypair()
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=0,
        ts=int(time.time()),
        body={"document_root": "00" * 32, "source_uri": "https://x"},
    )
    signed = env.sign_with(sign_priv)
    env2, sig2 = WireEnvelope.from_signed(signed)
    assert env == env2
    assert env.canonical_bytes() == env2.canonical_bytes()
    assert isinstance(sig2, bytes)
    assert len(sig2) == 64  # Ed25519 signature length


def test_from_signed_rejects_unknown_type():
    sign_priv, _ = generate_signing_keypair()
    env = WireEnvelope(
        type="ANNOUNCE_ROOT",  # valid
        sender_id="alice",
        epoch_id=0,
        ts=0,
        body={},
    )
    signed = env.sign_with(sign_priv)
    signed["envelope"]["type"] = "GARBAGE"
    with pytest.raises(ValueError, match="unknown envelope type"):
        WireEnvelope.from_signed(signed)


def test_from_signed_rejects_wrong_version():
    sign_priv, _ = generate_signing_keypair()
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=0,
        ts=0,
        body={},
    )
    signed = env.sign_with(sign_priv)
    signed["envelope"]["v"] = 99
    with pytest.raises(ValueError, match="unsupported wire version"):
        WireEnvelope.from_signed(signed)


def test_from_signed_rejects_missing_envelope():
    with pytest.raises(ValueError, match="missing envelope"):
        WireEnvelope.from_signed({"sig_b64": "AA=="})


def test_from_signed_rejects_bad_base64_sig():
    sign_priv, _ = generate_signing_keypair()
    env = WireEnvelope(type=TYPE_ANNOUNCE_ROOT, sender_id="x", epoch_id=0, ts=0, body={})
    signed = env.sign_with(sign_priv)
    signed["sig_b64"] = "not-base64!!"
    with pytest.raises(ValueError, match="not valid base64"):
        WireEnvelope.from_signed(signed)


def test_all_types_constant_completeness():
    """If someone adds a new TYPE_* constant they must add it to ALL_TYPES."""
    expected = {
        "ANNOUNCE_ROOT",
        "ANNOUNCE_DERIVATION",
        "ANNOUNCE_PROVIDENCE",
        "ANNOUNCE_FALSIFICATION",
        "REQUEST_BODY",
        "DELIVER_BODY",
    }
    assert ALL_TYPES == expected


def test_paths_are_namespaced_under_mesh():
    """Defends against accidental collisions with future top-level paths."""
    for p in (PATH_ANNOUNCE, PATH_INFO, PATH_REQUEST):
        assert p.startswith("/mesh/")


# ---------------------------------------------------------------------------
# verify_envelope_sig — roster-bound check
# ---------------------------------------------------------------------------


def test_verify_envelope_sig_accepts_known_member(tmp_path):
    alice_db, _bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)

    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=int(time.time()),
        body={"document_root": "11" * 32, "source_uri": "https://x"},
    )
    sig_payload = env.canonical_bytes()
    from arborist.mesh.crypto import sign as _sign
    sig = _sign(alice_id.sign_priv, sig_payload)

    a_conn = connect(alice_db)
    try:
        assert verify_envelope_sig(a_conn, env, sig) is True
    finally:
        a_conn.close()


def test_verify_envelope_sig_rejects_tampered_body(tmp_path):
    alice_db, _bob_db, alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)

    env_orig = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=100,
        body={"document_root": "aa" * 32, "source_uri": "https://x"},
    )
    from arborist.mesh.crypto import sign as _sign
    sig = _sign(alice_id.sign_priv, env_orig.canonical_bytes())

    env_tampered = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",
        epoch_id=1,
        ts=100,
        body={"document_root": "bb" * 32, "source_uri": "https://x"},  # changed
    )
    a_conn = connect(alice_db)
    try:
        assert verify_envelope_sig(a_conn, env_tampered, sig) is False
    finally:
        a_conn.close()


def test_verify_envelope_sig_rejects_unknown_sender(tmp_path):
    alice_db, _bob_db, _alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    sign_priv, _sign_pub = generate_signing_keypair()
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="mallory",  # not in roster
        epoch_id=1,
        ts=0,
        body={"document_root": "00" * 32, "source_uri": "https://x"},
    )
    from arborist.mesh.crypto import sign as _sign
    sig = _sign(sign_priv, env.canonical_bytes())
    a_conn = connect(alice_db)
    try:
        assert verify_envelope_sig(a_conn, env, sig) is False
    finally:
        a_conn.close()


def test_verify_envelope_sig_rejects_wrong_key_for_known_member(tmp_path):
    """Mallory steals alice's member_id but signs with her own key."""
    alice_db, _bob_db, _alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    rogue_priv, _rogue_pub = generate_signing_keypair()
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="alice",  # claims to be alice
        epoch_id=1,
        ts=0,
        body={"document_root": "00" * 32, "source_uri": "https://x"},
    )
    from arborist.mesh.crypto import sign as _sign
    sig = _sign(rogue_priv, env.canonical_bytes())
    a_conn = connect(alice_db)
    try:
        assert verify_envelope_sig(a_conn, env, sig) is False
    finally:
        a_conn.close()


# ---------------------------------------------------------------------------
# Server.handle_announce — direct (no HTTP)
# ---------------------------------------------------------------------------


def test_handle_announce_writes_audit_event(tmp_path):
    """Bob sends ANNOUNCE_ROOT to Alice's server; Alice's chain gains an event."""
    from arborist.mesh.wire import MeshWireServer

    alice_db, _bob_db, _alice_id, bob_id = _bootstrap_two_peer_db(tmp_path)

    a_conn = connect(alice_db)
    try:
        before = a_conn.execute("SELECT COUNT(*) FROM audit_events").fetchone()[0]
    finally:
        a_conn.close()

    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="bob",
        epoch_id=1,
        ts=int(time.time()),
        body={
            "document_root": "ab" * 32,
            "source_uri": "https://en.wikipedia.org/wiki/Cloud_Strife",
            "chunking_version": "tok-512-v1",
            "canonicalization_version": "norm-v1",
            "schema_version": "v9.8.0",
        },
    )
    signed = env.sign_with(bob_id.sign_priv)

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

    a_conn = connect(alice_db)
    try:
        after = a_conn.execute("SELECT COUNT(*) FROM audit_events").fetchone()[0]
        last = a_conn.execute(
            "SELECT event_type, body FROM audit_events ORDER BY seq DESC LIMIT 1"
        ).fetchone()
    finally:
        a_conn.close()
    assert after == before + 1
    assert last["event_type"] == "mesh_received"
    import json as _json
    body_json = _json.loads(last["body"])
    assert body_json["wire_type"] == TYPE_ANNOUNCE_ROOT
    assert body_json["sender_id"] == "bob"


def test_handle_announce_rejects_bad_sig(tmp_path):
    """Sig from wrong key → 401, no audit event."""
    from arborist.mesh.wire import MeshWireServer

    alice_db, _bob_db, _alice_id, _bob_id = _bootstrap_two_peer_db(tmp_path)
    rogue_priv, _ = generate_signing_keypair()

    env = WireEnvelope(
        type=TYPE_ANNOUNCE_ROOT,
        sender_id="bob",
        epoch_id=1,
        ts=0,
        body={"document_root": "00" * 32, "source_uri": "x"},
    )
    from arborist.mesh.crypto import sign as _sign
    sig = _sign(rogue_priv, env.canonical_bytes())
    import base64 as _b64
    signed = {"envelope": env.__dict__, "sig_b64": _b64.b64encode(sig).decode()}

    a_conn = connect(alice_db)
    try:
        before = a_conn.execute("SELECT COUNT(*) FROM audit_events").fetchone()[0]
    finally:
        a_conn.close()

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

    a_conn = connect(alice_db)
    try:
        after = a_conn.execute("SELECT COUNT(*) FROM audit_events").fetchone()[0]
    finally:
        a_conn.close()
    assert after == before  # rejected: no event written


def test_handle_announce_rejects_request_body_type(tmp_path):
    """The /announce endpoint refuses REQUEST_BODY / DELIVER_BODY."""
    from arborist.mesh.wire import MeshWireServer

    alice_db, _bob_db, _alice_id, bob_id = _bootstrap_two_peer_db(tmp_path)
    env = WireEnvelope(
        type=TYPE_REQUEST_BODY,
        sender_id="bob",
        epoch_id=1,
        ts=0,
        body={"root": "00" * 32},
    )
    signed = env.sign_with(bob_id.sign_priv)

    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
        assert status == 400
        assert "rejects" in body["error"]
    finally:
        srv.stop()


def test_handle_request_404_on_unknown_root(tmp_path):
    """Asking for a doc the responder doesn't have returns 404 cleanly."""
    from arborist.mesh.wire import MeshWireServer

    alice_db, _bob_db, _alice_id, bob_id = _bootstrap_two_peer_db(tmp_path)
    env = WireEnvelope(
        type=TYPE_REQUEST_BODY,
        sender_id="bob",
        epoch_id=1,
        ts=0,
        body={"root": "ff" * 32},
    )
    signed = env.sign_with(bob_id.sign_priv)
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)  # wrong endpoint type → 400
        assert status == 400
        status, body = srv.handle_request(signed)
        assert status == 404
    finally:
        srv.stop()


def test_announce_falsification_round_trips(tmp_path):
    """Falsifications propagate as ANNOUNCE_FALSIFICATION."""
    from arborist.mesh.wire import MeshWireServer

    alice_db, _bob_db, _alice_id, bob_id = _bootstrap_two_peer_db(tmp_path)
    env = WireEnvelope(
        type=TYPE_ANNOUNCE_FALSIFICATION,
        sender_id="bob",
        epoch_id=1,
        ts=0,
        body={"cache_key": "ab" * 32, "reason": "answer was wrong: cited X but X is fictional"},
    )
    signed = env.sign_with(bob_id.sign_priv)
    srv = MeshWireServer(alice_db, host="127.0.0.1", port=0)
    try:
        status, body = srv.handle_announce(signed)
        assert status == 200
        assert body["wire_type"] == TYPE_ANNOUNCE_FALSIFICATION
    finally:
        srv.stop()