path: root/packages/meshbay-node/src/meshbay_node/transport/webrtc_server.py

"""
MeshBay Node — WebRTC DataChannel server for browser clients.

Browsers cannot use QUIC for NAT traversal (WebTransport doesn't allow choosing
the UDP source port — Port-Restricted Cone NAT requires exact port matching).
WebRTC DataChannel with ICE/STUN handles this automatically.

The MNP protocol (handshake, file_request, file_chunk, chat, etc.) runs
identically over WebRTC DataChannel as over QUIC streams. Same E2E encryption,
same message types, same msgpack wire format.

Wire format on the DataChannel:
  - Each message is length-prefixed msgpack (4-byte big-endian + msgpack payload)
  - Same as QUIC streams and TCP+TLS
  - DataChannel is ordered and reliable (SCTP over DTLS)

Signaling flow (handled externally by the hub):
  Browser → Hub : POST /v1/nodes/{id}/webrtc/offer  {sdp, ice_candidates}
  Hub → Node    : WS push {type: "webrtc_offer", sdp, ice_candidates, peer_id}
  Node → Hub    : WS push {type: "webrtc_answer", sdp, ice_candidates, peer_id}
  Hub → Browser : SSE/response {sdp, ice_candidates}
  After signaling, DataChannel is P2P — hub is out of the loop.
"""

import asyncio
import base64
import logging
import struct
from pathlib import Path
from typing import Any

import jwt
import msgpack
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCDataChannel
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

from meshbay_common import MNP_VERSION
from meshbay_common.crypto import pk_to_b64
from meshbay_common.webcrypto import chunk_key_aes, encrypt_chunk_aes
from meshbay_common.protocol import MNP
from meshbay_node.indexer import GroupIndex

log = logging.getLogger(__name__)

CHUNK_SIZE = 1024 * 1024
MAX_MSG = 64 * 1024 * 1024


def _pack(obj: dict) -> bytes:
    data = msgpack.packb(obj, use_bin_type=True)
    return struct.pack(">I", len(data)) + data


class _DataChannelBuffer:
    """Accumulate DataChannel messages and extract length-prefixed msgpack."""

    def __init__(self):
        self._buf = bytearray()

    def feed(self, data: bytes):
        self._buf.extend(data)

    def messages(self):
        while len(self._buf) >= 4:
            length = struct.unpack(">I", self._buf[:4])[0]
            if length > MAX_MSG:
                raise ValueError(f"Message too large: {length}")
            if len(self._buf) < 4 + length:
                break
            msg_bytes = bytes(self._buf[4:4 + length])
            del self._buf[:4 + length]
            yield msgpack.unpackb(msg_bytes, raw=False)


