path: root/poc/spike6_gek.py

"""
MeshBay Spike 6 — GEK Distribution via X25519 + HKDF

Validates the full GEK lifecycle:

  1. Alice (admin/node) and Bob (member) register on hub
  2. Alice creates a group on hub
  3. Alice wraps the GEK for herself  → stores bundle on hub
  4. Alice fetches Bob's X25519 public key from hub
  5. Alice wraps the GEK for Bob      → stores bundle on hub
  6. Bob retrieves his bundle from hub
  7. Bob unwraps → recovers GEK
  8. Verify: recovered_gek == original_gek
  9. Bob decrypts content encrypted by Alice (Spike 5 chunk) with recovered GEK

The hub stores opaque encrypted blobs — it never sees the GEK in cleartext.
"""

import asyncio
import base64
import json
import os
import time
from pathlib import Path

import blake3
import httpx
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

HUB_URL    = "http://meshbay.org"
STATE_FILE = Path("node_state.json")
PASS = "✓"
FAIL = "✗"


# ── Key helpers ────────────────────────────────────────────────────────────────

def _sk_ed_to_b64(sk) -> str:
    return base64.b64encode(sk.private_bytes(
        serialization.Encoding.Raw, serialization.PrivateFormat.Raw,
        serialization.NoEncryption())).decode()

def _pk_to_b64(pk) -> str:
    return base64.b64encode(pk.public_bytes(
        serialization.Encoding.Raw, serialization.PublicFormat.Raw)).decode()


# ── GEK wrapping (ECIES-like) ─────────────────────────────────────────────────

def wrap_gek(gek_raw: bytes, pk_recipient_raw: bytes) -> dict:
    """
    Wrap GEK for a recipient using ephemeral X25519 key agreement + HKDF.

    Protocol (ECIES-like):
      1. Generate ephemeral keypair (sk_eph, pk_eph)
      2. shared = X25519(sk_eph, pk_recipient)
      3. wrap_key = HKDF(shared, salt=pk_eph, info="meshbay:gek_wrap:v1")
      4. wrapped = ChaCha20-Poly1305(wrap_key).encrypt(nonce, gek, aad=pk_recipient)
      5. Bundle = {pk_eph, nonce, wrapped}  ← stored on hub, opaque

    Security properties:
      - Only the recipient (who has sk_recipient) can unwrap
      - AAD binds the bundle to the specific recipient (prevents reassignment)
      - Ephemeral key ensures each bundle is unique even for the same GEK/recipient
    """
    sk_eph     = X25519PrivateKey.generate()
    pk_eph_raw = _pk_to_b64(sk_eph.public_key())

    pk_recipient = X25519PublicKey.from_public_bytes(pk_recipient_raw)
    shared = sk_eph.exchange(pk_recipient)

    wrap_key = HKDF(
        algorithm=hashes.SHA256(), length=32,
        salt=base64.b64decode(pk_eph_raw),
        info=b"meshbay:gek_wrap:v1"
    ).derive(shared)

    nonce   = os.urandom(12)
    wrapped = ChaCha20Poly1305(wrap_key).encrypt(nonce, gek_raw, pk_recipient_raw)

    return {
        "pk_eph_b64":  pk_eph_raw,
        "nonce_b64":   base64.b64encode(nonce).decode(),
        "wrapped_b64": base64.b64encode(wrapped).decode(),
    }


def unwrap_gek(bundle: dict, sk_recipient_raw: bytes, pk_recipient_raw: bytes) -> bytes:
    """
    Unwrap GEK using the recipient's X25519 private key.
    Mirrors wrap_gek() exactly.
    """
    pk_eph_raw = base64.b64decode(bundle["pk_eph_b64"])
    nonce      = base64.b64decode(bundle["nonce_b64"])
    wrapped    = base64.b64decode(bundle["wrapped_b64"])

    sk_recipient = X25519PrivateKey.from_private_bytes(sk_recipient_raw)
    pk_eph       = X25519PublicKey.from_public_bytes(pk_eph_raw)
    shared       = sk_recipient.exchange(pk_eph)

    wrap_key = HKDF(
        algorithm=hashes.SHA256(), length=32,
        salt=pk_eph_raw,
        info=b"meshbay:gek_wrap:v1"
    ).derive(shared)

    return ChaCha20Poly1305(wrap_key).decrypt(nonce, wrapped, pk_recipient_raw)


