path: root/examples/meshbay_session.py

"""
A member's session with a MeshBay node, shared by the examples.

The hub only signs you in and relays the WebRTC offer. Everything after that
goes straight to the node, encrypted under the group key.
"""

import asyncio
import base64
import hashlib
import json
import os
import struct
import sys
import time

import httpx
import msgpack
from aiortc import RTCConfiguration, RTCIceServer, RTCPeerConnection, RTCSessionDescription
from argon2.low_level import Type, hash_secret_raw
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.serialization import load_der_private_key
from meshbay_common import MNP_VERSION
from meshbay_common.crypto import pk_to_b64, unwrap_gek_aes
from meshbay_common.groupbox import PURPOSE_ACK, PURPOSE_INDEX, unseal
from meshbay_common.handshake import (
    MNP_MIN_SUPPORTED,
    challenge_transcript,
    check_version,
    handshake_transcript,
    make_proof,
    verify_proof,
    webrtc_binding,
)
from meshbay_common.join import join_transcript


def b64(data: bytes) -> str:
    return base64.b64encode(data).decode()


def unb64(text: str) -> bytes:
    return base64.b64decode(text)


def hkdf(key: bytes, info: str) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=32, salt=None,
                info=info.encode()).derive(key)


# Keys derived from the passphrase, the same way as keyderive.js

def auth_key(username: str, passphrase: str) -> str:
    """What the hub checks at sign-in. The passphrase itself never leaves here."""
    salt = hashlib.sha256(f"meshbay:auth:v1:{username}".encode()).digest()
    return b64(PBKDF2HMAC(algorithm=hashes.SHA512(), length=32, salt=salt,
                          iterations=600_000).derive(passphrase.encode()))


def bundle_master_key(username: str, passphrase: str, user_id: str, pepper_b64: str) -> bytes:
    """Argon2id of the passphrase, mixed with the pepper the hub returns at sign-in."""
    salt = hashlib.sha256(f"meshbay:bundle:v2:{username}".encode()).digest()[:16]
    a = hash_secret_raw(passphrase.encode(), salt, time_cost=3, memory_cost=131072,
                        parallelism=1, hash_len=32, type=Type.ID)
    return hkdf(a + unb64(pepper_b64), f"meshbay:bundle-master:v3|{user_id}")


def open_identity(bundle_b64: str, master: bytes, user_id: str, node_pk: str):
    """Open your identity keys for this node. The node stores them sealed."""
    raw = unb64(bundle_b64)  # "MBK3", pepper version, nonce, ciphertext
    key = hkdf(master, f"meshbay:bundle:v3|node|{node_pk}")
    aad = f"meshbay:bundle:v3|{user_id}|{node_pk}".encode()
    keys = json.loads(AESGCM(key).decrypt(raw[5:17], raw[17:], aad))
    sk_ed = load_der_private_key(unb64(keys["skEd"]), password=None)
    sk_x = load_der_private_key(unb64(keys["skX"]), password=None)
    return sk_ed, sk_x


# MNP over a WebRTC DataChannel

class Channel:
    """Frames are a 4-byte big-endian length followed by msgpack."""

    def __init__(self, dc):
        self.dc = dc
        self.buf = bytearray()
        self.queue: asyncio.Queue = asyncio.Queue()
        dc.on("message", self._on_message)

    def _on_message(self, data):
        self.buf.extend(data if isinstance(data, bytes) else data.encode())
        while len(self.buf) >= 4:
            size = struct.unpack(">I", self.buf[:4])[0]
            if len(self.buf) < 4 + size:
                break
            self.queue.put_nowait(msgpack.unpackb(bytes(self.buf[4:4 + size]), raw=False))
            del self.buf[:4 + size]

    def send(self, msg: dict) -> None:
        data = msgpack.packb(msg, use_bin_type=True)
        self.dc.send(struct.pack(">I", len(data)) + data)

    async def recv(self, *types: str) -> dict:
        """The next message of one of these types. Pushes from the node are skipped."""
        while True:
            msg = await asyncio.wait_for(self.queue.get(), timeout=60)
            if msg.get("type") == "error":
                sys.exit(f"Node refused: {msg.get('detail') or msg}")
            if msg.get("type") in types:
                return msg


def dtls_fingerprint(sdp: str) -> bytes:
    for line in sdp.splitlines():
        if line.startswith("a=fingerprint:sha-256 "):
            return bytes.fromhex(line.split(" ", 1)[1].strip().replace(":", ""))
    sys.exit("No DTLS fingerprint in the SDP")


class Session:
    """An authenticated connection to one node, for one group."""

    def __init__(self, pc, ch: Channel, group_id: str, gek: bytes):
        self.pc = pc
        self.ch = ch
        self.group_id = group_id
        self.gek = gek

    async def index(self) -> list[dict]:
        """The group's files. Each entry's path is its root name, then its folders."""
        self.ch.send({"type": "index_sync", "v": "0.1"})
        msg = await self.ch.recv("index_sync")
        return unseal(self.gek, PURPOSE_INDEX, "index_sync", self.group_id, msg)["entries"]

    async def open_transfer(self, kind: str, size: int, chunks: int) -> str:
        """Ask the node for a transfer slot and wait until it is granted."""
        tr = os.urandom(16).hex()
        self.ch.send({"type": "transfer_open", "v": "0.1", "tr": tr, "kind": kind,
                      "bytes": size, "chunks": chunks})
        while (await self.ch.recv("transfer_state"))["state"] != "granted":
            print("Waiting for a free transfer slot on the node...")
        return tr

    def close_transfer(self, tr: str) -> None:
        self.ch.send({"type": "transfer_close", "v": "0.1", "tr": tr, "reason": "done"})

    async def close(self) -> None:
        await self.pc.close()