class WebRTCPeerSession:
    """One WebRTC peer connection, handling MNP over a DataChannel."""

    def __init__(self, pc: RTCPeerConnection, node_ctx: dict):
        self._pc = pc
        self._ctx = node_ctx
        self._channel: RTCDataChannel | None = None
        self._buffer = _DataChannelBuffer()
        self._user_id: str | None = None
        self._group_id: str | None = None

    def _setup_channel(self, channel: RTCDataChannel) -> None:
        self._channel = channel

        @channel.on("message")
        def on_message(message):
            if isinstance(message, str):
                message = message.encode()
            self._buffer.feed(message)
            for msg in self._buffer.messages():
                self._handle_message(msg)

    def _handle_message(self, msg: dict) -> None:
        mtype = msg.get("type")
        log.debug("WebRTC recv: %s", mtype)
        try:
            if mtype == MNP.HANDSHAKE:
                self._do_handshake(msg)
            elif self._user_id is None:
                self._send({"type": "error", "detail": "Handshake required"})
            elif mtype == MNP.INDEX_SYNC:
                self._do_index_sync()
            elif mtype == MNP.FILE_REQUEST:
                self._do_file_request(msg)
            elif mtype == MNP.STREAM_SEGMENT:
                self._do_stream_segment(msg)
            elif mtype == MNP.GEK_REQUEST:
                self._do_gek_request()
            elif mtype == MNP.CHAT_MESSAGE:
                self._do_chat_message(msg)
            elif mtype == MNP.CHAT_HISTORY:
                self._do_chat_history(msg)
            elif mtype == MNP.FILE_UPLOAD:
                self._do_file_upload(msg)
            else:
                log.warning("Unknown MNP message type on DataChannel: %s", mtype)
        except Exception as e:
            log.error("Error handling %s on DataChannel: %s", mtype, e)
            self._send({"type": "error", "detail": str(e)})

    def _do_handshake(self, msg: dict) -> None:
        token = msg.get("token", "")
        group_id = msg.get("group_id", "")
        try:
            decoded = jwt.decode(token, self._ctx["hub_pk_pem"], algorithms=["EdDSA"])
        except Exception as e:
            self._send({"type": "error", "detail": f"Invalid JWT: {e}"})
            return

        denylist = self._ctx.get("denylist")
        if denylist and denylist.is_denied(decoded.get("sub", ""), decoded.get("jti", "")):
            self._send({"type": "error", "detail": "Token revoked"})
            return

        if group_id and group_id not in decoded.get("groups", []):
            self._send({"type": "error", "detail": "Not a member of this group"})
            return

        if group_id and "groups" in self._ctx and group_id not in self._ctx["groups"]:
            self._send({"type": "error", "detail": "Group not hosted on this node"})
            return

        self._user_id = decoded["sub"]
        self._group_id = group_id

        peers = self._ctx.get("_peers")
        if peers is not None:
            peers[self._user_id] = self

        log.info("WebRTC handshake OK — user=%s group=%s",
                 self._user_id[:8], group_id[:8] if group_id else "none")
        self._send({
            "type": MNP.HANDSHAKE_ACK,
            "v": MNP_VERSION,
            "node_pk": pk_to_b64(self._ctx["sk_node"].public_key()),
        })

    def _group_ctx(self) -> dict:
        if "groups" in self._ctx and self._group_id:
            return self._ctx["groups"][self._group_id]
        return self._ctx

    def _do_index_sync(self) -> None:
        ctx = self._group_ctx()
        idx = ctx["index"]
        entries = [
            {
                "id": e.id, "name": e.name, "path": e.path,
                "size": e.size, "type": e.type, "added_at": e.added_at,
            }
            for e in idx.entries
        ]
        self._send({
            "type": MNP.INDEX_SYNC,
            "v": MNP_VERSION,
            "group_id": idx.group_id,
            "version": idx.version,
            "entries": entries,
        })

    def _do_gek_request(self) -> None:
        ctx = self._group_ctx()
        gek = ctx.get("gek")
        if not gek:
            self._send({"type": "error", "detail": "No GEK available"})
            return
        self._send({
            "type": MNP.GEK_RESPONSE,
            "v": MNP_VERSION,
            "gek_b64": base64.b64encode(gek).decode(),
        })

    def _do_file_request(self, msg: dict) -> None:
        ctx = self._group_ctx()
        file_id = msg["file_id"]
        chunk_index = msg["chunk_index"]
        entry = ctx["index"].get_entry(file_id)
        if not entry:
            log.warning("File not found: %s", file_id[:16])
            self._send({"type": "error", "detail": "File not found"})
            return

        file_path = ctx["shared_root"] / entry.path / entry.name
        if not file_path.exists():
            self._send({"type": "error", "detail": "File not on disk"})
            return

        file_hash = bytes.fromhex(entry.id)
        chunk_data = _read_and_encrypt(
            self._ctx["sk_node"],
            ctx["gek"],
            file_path,
            chunk_index,
            file_hash,
        )
        self._send(chunk_data)

    def _do_stream_segment(self, msg: dict) -> None:
        ctx = self._group_ctx()
        file_id = msg["file_id"]
        segment_index = msg["segment_index"]
        segment_duration = msg.get("segment_duration", 4)

        entry = ctx["index"].get_entry(file_id)
        if not entry:
            self._send({"type": "error", "detail": "File not found"})
            return

        file_path = ctx["shared_root"] / entry.path / entry.name
        if not file_path.exists():
            self._send({"type": "error", "detail": "File not on disk"})
            return

        import subprocess
        try:
            result = subprocess.run(
                ["ffmpeg", "-hide_banner", "-loglevel", "error",
                 "-ss", str(segment_index * segment_duration),
                 "-i", str(file_path),
                 "-t", str(segment_duration),
                 "-c:v", "copy", "-c:a", "copy",
                 "-f", "mpegts", "pipe:1"],
                capture_output=True, timeout=30,
            )
            if result.returncode != 0 or not result.stdout:
                self._send({"type": "error", "detail": "Segment extraction failed"})
                return
            segment_data = result.stdout
        except Exception:
            self._send({"type": "error", "detail": "Segment extraction failed"})
            return

        self._send({
            "type": MNP.STREAM_SEGMENT,
            "v": MNP_VERSION,
            "file_id": file_id,
            "segment_index": segment_index,
            "data_b64": base64.b64encode(segment_data).decode(),
            "size": len(segment_data),
        })

    def _do_chat_message(self, msg: dict) -> None:
        chat_store = self._ctx.get("chat_store")
        payload = msg.get("payload", "")
        sender_name = msg.get("sender_name", "")
        if sender_name:
            self._ctx.setdefault("_user_names", {})[self._user_id] = sender_name
        if chat_store:
            raw = payload.encode() if isinstance(payload, str) else payload
            asyncio.ensure_future(chat_store.save_message(
                sender_id=msg.get("sender_id", self._user_id),
                iteration=msg.get("iteration", 0),
                payload=raw,
                thread_id=msg.get("thread_id"),
            ))

        peers = self._ctx.get("_peers", {})
        broadcast = {
            "type": MNP.CHAT_MESSAGE,
            "v": MNP_VERSION,
            "sender_id": msg.get("sender_id", self._user_id),
            "sender_name": sender_name,
            "payload": payload,
            "thread_id": msg.get("thread_id"),
            "timestamp": __import__("time").time(),
        }
        for uid, session in list(peers.items()):
            if uid != self._user_id and session is not self:
                try:
                    session._send(broadcast)
                except Exception:
                    pass

        self._send({"type": "ack", "v": MNP_VERSION})

    def _do_chat_history(self, msg: dict) -> None:
        chat_store = self._ctx.get("chat_store")
        if not chat_store:
            self._send({
                "type": MNP.CHAT_HISTORY_RESPONSE,
                "v": MNP_VERSION,
                "messages": [],
            })
            return

        since = msg.get("since", 0)
        limit = msg.get("limit", 100)
        asyncio.ensure_future(self._send_chat_history(chat_store, since, limit))

    async def _send_chat_history(self, chat_store, since: float, limit: int) -> None:
        msgs = await chat_store.get_messages(since=since, limit=limit)
        names = self._ctx.get("_user_names", {})
        self._send({
            "type": MNP.CHAT_HISTORY_RESPONSE,
            "v": MNP_VERSION,
            "messages": [
                {
                    "id": m.id,
                    "sender_id": m.sender_id,
                    "sender_name": names.get(m.sender_id, ""),
                    "payload": m.payload.decode("utf-8", errors="replace")
                    if isinstance(m.payload, bytes) else m.payload,
                    "timestamp": m.timestamp,
                    "thread_id": m.thread_id,
                }
                for m in msgs
            ],
        })

    def _do_file_upload(self, msg: dict) -> None:
        ctx = self._group_ctx()
        filename = msg.get("filename", "")
        chunk_index = msg.get("chunk_index", 0)
        total_chunks = msg.get("total_chunks", 1)
        data = msg.get("data")

        if not filename or data is None:
            self._send({"type": "error", "detail": "Missing filename or data"})
            return

        shared_root = ctx.get("shared_root")
        if not shared_root:
            self._send({"type": "error", "detail": "No shared directory"})
            return

        upload_dir = shared_root / ".uploads"
        upload_dir.mkdir(exist_ok=True)
        safe_name = filename.replace("/", "_").replace("\\", "_").replace("..", "_")
        tmp_path = upload_dir / f"{safe_name}.part"

        if isinstance(data, str):
            chunk_bytes = base64.b64decode(data)
        else:
            chunk_bytes = bytes(data)

        mode = "ab" if chunk_index > 0 else "wb"
        with open(tmp_path, mode) as f:
            f.write(chunk_bytes)

        self._send({
            "type": MNP.FILE_UPLOAD_ACK,
            "v": MNP_VERSION,
            "chunk_index": chunk_index,
            "filename": filename,
        })

        if chunk_index + 1 >= total_chunks:
            final_path = shared_root / safe_name
            tmp_path.rename(final_path)
            log.info("Upload complete: %s (%d chunks)", safe_name, total_chunks)

    def _send(self, obj: dict) -> None:
        if self._channel and self._channel.readyState == "open":
            self._channel.send(_pack(obj))
        else:
            log.warning("WebRTC send skipped: channel=%s",
                        self._channel.readyState if self._channel else "none")

    async def close(self) -> None:
        peers = self._ctx.get("_peers")
        if peers and self._user_id:
            peers.pop(self._user_id, None)
        await self._pc.close()


