218 lines
8.5 KiB
Python
218 lines
8.5 KiB
Python
#!/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 storage_filter(vendor=None, product=None):
|
|
"""Allow only USB Mass Storage, optionally narrowed to one VID/PID."""
|
|
if vendor is None or product is None:
|
|
return "0x08:-1:-1:-1:1|-1:-1:-1:-1:0"
|
|
return f"0x08:0x{vendor}:0x{product}:-1:1|-1:-1:-1:-1:0"
|
|
|
|
|
|
def attach(qmp, vm, port, bus, vendor=None, product=None):
|
|
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": storage_filter(vendor, product),
|
|
})
|
|
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 ensure(qmp, vm, port, bus, vendor=None, product=None):
|
|
"""Attach the runtime endpoint if the VM is running and it is absent."""
|
|
channel, device = ids(vm)
|
|
if device not in peripherals(qmp):
|
|
attach(qmp, vm, port, bus, vendor, product)
|
|
|
|
|
|
def probe_status(port, host="127.0.0.1", timeout=0.7):
|
|
"""Return whether the box has an exportable storage device ready."""
|
|
try:
|
|
with socket.create_connection((host, port), timeout=timeout) as stream:
|
|
stream.settimeout(timeout)
|
|
return stream.recv(32).strip() == b"ready"
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def watch(vm, port, bus, vendor=None, product=None, interval=3,
|
|
status_port=None, remove_after=3):
|
|
"""Track the box device and keep the guest endpoint in sync.
|
|
|
|
With status_port, the QEMU endpoint is present only while the box reports
|
|
an exportable physical storage device. A short missing-device debounce
|
|
prevents a transient udev/export restart from causing a guest hot-unplug.
|
|
Without status_port this retains the legacy always-present behavior.
|
|
"""
|
|
missing = 0
|
|
while True:
|
|
try:
|
|
ready = True if status_port is None else probe_status(status_port)
|
|
qmp = Qmp(vm)
|
|
try:
|
|
channel, device = ids(vm)
|
|
present = device in peripherals(qmp)
|
|
if ready:
|
|
missing = 0
|
|
if not present:
|
|
attach(qmp, vm, port, bus, vendor, product)
|
|
else:
|
|
missing += 1
|
|
if present and missing >= remove_after:
|
|
detach(qmp, vm)
|
|
finally:
|
|
qmp.close()
|
|
except (OSError, RuntimeError) as exc:
|
|
print(f"waiting for VM/QMP: {exc}", flush=True)
|
|
time.sleep(interval)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("action", choices=["status", "attach", "detach", "watch"])
|
|
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("--any-storage", action="store_true",
|
|
help="Allow any pure USB Mass Storage device")
|
|
parser.add_argument("--status-port", type=int,
|
|
help="loopback port forwarded to the box readiness endpoint")
|
|
parser.add_argument("--interval", type=float, default=3,
|
|
help="watch polling interval in seconds")
|
|
parser.add_argument("--remove-after", type=int, default=3,
|
|
help="consecutive empty probes before guest removal")
|
|
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.status_port is not None and not 1 <= args.status_port <= 65535:
|
|
parser.error("invalid status port")
|
|
if args.interval <= 0 or args.remove_after < 1:
|
|
parser.error("interval must be positive and remove-after must be at least one")
|
|
if args.action in ("attach", "watch"):
|
|
exact = all(re.fullmatch(r"[0-9a-fA-F]{4}", x or "") for x in [args.vendor, args.product])
|
|
if args.any_storage and (args.vendor or args.product):
|
|
parser.error("--any-storage cannot be combined with --vendor/--product")
|
|
if not args.any_storage and not exact:
|
|
parser.error("attach/watch requires --any-storage or exact --vendor and --product")
|
|
elif args.any_storage or args.vendor or args.product:
|
|
parser.error("--vendor/--product/--any-storage are only valid for attach/watch")
|
|
if args.action == "watch":
|
|
watch(args.vm, args.port, args.bus,
|
|
None if args.any_storage else args.vendor.lower(),
|
|
None if args.any_storage else args.product.lower(),
|
|
interval=args.interval, status_port=args.status_port,
|
|
remove_after=args.remove_after)
|
|
return
|
|
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,
|
|
None if args.any_storage else args.vendor.lower(),
|
|
None if args.any_storage else 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()
|