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 base64
import json
import logging
import time
import uuid
from datetime import datetime, timezone
from typing import Any

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

import jwt
from meshbay_hub.auth import hub_public_key_pem, decode_access_token
from meshbay_hub.api.deps import get_current_user, require_admin
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


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]



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(timezone.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_sk_pem, _hub_id
    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) -> None:
    """Node informs hub that a chat message was posted — create notifications for offline members."""
    if not group_id:
        return
    try:
        from meshbay_hub.db.engine import get_session_factory
        from meshbay_hub.db.models import GroupMember, Group
        from meshbay_hub.api.notifications import create_notification

        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 _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()}

    claimed = set(claimed_groups or authorized)
    return claimed_id, sorted(authorized & claimed)


@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": "..."}
        raw = await ws.receive_text()
        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)
            elif msg.get("type") == "update_groups":
                new_gids = msg.get("group_ids", [])
                _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":
                asyncio.ensure_future(_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", ""),
                ))

    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:
            _connected_nodes.pop(node_id, None)
            _node_groups.pop(node_id, None)


# ── 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 asyncio.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
    }