def _read_and_encrypt(
    sk_node: Ed25519PrivateKey,
    gek: bytes,
    file_path: Path,
    chunk_index: int,
    file_hash: bytes,
) -> dict:
    with open(file_path, "rb") as f:
        f.seek(chunk_index * CHUNK_SIZE)
        plaintext = f.read(CHUNK_SIZE)

    ckey = chunk_key_aes(gek, file_hash, chunk_index)
    nonce, ct = encrypt_chunk_aes(ckey, plaintext)

    return {
        "type": MNP.FILE_CHUNK,
        "v": MNP_VERSION,
        "chunk_index": chunk_index,
        "plaintext_size": len(plaintext),
        "nonce": nonce,
        "ct": ct,
    }


class WebRTCTransport:
    """
    Manages WebRTC peer connections for browser clients.

    Usage:
        transport = WebRTCTransport(sk_node, hub_pk_pem, gek, shared_root, index)
        answer_sdp = await transport.handle_offer(offer_sdp, peer_id)
        # Return answer_sdp to the browser via hub signaling
    """

    def __init__(
        self,
        sk_node: Ed25519PrivateKey,
        hub_pk_pem: bytes,
        gek: bytes,
        shared_root: Path,
        index: GroupIndex,
        groups: dict[str, dict] | None = None,
        denylist: Any | None = None,
        stun_servers: list[str] | None = None,
    ):
        self._ctx: dict[str, Any] = {
            "sk_node": sk_node,
            "hub_pk_pem": hub_pk_pem,
            "gek": gek,
            "shared_root": shared_root,
            "index": index,
            "_peers": {},
        }
        if groups:
            self._ctx["groups"] = groups
        if denylist:
            self._ctx["denylist"] = denylist
        self._stun = stun_servers or ["stun:stun.l.google.com:19302"]
        self._sessions: dict[str, WebRTCPeerSession] = {}

    async def handle_offer(
        self, offer_sdp: str, peer_id: str,
    ) -> tuple[str, list[dict]]:
        """
        Process a WebRTC SDP offer from a browser client.

        Returns (answer_sdp, ice_candidates) to relay back via hub signaling.
        ICE candidates are embedded in the SDP (aiortc gathers before returning).
        """
        from aiortc import RTCIceServer, RTCConfiguration

        config = RTCConfiguration(
            iceServers=[RTCIceServer(urls=s) for s in self._stun] if self._stun else []
        )
        pc = RTCPeerConnection(configuration=config)
        session = WebRTCPeerSession(pc, self._ctx)
        self._sessions[peer_id] = session

        @pc.on("datachannel")
        def on_datachannel(channel: RTCDataChannel):
            log.info("WebRTC DataChannel opened: %s (peer=%s)", channel.label, peer_id)
            session._setup_channel(channel)

        @pc.on("connectionstatechange")
        async def on_state_change():
            state = pc.connectionState
            log.info("WebRTC connection state: %s (peer=%s)", state, peer_id)
            if state in ("failed", "closed"):
                self._sessions.pop(peer_id, None)

        offer = RTCSessionDescription(sdp=offer_sdp, type="offer")
        await pc.setRemoteDescription(offer)
        answer = await pc.createAnswer()
        await pc.setLocalDescription(answer)

        log.info("WebRTC answer ready for peer=%s", peer_id)
        return pc.localDescription.sdp, []

    async def close_peer(self, peer_id: str) -> None:
        session = self._sessions.pop(peer_id, None)
        if session:
            await session.close()

    async def close_all(self) -> None:
        for session in list(self._sessions.values()):
            await session.close()
        self._sessions.clear()

    @property
    def active_peers(self) -> int:
        return len(self._sessions)