path: root/packages/meshbay-hub/src/meshbay_hub/api/revocation.py

"""
MeshBay Hub — Revocation system.

Two components:
  1. WebSocket endpoint /v1/nodes/ws
     Nodes connect at startup and keep the connection alive.
     Hub pushes signed revocation tokens when a user or group is revoked.

  2. Admin endpoint POST /v1/admin/revoke
     Hub operator revokes a user or group.
     Signed revocation token is broadcast to all connected nodes.

Revocation token format (signed Ed25519):
  {
    "type":       "revocation",
    "target":     "user" | "group",
    "target_id":  "<user_id or group_id>",
    "reason":     "<reason string>",
    "revoked_at": <unix timestamp>,
    "jti":        "<uuid4>",
  }

Nodes verify the token with the hub's public key (already cached at startup).
On receipt: immediately refuse JWT tokens matching the revoked user_id,
and close active connections for that user.
"""

import asyncio
import json
import logging
import time
import uuid
from datetime import UTC, datetime

import jwt
from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, WebSocketDisconnect
from meshbay_common.background import spawn
from pydantic import BaseModel
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from meshbay_hub.api.deps import get_current_user, require_admin
from meshbay_hub.auth import decode_access_token
from meshbay_hub.db.engine import get_db
from meshbay_hub.db.models import Group, GroupMember, IPLog, Node, User

log = logging.getLogger(__name__)

router = APIRouter(tags=["revocation"])

# ── Connected node registry ───────────────────────────────────────────────────

_connected_nodes: dict[str, WebSocket] = {}   # node_id → websocket
_node_groups: dict[str, list[str]] = {}        # node_id → [group_id, ...]
_punch_events: dict[str, asyncio.Event] = {}  # node_id → signaling event

# How long an unauthenticated socket may stay open before saying who it is. The
# node sends its auth message the moment the connection opens; anything that
# has not spoken by now is holding a socket and a task for nothing.
NODE_WS_AUTH_TIMEOUT = 10.0


def is_node_connected(node_id: str) -> bool:
    return node_id in _connected_nodes


def get_connected_node_count() -> int:
    return len(_connected_nodes)


def get_online_nodes_for_group(group_id: str) -> list[str]:
    return [nid for nid, gids in _node_groups.items() if group_id in gids]


# A chat_notify costs one query per member of the group plus a write for each,
# and nothing on the node's side paces it. Without a budget, one node can keep
# the hub's database busy on behalf of a group it belongs to — a cost borne by
# every other group on the instance. Generous enough that a lively conversation
# never meets it: notifications aggregate to one row per person per group, so
# the useful rate is far below this.
NOTIFY_BURST = 30            # messages
NOTIFY_WINDOW_SECONDS = 60

_notify_window: dict[str, tuple[float, int]] = {}   # node_id → (window start, count)


def forget_node(node_id: str) -> None:
    """Drop everything a disconnected node's socket owned.

    One function rather than three lines in a `finally`, so that what a
    disconnect does can be asserted by running it instead of by re-enacting it
    — a test that re-enacts cleanup tests its own re-enactment, and would not
    have noticed `_notify_window` being added here.

    Which is the point: `_notify_window` is **not** cleared. Dropping it would
    make reconnecting the way to refill the budget, and a node's token stays
    valid for an hour. It expires by time, in `_notify_budget`.
    """
    _connected_nodes.pop(node_id, None)
    _node_groups.pop(node_id, None)


def _notify_budget(node_id: str) -> bool:
    """True if this node may send one more chat_notify now."""
    now = time.monotonic()
    if len(_notify_window) > 1000:
        # Swept here rather than on disconnect, which would let a node refill
        # its budget by reconnecting — the same token stays valid for an hour.
        for nid, (started, _) in list(_notify_window.items()):
            if now - started >= NOTIFY_WINDOW_SECONDS:
                _notify_window.pop(nid, None)
    start, count = _notify_window.get(node_id, (now, 0))
    if now - start >= NOTIFY_WINDOW_SECONDS:
        start, count = now, 0
    if count >= NOTIFY_BURST:
        _notify_window[node_id] = (start, count)
        return False
    _notify_window[node_id] = (start, count + 1)
    return True



async def _mark_hosted(group_ids: list[str]) -> None:
    """Stamp the first time a node announced it hosts each of these groups.

    `group_ids` is already narrowed to what this node may claim — the caller
    derives it from the database and a node can only shrink the set, never widen
    it (finding C2) — so being announced here is evidence the group has a host.

    Set once. A node going offline does not un-host a group, and re-stamping on
    every reconnection would make `hosted_at` a "last seen" field, which is what
    the in-memory registry is already for.
    """
    from meshbay_hub.db.engine import get_session_factory
    from meshbay_hub.db.models import Group

    if not group_ids:
        return
    try:
        async with get_session_factory()() as db:
            await db.execute(
                update(Group)
                .where(Group.id.in_(group_ids), Group.hosted_at.is_(None))
                .values(hosted_at=datetime.now(UTC)))
            await db.commit()
    except Exception as e:
        # A group that stays unhosted in the table is visible to its owner and
        # collected later; failing the socket over it would take the node down.
        log.warning("Could not mark groups hosted: %s", e)