# ── Main ──────────────────────────────────────────────────────────────────────

async def main():
    print("\n=== MeshBay Spike 6 — GEK Distribution ===\n")

    # Load Alice's state (admin/node operator from Spike 3/5)
    state = json.loads(STATE_FILE.read_text()) if STATE_FILE.exists() else {}

    sk_alice_ed  = Ed25519PrivateKey.from_private_bytes(
        base64.b64decode(state["sk_ed25519_b64"]))
    sk_alice_x   = X25519PrivateKey.from_private_bytes(
        base64.b64decode(state["sk_x25519_b64"]))
    pk_alice_x_raw = base64.b64decode(_pk_to_b64(sk_alice_x.public_key()))

    gek_raw = base64.b64decode(state["gek_b64"])
    print(f"  Alice PK (X25519): {_pk_to_b64(sk_alice_x.public_key())[:24]}...")
    print(f"  GEK (to protect):  {base64.b64encode(gek_raw).decode()[:24]}...")

    # Load or generate Bob's keypairs (persist so re-runs are idempotent)
    bob_state_file = Path("bob_state.json")
    if bob_state_file.exists():
        bob_st   = json.loads(bob_state_file.read_text())
        sk_bob_ed = Ed25519PrivateKey.from_private_bytes(base64.b64decode(bob_st["sk_ed"]))
        sk_bob_x  = X25519PrivateKey.from_private_bytes(base64.b64decode(bob_st["sk_x"]))
    else:
        sk_bob_ed = Ed25519PrivateKey.generate()
        sk_bob_x  = X25519PrivateKey.generate()
        sk_bob_x_raw = sk_bob_x.private_bytes(
            serialization.Encoding.Raw, serialization.PrivateFormat.Raw,
            serialization.NoEncryption())
        bob_state_file.write_text(json.dumps({
            "sk_ed": _sk_ed_to_b64(sk_bob_ed),
            "sk_x":  base64.b64encode(sk_bob_x_raw).decode(),
        }))
    pk_bob_x_raw = base64.b64decode(_pk_to_b64(sk_bob_x.public_key()))

    async with httpx.AsyncClient(timeout=15) as c:

        # ── Step 1: Register Alice (may already exist → 409 ok) ───────────────
        print("\n[ 1 ] Register Alice and Bob on hub")
        r = await c.post(f"{HUB_URL}/v1/users/register", json={
            "username": state.get("username", "node_cbesson"),
            "password": state.get("password", "nodepass42!"),
            "pk_user_ed25519": _pk_to_b64(sk_alice_ed.public_key()),
            "pk_user_x25519":  _pk_to_b64(sk_alice_x.public_key()),
        })
        if r.status_code in (201, 409):
            print(f"  {PASS} Alice: {'registered' if r.status_code == 201 else 'already exists'}")
        else:
            raise RuntimeError(f"Alice register failed: {r.text}")

        # ── Register Bob ───────────────────────────────────────────────────────
        r = await c.post(f"{HUB_URL}/v1/users/register", json={
            "username": "bob_member",
            "password": "bobpass42!",
            "pk_user_ed25519": _pk_to_b64(sk_bob_ed.public_key()),
            "pk_user_x25519":  _pk_to_b64(sk_bob_x.public_key()),
        })
        assert r.status_code in (201, 409), f"Bob register failed: {r.text}"
        print(f"  {PASS} Bob: {'registered' if r.status_code == 201 else 'already exists'}")

        # ── Login both ─────────────────────────────────────────────────────────
        r = await c.post(f"{HUB_URL}/v1/users/login", json={
            "username": state.get("username", "node_cbesson"),
            "password": state.get("password", "nodepass42!")})
        r.raise_for_status()
        alice_token = r.json()["access_token"]
        alice_hdrs  = {"Authorization": f"Bearer {alice_token}"}

        r = await c.post(f"{HUB_URL}/v1/users/login", json={
            "username": "bob_member", "password": "bobpass42!"})
        r.raise_for_status()
        bob_token = r.json()["access_token"]
        bob_hdrs  = {"Authorization": f"Bearer {bob_token}"}
        print(f"  {PASS} Both logged in")

        # ── Step 2: Alice creates a group ──────────────────────────────────────
        print("\n[ 2 ] Alice creates group 'test-group'")
        r = await c.post(f"{HUB_URL}/v1/groups",
            json={"name": "test-group"}, headers=alice_hdrs)
        r.raise_for_status()
        group_id = r.json()["group_id"]
        print(f"  {PASS} Group created: group_id={group_id[:8]}...")

        # ── Step 3: Alice wraps GEK for herself and stores on hub ──────────────
        print("\n[ 3 ] Alice wraps GEK for herself → stores on hub")
        t0 = time.perf_counter()
        bundle_alice = wrap_gek(gek_raw, pk_alice_x_raw)
        print(f"  {PASS} Wrapped in {(time.perf_counter()-t0)*1000:.2f}ms")
        print(f"         pk_eph   : {bundle_alice['pk_eph_b64'][:24]}...")
        wrapped = bundle_alice["wrapped_b64"]
        print(f"         wrapped  : {wrapped[:24]}... "
              f"({len(base64.b64decode(wrapped))}B)")

        r = await c.post(
            f"{HUB_URL}/v1/groups/{group_id}/members/{state.get('username','node_cbesson')}/gek",
            json=bundle_alice, headers=alice_hdrs)
        r.raise_for_status()
        print(f"  {PASS} Bundle stored on hub for Alice")

        # ── Step 4: Alice fetches Bob's public key from hub ───────────────────
        print("\n[ 4 ] Alice fetches Bob's X25519 public key from hub")
        r = await c.get(f"{HUB_URL}/v1/users/bob_member/pubkeys", headers=alice_hdrs)
        r.raise_for_status()
        bob_pubkeys    = r.json()
        pk_bob_x_from_hub = base64.b64decode(bob_pubkeys["pk_x25519"])
        assert pk_bob_x_from_hub == pk_bob_x_raw, "Bob's PK from hub doesn't match!"
        print(f"  {PASS} Bob's X25519 PK fetched: {bob_pubkeys['pk_x25519'][:24]}...")

        # ── Step 5: Alice wraps GEK for Bob and stores on hub ─────────────────
        print("\n[ 5 ] Alice wraps GEK for Bob → stores on hub")
        t0 = time.perf_counter()
        bundle_bob = wrap_gek(gek_raw, pk_bob_x_from_hub)
        wrap_ms = (time.perf_counter()-t0)*1000
        print(f"  {PASS} Wrapped in {wrap_ms:.2f}ms")
        print(f"         pk_eph  : {bundle_bob['pk_eph_b64'][:24]}...")
        print("         (different from Alice's bundle — ephemeral key is unique)")

        r = await c.post(
            f"{HUB_URL}/v1/groups/{group_id}/members/bob_member/gek",
            json=bundle_bob, headers=alice_hdrs)
        r.raise_for_status()
        print(f"  {PASS} Bundle stored on hub for Bob")

        # ── Step 6: Bob retrieves his bundle from hub ──────────────────────────
        print("\n[ 6 ] Bob retrieves his GEK bundle from hub")
        r = await c.get(f"{HUB_URL}/v1/groups/{group_id}/gek", headers=bob_hdrs)
        r.raise_for_status()
        retrieved_bundle = r.json()
        assert retrieved_bundle["pk_eph_b64"]  == bundle_bob["pk_eph_b64"]
        assert retrieved_bundle["wrapped_b64"] == bundle_bob["wrapped_b64"]
        print(f"  {PASS} Bundle retrieved (matches what Alice stored)")

        # ── Step 7: Bob unwraps GEK ────────────────────────────────────────────
        print("\n[ 7 ] Bob unwraps GEK using his X25519 private key")
        t0 = time.perf_counter()
        recovered_gek = unwrap_gek(
            retrieved_bundle,
            sk_bob_x.private_bytes(serialization.Encoding.Raw,
                                   serialization.PrivateFormat.Raw,
                                   serialization.NoEncryption()),
            pk_bob_x_raw
        )
        unwrap_ms = (time.perf_counter()-t0)*1000
        print(f"  {PASS} Unwrapped in {unwrap_ms:.2f}ms")

        # ── Step 8: Verify recovered GEK == original ──────────────────────────
        print("\n[ 8 ] Verify recovered GEK matches original")
        if recovered_gek == gek_raw:
            print(f"  {PASS} recovered_gek == original_gek  ← KEY RESULT")
            print(f"         GEK: {base64.b64encode(recovered_gek).decode()[:24]}...")
        else:
            print(f"  {FAIL} GEK MISMATCH — crypto error")
            return

        # ── Step 9: Bob decrypts a chunk encrypted by Alice ───────────────────
        print("\n[ 9 ] Bob decrypts content Alice encrypted with the GEK (Spike 5 validation)")
        test_data = b"Secret group content: " + os.urandom(64)
        file_hash = blake3.blake3(test_data).digest()

        # Alice encrypts (node side)
        chunk_key = HKDF(
            algorithm=hashes.SHA256(), length=32, salt=None,
            info=b"file:" + file_hash + b":chunk:" + (0).to_bytes(4, 'big')
        ).derive(gek_raw)
        nonce = os.urandom(12)
        ciphertext = ChaCha20Poly1305(chunk_key).encrypt(nonce, test_data, None)
        print(f"  Alice encrypted {len(test_data)}B → {len(ciphertext)}B ciphertext")

        # Bob decrypts (client side, using recovered GEK)
        chunk_key_bob = HKDF(
            algorithm=hashes.SHA256(), length=32, salt=None,
            info=b"file:" + file_hash + b":chunk:" + (0).to_bytes(4, 'big')
        ).derive(recovered_gek)   # ← uses recovered GEK, not original
        plaintext = ChaCha20Poly1305(chunk_key_bob).decrypt(nonce, ciphertext, None)

        if plaintext == test_data:
            print(f"  {PASS} Bob decrypted content correctly using recovered GEK")
        else:
            print(f"  {FAIL} Decryption produced wrong plaintext")
            return

        # ── Tampering check ───────────────────────────────────────────────────
        print("\n[ 10 ] Security: wrong private key cannot unwrap ─────────────────")
        sk_eve  = X25519PrivateKey.generate()
        pk_eve  = base64.b64decode(_pk_to_b64(sk_eve.public_key()))
        try:
            _ = unwrap_gek(
                retrieved_bundle,
                sk_eve.private_bytes(serialization.Encoding.Raw,
                                     serialization.PrivateFormat.Raw,
                                     serialization.NoEncryption()),
                pk_eve
            )
            print(f"  {FAIL} Wrong key should have been rejected!")
        except Exception:
            print(f"  {PASS} Wrong private key correctly rejected (AEAD auth failed)")

        print("\n" + "="*55)
        print("Spike 6 COMPLETE — GEK distribution validated.")
        print(f"  wrap_gek    : {wrap_ms:.2f} ms")
        print(f"  unwrap_gek  : {unwrap_ms:.2f} ms")
        print("  Hub role    : stores opaque bundle, never sees GEK in clear")
        print("  Security    : wrong key rejected by AEAD authentication tag")

asyncio.run(main())