devplacepy/devplacepy/routers/pubsub.py

# retoor <retoor@molodetz.nl>

import json
import logging
from datetime import datetime, timezone

from fastapi import APIRouter, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse

from devplacepy.services.manager import service_manager
from devplacepy.services.pubsub import policy
from devplacepy.services.pubsub.hub import pubsub

from devplacepy.routers.dbapi._shared import error, read_body

logger = logging.getLogger(__name__)
router = APIRouter()


def _require_privileged(request):
    actor = policy.resolve_actor(request)
    if not actor.privileged:
        return None
    return actor


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _message(topic: str, data) -> dict:
    return {"type": "message", "topic": topic, "data": data, "ts": _now()}


@router.websocket("/ws")
async def pubsub_ws(websocket: WebSocket):
    await websocket.accept()
    svc = service_manager.get_service("pubsub")
    if svc is None or not svc.is_enabled():
        await websocket.close(code=1013)
        return
    if not service_manager.owns_lock():
        await websocket.close(code=4013)
        return
    actor = policy.resolve_actor(websocket)
    if actor.kind == "guest" and not policy.guests_enabled():
        await websocket.close(code=1008)
        return
    await websocket.send_json({"type": "ready", "actor": actor.kind})
    try:
        while True:
            data = await websocket.receive_json()
            kind = data.get("type")
            topic = str(data.get("topic", "")).strip()
            if kind in ("subscribe", "unsubscribe", "publish") and not policy.valid_topic(
                topic
            ):
                await websocket.send_json({"type": "error", "message": "Invalid topic name."})
                continue
            if kind == "subscribe":
                if policy.can_subscribe(actor, topic):
                    pubsub.subscribe(topic, websocket)
                    await websocket.send_json({"type": "subscribed", "topic": topic})
                else:
                    await websocket.send_json(
                        {"type": "error", "message": f"Not allowed to subscribe to {topic}."}
                    )
            elif kind == "unsubscribe":
                pubsub.unsubscribe(topic, websocket)
                await websocket.send_json({"type": "unsubscribed", "topic": topic})
            elif kind == "publish":
                if not policy.can_publish(actor, topic):
                    await websocket.send_json(
                        {"type": "error", "message": f"Not allowed to publish to {topic}."}
                    )
                    continue
                payload = data.get("data")
                if len(json.dumps(payload, default=str)) > policy.MAX_PAYLOAD_BYTES:
                    await websocket.send_json({"type": "error", "message": "Payload too large."})
                    continue
                delivered = await pubsub.publish(topic, _message(topic, payload))
                await websocket.send_json(
                    {"type": "ack", "topic": topic, "delivered": delivered}
                )
            else:
                await websocket.send_json({"type": "error", "message": "Unknown frame type."})
    except WebSocketDisconnect:
        pass
    except Exception:
        logger.exception("pubsub websocket loop failed")
    finally:
        pubsub.drop_socket(websocket)


@router.post("/publish")
async def pubsub_http_publish(request: Request):
    caller = _require_privileged(request)
    if caller is None:
        return error(403, "Admin or internal access required.")
    if not service_manager.owns_lock():
        return error(409, "Pub/sub is served by the lock-owner worker. Retry.")
    body = await read_body(request)
    topic = str(body.get("topic", "")).strip()
    if not policy.valid_topic(topic):
        return error(400, "Invalid topic name.")
    payload = body.get("data")
    delivered = await pubsub.publish(topic, _message(topic, payload))
    from devplacepy.services.audit import record as audit

    audit.record(
        request,
        "pubsub.publish",
        target_type="topic",
        target_uid=topic,
        summary=f"{caller.username or caller.kind} published to {topic}",
        metadata={"topic": topic, "delivered": delivered, "caller": caller.kind},
    )
    return JSONResponse({"topic": topic, "delivered": delivered})


@router.get("/topics")
async def pubsub_list_topics(request: Request):
    if _require_privileged(request) is None:
        return error(403, "Admin or internal access required.")
    return JSONResponse({"topics": pubsub.topics()})