async def broadcast_revocation(token: str) -> int:
    """Push a signed revocation token to all connected nodes. Returns count sent."""
    payload = json.dumps({"type": "revocation", "token": token})
    disconnected = []
    sent = 0
    for node_id, ws in _connected_nodes.items():
        try:
            await ws.send_text(payload)
            sent += 1
        except Exception:
            disconnected.append(node_id)
    for node_id in disconnected:
        _connected_nodes.pop(node_id, None)
    return sent


def _sign_revocation(target: str, target_id: str, reason: str) -> str:
    """Issue a signed revocation token (JWT EdDSA)."""
    from meshbay_hub.auth import _hub_id, _hub_sk_pem
    now = int(time.time())
    payload = {
        "type":       "revocation",
        "target":     target,       # "user" or "group"
        "target_id":  target_id,
        "reason":     reason,
        "revoked_at": now,
        "jti":        str(uuid.uuid4()),
        "iss":        _hub_id,
        "iat":        now,
    }
    return jwt.encode(payload, _hub_sk_pem, algorithm="EdDSA")


# ── WebSocket endpoint ────────────────────────────────────────────────────────

async def _handle_chat_notify(group_id: str, sender_name: str, sender_user_id: str,
                              *, node_id: str) -> None:
    """Node informs hub that a chat message was posted — create notifications for offline members.

    `group_id` arrives in the node's own message and is checked against what
    that node is registered for. Without the check, any connected node could
    write a notification to every member of **any** group on the hub, carrying
    a display string of its choosing, with its account having no relation to
    that group at all. Same shape as the empty group claim: something believed
    about a group the sender has nothing to do with.

    `node_id` is keyword-**required** rather than defaulted. A default here
    would mean "unchecked when the caller forgets", which is the failure this
    whole review is about.
    """
    if not group_id:
        return
    if group_id not in _node_groups.get(node_id, ()):
        log.warning("Node %s sent chat_notify for a group it does not host",
                    (node_id or "?")[:8])
        return
    try:
        from meshbay_hub.api.notifications import create_notification
        from meshbay_hub.db.engine import get_session_factory
        from meshbay_hub.db.models import Group, GroupMember

        async with get_session_factory()() as db:
            group = await db.get(Group, group_id)
            if not group:
                return
            result = await db.execute(
                select(GroupMember.user_id).where(GroupMember.group_id == group_id)
            )
            member_ids = [r[0] for r in result.all()]
            for uid in member_ids:
                if uid == sender_user_id:
                    continue
                await create_notification(
                    db, uid, "chat_message",
                    f"{sender_name or 'Someone'} posted in {group.name}",
                    link=f"#/group/{group_id}",
                    group_id=group_id,
                    # One line per conversation, moved to when it last spoke.
                    aggregate=True,
                )
            await db.commit()
    except Exception as e:
        log.warning("Chat notify failed: %s", e)


async def _reject(ws: WebSocket, detail: str, code: int) -> None:
    await ws.send_text(json.dumps({"type": "error", "detail": detail}))
    await ws.close(code=code)


async def _authorized_groups(user_id: str) -> set[str]:
    """The groups this account belongs to — the ceiling on what its nodes may claim.

    Read fresh rather than captured once at registration: a node WebSocket lives
    for hours, and a group joined in the meantime has to become claimable through
    `update_groups` without reconnecting.
    """
    from meshbay_hub.db.engine import get_session_factory

    async with get_session_factory()() as db:
        result = await db.execute(
            select(GroupMember.group_id).where(GroupMember.user_id == user_id))
        return {gid for (gid,) in result.all()}


def _claimable(claimed_groups, authorized: set[str]) -> list[str]:
    """What a node actually gets registered for. Two rules.

    A node may only *narrow* the set: `authorized` is the ceiling, or a node
    could advertise itself as a source for any group on the hub (finding C2).

    And an empty claim means **no groups**, never "all of them". This used to
    read `set(claimed_groups or authorized)`, so a node hosting nothing — which
    sends no `group_ids` at all — was registered as a host for every group its
    owner belonged to, other people's included. Such a node cannot serve any of
    them: it holds no GEK, and its own handshake refuses them with "Group not
    hosted on this node". But `/v1/groups/{id}/nodes` returns nodes in
    registration order, so once one of them won the reconnection race after a
    hub restart it became `nodes[0]` and captured the group's entire client
    traffic. Any member could take a group down for everyone, by accident,
    merely by leaving an unconfigured node running.
    """
    return sorted(authorized & set(claimed_groups or ()))