async def connect(hub_url: str, username: str, passphrase: str,
                  group_name: str, node_id: str = "") -> Session:
    async with httpx.AsyncClient(base_url=hub_url, timeout=30) as hub:

        # 1. Sign in to the hub.
        r = await hub.post("/v1/users/login", json={
            "username": username, "auth_key": auth_key(username, passphrase)})
        if r.status_code != 200:
            sys.exit(f"Sign-in refused: {r.text}")
        login = r.json()
        auth = {"Authorization": f"Bearer {login['access_token']}"}
        user_id = (await hub.get("/v1/users/me", headers=auth)).json()["user_id"]
        master = bundle_master_key(username, passphrase, user_id, login["bundle_pepper"])

        # 2. Find the group, and a node that serves it.
        groups = (await hub.get("/v1/groups/mine", headers=auth)).json()["groups"]
        group = next((g for g in groups if g["name"] == group_name), None)
        if group is None:
            sys.exit(f"You are not a member of a group named {group_name!r}")
        group_id = group["id"]
        nodes = (await hub.get(f"/v1/groups/{group_id}/nodes", headers=auth)).json()["nodes"]
        node = next((n for n in nodes if not node_id or n["node_id"] == node_id), None)
        if node is None:
            sys.exit("No node serving this group is online")

        # 3. Open a DataChannel to the node. The hub relays the offer and the
        #    answer, and nothing after that.
        pc = RTCPeerConnection(RTCConfiguration(
            iceServers=[RTCIceServer(urls="stun:stun.l.google.com:19302")]))
        dc = pc.createDataChannel("mnp", ordered=True)
        ch = Channel(dc)
        opened = asyncio.Event()
        dc.on("open", opened.set)
        await pc.setLocalDescription(await pc.createOffer())
        offer_sdp = pc.localDescription.sdp
        r = await hub.post(f"/v1/nodes/{node['node_id']}/webrtc/offer", headers=auth,
                           json={"sdp": offer_sdp, "ice_candidates": []})
        if r.status_code != 200:
            sys.exit(f"The node did not answer: {r.text}")
        answer_sdp = r.json()["sdp"]
        await pc.setRemoteDescription(RTCSessionDescription(sdp=answer_sdp, type="answer"))
        await asyncio.wait_for(opened.wait(), timeout=30)
        binding = webrtc_binding(dtls_fingerprint(offer_sdp), dtls_fingerprint(answer_sdp))

        # 4. Handshake. The node gets a short-lived token for this group only,
        #    never the hub session token.
        r = await hub.post("/v1/nodes/mnp-token", headers=auth,
                           json={"node_pk": node["pk_node"], "group_id": group_id})
        mnp_token = r.json()["mnp_token"]

    nonce_c = os.urandom(32)
    ch.send({"type": "handshake", "v": MNP_VERSION, "v_min": MNP_MIN_SUPPORTED,
             "token": mnp_token, "group_id": group_id, "nonce": b64(nonce_c)})
    challenge = await ch.recv("handshake_challenge")
    check_version(challenge["v"], challenge["v_min"])
    nonce_s = unb64(challenge["nonce"])
    node_pk = challenge["node_pk"]
    # The node signs its challenge over this connection: node_pk is proved here.
    Ed25519PublicKey.from_public_bytes(unb64(node_pk)).verify(
        unb64(challenge["sig"]), challenge_transcript(group_id, nonce_c, nonce_s, binding))

    # 5. Your identity keys for this node, which it keeps sealed for you.
    ch.send({"type": "keypair_bundle_fetch", "v": "0.1"})
    kp = await ch.recv("keypair_bundle_resp")
    if not kp.get("found"):
        sys.exit("This node keeps no copy of your keys. Turn on Browser access in "
                 "the application's settings, or open the group once in a browser.")
    sk_ed, sk_x = open_identity(kp["bundle_enc"], master, user_id, node_pk)

    # 6. The group key. The node recognises your identity and wraps the key for it.
    pk_ed = pk_to_b64(sk_ed.public_key())
    pk_x = b64(sk_x.public_key().public_bytes_raw())
    ts = int(time.time())
    ch.send({"type": "join_request", "v": "0.1", "group_id": group_id,
             "pk_ed25519": pk_ed, "pk_x25519": pk_x, "code": "", "ts": ts,
             "sig": b64(sk_ed.sign(join_transcript(
                 node_pk_b64=node_pk, group_id=group_id, user_id=user_id,
                 pk_ed25519_b64=pk_ed, pk_x25519_b64=pk_x, nonce_node=nonce_s, ts=ts)))})
    joined = await ch.recv("join_result")
    if not joined.get("ok") or not joined.get("gek"):
        sys.exit(f"The node did not give the group key: {joined.get('reason')}")
    gek = unwrap_gek_aes(joined, sk_x.private_bytes_raw(),
                         sk_x.public_key().public_bytes_raw())

    # 7. Prove the group key, then check that the node proves it too.
    ch.send({"type": "handshake_response", "v": "0.1", "proof": b64(make_proof(
        gek, "client", group_id, nonce_c, nonce_s, binding))})
    ack = await ch.recv("handshake_ack")
    if not verify_proof(gek, unb64(ack["proof"]), "node", group_id, nonce_c, nonce_s, binding):
        sys.exit("The node does not hold the group key")
    Ed25519PublicKey.from_public_bytes(unb64(ack["node_pk"])).verify(
        unb64(ack["sig"]), handshake_transcript("node", group_id, nonce_c, nonce_s, binding))
    if ack["node_pk"] != node_pk:
        sys.exit("The node changed identity during the handshake")
    unseal(gek, PURPOSE_ACK, "handshake_ack", group_id, ack)

    return Session(pc, ch, group_id, gek)