proxmox-api-simulator/app/handlers/lxc_extra.py

"""Remaining LXC console / RRD / volume helpers with durable state."""

from __future__ import annotations

import json
import secrets
from typing import Any, cast

from fastapi import Request

from app.api.registry import HandlerRegistry
from app.config import Settings
from app.handlers.lxc import _database, _lxc_resource, _state, _values
from app.security.auth import issue_ticket


def _settings(request: Request) -> Settings:
    return cast(Settings, request.app.state.settings)


async def _save_state(request: Request, resource_id: Any, state: dict[str, Any]) -> None:
    await _database(request).pool.execute(
        "UPDATE resources SET state=$2::jsonb, version=version+1, updated_at=now() WHERE id=$1",
        resource_id,
        json.dumps(state, sort_keys=True),
    )


async def _save_config(request: Request, resource_id: Any, config: dict[str, Any]) -> None:
    await _database(request).pool.execute(
        "UPDATE containers SET config=$2::jsonb WHERE resource_id=$1",
        resource_id,
        json.dumps(config, sort_keys=True),
    )


def register_lxc_extra_handlers(registry: HandlerRegistry) -> None:
    async def interfaces(request: Request, inputs: dict[str, Any]) -> list[dict[str, Any]]:
        values = _values(inputs)
        resource = await _lxc_resource(request, str(values["node"]), str(values["vmid"]))
        state = _state(resource["state"])
        ifaces = state.get("interfaces")
        return [dict(item) for item in ifaces] if isinstance(ifaces, list) else []

    async def move_volume(request: Request, inputs: dict[str, Any]) -> str:
        from app.handlers.lxc import _create_task

        values = _values(inputs)
        node, vmid = str(values["node"]), str(values["vmid"])
        resource = await _lxc_resource(request, node, vmid)
        volume = str(values.get("volume") or values.get("disk") or "rootfs")
        storage = str(values.get("storage") or "local-lvm")
        config = _state(resource["config"])
        current = str(config.get(volume) or "")
        if current:
            # rewrite storage prefix when present
            rest = current.split(":", 1)[1] if ":" in current else current
            config[volume] = f"{storage}:{rest}"
            await _save_config(request, resource["id"], config)
        state = _state(resource["state"])
        moves = state.setdefault("volume_moves", [])
        if not isinstance(moves, list):
            moves = state["volume_moves"] = []
        moves.append({"volume": volume, "storage": storage})
        await _save_state(request, resource["id"], state)
        return await _create_task(
            request,
            node=node,
            vmid=vmid,
            task_type="lxc-move-volume",
            payload={
                "resource_id": str(resource["id"]),
                "volume": volume,
                "storage": storage,
            },
        )

    async def rrd(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        values = _values(inputs)
        resource = await _lxc_resource(request, str(values["node"]), str(values["vmid"]))
        state = _state(resource["state"])
        rrd_state = state.get("rrd")
        return dict(rrd_state) if isinstance(rrd_state, dict) else {}

    async def rrddata(request: Request, inputs: dict[str, Any]) -> list[dict[str, Any]]:
        values = _values(inputs)
        resource = await _lxc_resource(request, str(values["node"]), str(values["vmid"]))
        state = _state(resource["state"])
        series = state.get("rrddata")
        return [dict(item) for item in series] if isinstance(series, list) else []

    async def _console(request: Request, inputs: dict[str, Any], kind: str) -> dict[str, Any]:
        values = _values(inputs)
        resource = await _lxc_resource(request, str(values["node"]), str(values["vmid"]))
        key = _settings(request).ticket_signing_key.get_secret_value().encode()
        ticket = issue_ticket(str(request.state.principal), key)
        port = 6900 + int(values["vmid"]) % 1000
        state = _state(resource["state"])
        consoles = state.setdefault("consoles", {})
        payload = {
            "type": kind,
            "port": port,
            "ticket": ticket,
            "upid": (
                f"UPID:{values['node']}:{secrets.token_hex(4)}:"
                f"{kind}:{values['vmid']}:{request.state.principal}:"
            ),
            "user": str(request.state.principal),
        }
        consoles[kind] = {k: v for k, v in payload.items() if k != "ticket"}
        await _save_state(request, resource["id"], state)
        return payload

    async def vncproxy(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        return await _console(request, inputs, "vnc")

    async def spiceproxy(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        return await _console(request, inputs, "spice")

    async def termproxy(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        return await _console(request, inputs, "term")

    async def mtunnel(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        return await _console(request, inputs, "mtunnel")

    async def _ws(request: Request, inputs: dict[str, Any], kind: str) -> dict[str, Any]:
        values = _values(inputs)
        resource = await _lxc_resource(request, str(values["node"]), str(values["vmid"]))
        state = _state(resource["state"])
        console = state.get("consoles", {}).get(kind) or {"port": 6900}
        key = _settings(request).ticket_signing_key.get_secret_value().encode()
        return {
            "port": console.get("port", 6900),
            "ticket": issue_ticket(str(request.state.principal), key),
        }

    async def vncwebsocket(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        return await _ws(request, inputs, "vnc")

    async def mtunnelwebsocket(request: Request, inputs: dict[str, Any]) -> dict[str, Any]:
        return await _ws(request, inputs, "mtunnel")

    base = "/nodes/{node}/lxc/{vmid}"
    registry.register(f"{base}/interfaces", "GET", interfaces)
    registry.register(f"{base}/move_volume", "POST", move_volume)
    registry.register(f"{base}/rrd", "GET", rrd)
    registry.register(f"{base}/rrddata", "GET", rrddata)
    registry.register(f"{base}/vncproxy", "POST", vncproxy)
    registry.register(f"{base}/spiceproxy", "POST", spiceproxy)
    registry.register(f"{base}/termproxy", "POST", termproxy)
    registry.register(f"{base}/mtunnel", "POST", mtunnel)
    registry.register(f"{base}/vncwebsocket", "GET", vncwebsocket)
    registry.register(f"{base}/mtunnelwebsocket", "GET", mtunnelwebsocket)