async def _authorize_node_ws(token: str, claimed_id: str, claimed_groups) -> tuple:
    """
    Resolve a node WS registration against the database.

    Returns (node_id, group_ids) on success, or (None, error_detail) on refusal.
    Uses a short-lived session on purpose: a node WebSocket lives for hours, and a
    request-scoped dependency would pin a PostgreSQL connection for its whole
    lifetime, exhausting the pool once a handful of nodes connect.
    """
    from meshbay_hub.db.engine import get_session_factory

    try:
        decoded = decode_access_token(token)
    except Exception as e:
        return None, str(e)

    if decoded.get("scope") != "node":
        return None, "Node-scoped token required"

    user_id = decoded.get("sub", "")
    if not claimed_id:
        return None, "node_id required"

    async with get_session_factory()() as db:
        node = await db.get(Node, claimed_id)
        if node is None or node.user_id != user_id:
            log.warning("Rejected WS registration for node %s by user %s",
                        claimed_id[:8], (user_id or "?")[:8])
            return None, "node_id does not belong to this account"

        user = await db.get(User, user_id)
        if user is None or user.status != "active":
            return None, "Account not active"

        # Groups come from the database. The node may narrow the set to what it
        # actually hosts, but it cannot widen it to groups it is not a member of —
        # otherwise it could advertise itself as a source for any group on the hub.
        result = await db.execute(
            select(GroupMember.group_id).where(GroupMember.user_id == user_id))
        authorized = {gid for (gid,) in result.all()}

    return claimed_id, _claimable(claimed_groups, authorized)


@router.websocket("/v1/nodes/ws")
async def node_websocket(ws: WebSocket):
    """
    Persistent WebSocket connection for nodes.

    Finding C2: this used to take `node_id` and `group_ids` straight from the
    client's first message, with no check that the authenticated user owned that
    node. Any registered user could connect with an ordinary browser token, claim a
    victim node's id, and overwrite its entry in `_connected_nodes`. Every WebRTC
    offer for that node was then relayed to the attacker, who answered with their
    own SDP — a full node impersonation, and the DTLS channel binding does not help
    because the attacker is the endpoint rather than a relay. The attacker received
    the victim's encrypted keypair bundle, their chat, and their uploads.

    Identity now comes from the token and the database, never from the message.
    """
    await ws.accept()
    node_id: str | None = None

    try:
        # Auth: expect {"type": "auth", "token": "<jwt>", "node_id": "..."}
        # Bounded: the socket is accepted before anyone is authenticated, so an
        # unbounded wait is a connection any stranger can hold open for ever.
        try:
            raw = await asyncio.wait_for(ws.receive_text(), NODE_WS_AUTH_TIMEOUT)
        except TimeoutError:
            await _reject(ws, "Authentication timed out", 4001)
            return
        msg = json.loads(raw)
        if msg.get("type") != "auth" or "token" not in msg:
            await _reject(ws, "Send auth first", 4001)
            return

        try:
            decoded = decode_access_token(msg["token"])
        except Exception as e:
            await _reject(ws, str(e), 4001)
            return

        claimed_id = msg.get("node_id") or ""

        # Refuse to displace a live registration rather than silently overwriting it.
        if claimed_id and claimed_id in _connected_nodes:
            await _reject(ws, "Node already connected", 4009)
            return

        resolved_id, result = await _authorize_node_ws(
            msg["token"], claimed_id, msg.get("group_ids"))
        if resolved_id is None:
            await _reject(ws, result, 4003)
            return
        group_ids = result

        user_id = decoded.get("sub", "")
        node_id = resolved_id
        _connected_nodes[node_id] = ws
        _node_groups[node_id] = group_ids
        await _mark_hosted(group_ids)
        log.info("Node WS connected: %s (user=%s, groups=%d)",
                 node_id[:8], user_id[:8], len(group_ids))
        await ws.send_text(json.dumps({"type": "auth_ok", "node_id": node_id}))

        # Message loop — handle ping, punch_ready, etc.
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            if msg.get("type") == "ping":
                await ws.send_text(json.dumps({"type": "pong"}))
            elif msg.get("type") == "punch_ready":
                event = _punch_events.get(node_id)
                if event:
                    event.set()
            elif msg.get("type") == "webrtc_answer":
                from meshbay_hub.api.signaling import handle_webrtc_answer
                handle_webrtc_answer(msg, node_id)
            elif msg.get("type") == "update_groups":
                # Through the same gate as the registration above. This used to
                # assign the message's list verbatim, so the ceiling that makes
                # C2 hold at authentication could be stepped over one message
                # later: a node had only to reload to claim any group on the hub.
                new_gids = _claimable(msg.get("group_ids"),
                                      await _authorized_groups(user_id))
                _node_groups[node_id] = new_gids
                await _mark_hosted(new_gids)
                log.info("Node %s updated groups: %d", node_id[:8], len(new_gids))
            elif msg.get("type") == "chat_notify":
                if not _notify_budget(node_id):
                    log.warning("Node %s exceeded its chat_notify rate", node_id[:8])
                    continue
                spawn(_handle_chat_notify(
                    msg.get("group_id", ""),
                    msg.get("sender_name", ""),
                    # The author, as the node authenticated them — not
                    # decoded["sub"], which is the machine's own account and
                    # made this filter miss everyone except the operator. A node
                    # that lied here could only suppress one notification, which
                    # is the same power it has by not sending the message at all.
                    msg.get("sender_user_id", ""),
                    node_id=node_id,
                ))

    except WebSocketDisconnect:
        log.info("Node WS disconnected: %s", (node_id or "unknown")[:8])
    except Exception as e:
        log.error("Node WS error: %s", e)
    finally:
        if node_id:
            forget_node(node_id)


