#!/usr/bin/env python3
"""Run on the PVE node as an administrator; the terminal never calls PVE APIs.

Runtime-only USB redirection, carried over a separately authenticated SSH tunnel.
No VM config writes, drivers, controller replacement or VM restarts.
"""
import argparse
import json
import re
import socket
import time


class Qmp:
    def __init__(self, vm):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.settimeout(10)
        self.sock.connect(f"/var/run/qemu-server/{vm}.qmp")
        self.stream = self.sock.makefile("rwb", buffering=0)
        greeting = json.loads(self.stream.readline())
        if "QMP" not in greeting:
            raise RuntimeError("Not a QMP socket")
        self.sequence = 0
        self.call("qmp_capabilities")

    def call(self, command, arguments=None):
        self.sequence += 1
        ident = self.sequence
        message = {"execute": command, "id": ident}
        if arguments is not None:
            message["arguments"] = arguments
        self.stream.write(json.dumps(message).encode() + b"\n")
        while True:
            raw = self.stream.readline()
            if not raw:
                raise RuntimeError("QMP disconnected")
            reply = json.loads(raw)
            if reply.get("id") != ident:
                continue
            if "error" in reply:
                raise RuntimeError(f"{command}: {reply['error']}")
            return reply["return"]

    def close(self):
        self.stream.close()
        self.sock.close()


def ids(vm):
    return f"vdi_usb_box_{vm}", f"vdi_usb_redir_{vm}"


def peripherals(qmp):
    return {x["name"] for x in qmp.call("qom-list", {"path": "/machine/peripheral"})}


def attach(qmp, vm, port, bus, vendor, product):
    channel, device = ids(vm)
    channels = {x["label"]: x for x in qmp.call("query-chardev")}
    if channel in channels or device in peripherals(qmp):
        raise RuntimeError("POC endpoint already exists; inspect status before changing it")
    qmp.call("chardev-add", {
        "id": channel,
        "backend": {"type": "socket", "data": {
            "addr": {"type": "inet", "data": {"host": "127.0.0.1", "port": str(port)}},
            "server": False, "reconnect-ms": 1000,
        }},
    })
    try:
        qmp.call("device_add", {
            "driver": "usb-redir", "id": device, "chardev": channel, "bus": bus,
            # QEMU uses ':' within rules; spice-client uses ',' instead.
            "filter": f"0x08:0x{vendor}:0x{product}:-1:1|-1:-1:-1:-1:0",
        })
    except Exception:
        # Roll back only the channel created by this invocation.
        qmp.call("chardev-remove", {"id": channel})
        raise


def detach(qmp, vm):
    channel, device = ids(vm)
    if device in peripherals(qmp):
        qmp.call("device_del", {"id": device})
        deadline = time.monotonic() + 8
        while device in peripherals(qmp):
            if time.monotonic() >= deadline:
                raise RuntimeError("USB removal pending; leave its channel intact and inspect guest")
            time.sleep(0.2)
    if any(x["label"] == channel for x in qmp.call("query-chardev")):
        qmp.call("chardev-remove", {"id": channel})


def status(qmp, vm):
    channel, device = ids(vm)
    return {
        "vm": vm,
        "device_present": device in peripherals(qmp),
        "channels": [x for x in qmp.call("query-chardev") if x["label"] == channel],
        "usb": qmp.call("human-monitor-command", {"command-line": "info usb"}),
        "vm_status": qmp.call("query-status"),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["status", "attach", "detach"])
    parser.add_argument("--vm", type=int, required=True)
    parser.add_argument("--port", type=int, default=40001)
    parser.add_argument("--bus", default="ehci-2.0")
    parser.add_argument("--vendor")
    parser.add_argument("--product")
    parser.add_argument("--guest-ejected", action="store_true",
                        help="Confirm the guest has safely ejected the USB device before detach")
    args = parser.parse_args()
    if args.vm < 100 or not 1 <= args.port <= 65535:
        parser.error("invalid VM ID or local tunnel port")
    if args.action == "attach" and not all(
            re.fullmatch(r"[0-9a-fA-F]{4}", x or "") for x in [args.vendor, args.product]):
        parser.error("attach requires an exact four-digit --vendor and --product")
    if args.action == "detach" and not args.guest_ejected:
        parser.error("safely eject in the guest first, then pass --guest-ejected")
    qmp = Qmp(args.vm)
    try:
        if args.action == "attach":
            attach(qmp, args.vm, args.port, args.bus, args.vendor.lower(), args.product.lower())
        elif args.action == "detach":
            detach(qmp, args.vm)
        print(json.dumps(status(qmp, args.vm), ensure_ascii=False, indent=2))
    finally:
        qmp.close()


if __name__ == "__main__":
    main()