# ── Admin revocation endpoint ─────────────────────────────────────────────────

class IncomingRequest(BaseModel):
    peer_ip:   str
    peer_port: int


@router.post("/v1/nodes/{node_id}/incoming", status_code=200)
async def notify_incoming(
    node_id: str,
    body: IncomingRequest,
    request: Request,
    current_user: User = Depends(get_current_user),
):
    """
    Signal a node that a client wants to connect (NAT punch coordination).
    Hub forwards the request via WebSocket; node punches NAT and replies punch_ready.

    Finding H6: peer_ip was taken verbatim, so any authenticated user could make an
    arbitrary node emit UDP packets to an address of their choosing — a small
    reflection primitive using someone else's machine. The probe target must now be
    the caller's own source address.
    """
    from meshbay_hub.api.netutil import client_ip

    caller_ip = client_ip(request)
    if body.peer_ip != caller_ip:
        raise HTTPException(
            status_code=403,
            detail="peer_ip must match the requesting address")
    if not (1 <= body.peer_port <= 65535):
        raise HTTPException(status_code=422, detail="Invalid peer_port")

    ws = _connected_nodes.get(node_id)
    if not ws:
        raise HTTPException(status_code=404, detail="Node not connected")

    event = asyncio.Event()
    _punch_events[node_id] = event

    await ws.send_text(json.dumps({
        "type": "client_incoming",
        "peer_ip": body.peer_ip,
        "peer_port": body.peer_port,
    }))

    try:
        await asyncio.wait_for(event.wait(), timeout=5.0)
    except TimeoutError:
        raise HTTPException(status_code=504, detail="Node did not respond in time")
    finally:
        _punch_events.pop(node_id, None)

    return {"status": "ready", "node_id": node_id}


class RevokeRequest(BaseModel):
    target:    str   # "user" or "group"
    target_id: str
    reason:    str = "policy_violation"


@router.post("/v1/admin/revoke", status_code=200)
async def admin_revoke(
    body: RevokeRequest,
    current_user: User = Depends(require_admin),
    db: AsyncSession = Depends(get_db),
):
    """
    Revoke a user or group. Admin only (user must be hub admin — user_id in config).
    Issues a signed revocation token and broadcasts to all connected nodes.
    Also marks the target as revoked in the database.
    """
    if body.target not in ("user", "group"):
        raise HTTPException(status_code=422, detail="target must be 'user' or 'group'")

    # Mark as revoked in DB
    if body.target == "user":
        obj = await db.get(User, body.target_id)
        if not obj:
            raise HTTPException(status_code=404, detail="User not found")
        obj.status = "revoked"
    else:
        obj = await db.get(Group, body.target_id)
        if not obj:
            raise HTTPException(status_code=404, detail="Group not found")
        obj.status = "revoked"

    db.add(IPLog(
        user_id=current_user.id,
        event=f"revoke_{body.target}",
        ip_address="admin",
        detail=f"{body.target_id}: {body.reason}",
    ))
    await db.commit()

    # Issue and broadcast signed revocation token
    rev_token = _sign_revocation(body.target, body.target_id, body.reason)
    sent = await broadcast_revocation(rev_token)

    log.warning("Revoked %s %s — broadcast to %d nodes", body.target, body.target_id[:8], sent)
    return {
        "status":      "revoked",
        "target":      body.target,
        "target_id":   body.target_id,
        "nodes_notified": sent,
        "token":       rev_token,   # admin can store this for manual distribution
    }