#!/usr/bin/env python3 """Install USB Mass Storage redirection using a restricted SSH tunnel. Requires paramiko on the admin machine and rustc/cargo/usbredirect on the box. PVE authentication uses SSH keys or VDI_PVE_SSH_PASSWORD (never stored in files). Use --any-storage to accept any pure USB Mass Storage device, or provide an exact VID/PID/serial identity for the legacy single-device policy. The deployment installs a QMP watcher so the runtime USB endpoint is restored after a VM/PVE restart. """ import argparse import json import os from pathlib import Path import re import shlex import subprocess import tempfile import time def box_run(box, command, stdin=None): result = subprocess.run( ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", box, command], input=stdin, capture_output=True, text=True, encoding="utf-8", timeout=120) if result.returncode: raise RuntimeError(result.stderr + result.stdout) return result.stdout def pve_run(client, command): _, out, err = client.exec_command(command, timeout=60) stdout, stderr = out.read().decode(), err.read().decode() if out.channel.recv_exit_status(): raise RuntimeError(stderr + stdout) return stdout def install(args, password=None): import paramiko exact = all(getattr(args, name, None) for name in ("vendor", "product", "serial")) if args.any_storage and exact: raise ValueError("--any-storage cannot be combined with an exact device identity") if not args.any_storage and not exact: raise ValueError("choose --any-storage or provide --vendor, --product and --serial") if exact: if not re.fullmatch(r"[0-9a-f]{4}", args.vendor) or not re.fullmatch(r"[0-9a-f]{4}", args.product): raise ValueError("VID/PID must be four lowercase hex digits") if not re.fullmatch(r"[A-Za-z0-9_.-]+", args.serial) or not args.serial: raise ValueError("installer supports nonempty alphanumeric serials only") if not re.fullmatch(r"[a-z_][a-z0-9_-]*@[A-Za-z0-9.-]+", args.box): raise ValueError("invalid box SSH destination") if args.vm < 100: raise ValueError("invalid VM ID") if box_run(args.box, "systemctl is-active vdi-usb-monitor.service || true").strip() == "active": raise RuntimeError("Safely eject the guest USB and stop the box export service before redeployment") repo = Path(__file__).resolve().parents[1] box_host = args.box.split("@", 1)[1] # Copy the existing trusted host keys; never trust an unauthenticated keyscan. known = subprocess.run(["ssh-keygen", "-F", box_host], capture_output=True, text=True, check=True).stdout known = "\n".join(line for line in known.splitlines() if not line.startswith("#")) + "\n" if not known.strip(): raise RuntimeError("Box host key must already be trusted by OpenSSH") client = paramiko.SSHClient() client.load_system_host_keys(str(Path.home() / ".ssh" / "known_hosts")) client.set_missing_host_key_policy(paramiko.RejectPolicy()) client.connect(args.pve, username="root", password=password or os.getenv("VDI_PVE_SSH_PASSWORD"), timeout=10, allow_agent=True, look_for_keys=True) try: stamp = time.strftime("%Y%m%d-%H%M%S") host_dir = f"/etc/vdi-usb-poc/vm{args.vm}" pve_run(client, f"install -d -m 700 {host_dir}; test -f {host_dir}/id_ed25519 || " f"ssh-keygen -q -t ed25519 -N '' -C vdi-usb-poc-vm{args.vm} -f {host_dir}/id_ed25519") pubkey = pve_run(client, f"cat {host_dir}/id_ed25519.pub").strip() backup = f"/home/{args.box.split('@')[0]}/vdi-terminal/backups/usb-device-{stamp}" # append a restricted tunnel-only key while preserving existing access. auth_script = '''import pathlib,shutil,sys p=pathlib.Path.home()/'.ssh'/'authorized_keys' b=pathlib.Path(sys.argv[1]);b.mkdir(parents=True,exist_ok=True) if p.exists():shutil.copy2(p,b/'authorized_keys') p.parent.mkdir(mode=0o700,exist_ok=True) text=p.read_text() if p.exists() else '' line=sys.argv[2] if line not in text.splitlines():p.write_text(text.rstrip()+'\\n'+line+'\\n') p.chmod(0o600) ''' auth_line = (f'from="{args.pve}",restrict,port-forwarding,permitopen="127.0.0.1:40001",' f'permitopen="127.0.0.1:40002",' f'command="/bin/false" {pubkey}') box_run(args.box, f"python3 - {shlex.quote(backup)} {shlex.quote(auth_line)}", auth_script) stage = box_run(args.box, "mktemp -d /tmp/vdi-usb-deploy.XXXXXX").strip() service_text = (repo / "terminal/usb-monitor/etc/vdi-usb-monitor.service").read_text() if args.any_storage: service_exec = ("ExecStart=/usr/local/libexec/vdi-usb-monitor serve " "--any-storage --listen 127.0.0.1:40001 --status 127.0.0.1:40002") else: service_exec = (f"ExecStart=/usr/local/libexec/vdi-usb-monitor serve --vendor {args.vendor} " f"--product {args.product} --serial {args.serial} --listen 127.0.0.1:40001 " "--status 127.0.0.1:40002") service_text = re.sub(r"^ExecStart=.*$", service_exec, service_text, flags=re.MULTILINE) service_fd, service_path = tempfile.mkstemp(prefix="vdi-usb-monitor-", suffix=".service") os.close(service_fd) service_tmp = Path(service_path) service_tmp.write_text(service_text) try: subprocess.run(["scp", "-q", str(service_tmp), f"{args.box}:{stage}/vdi-usb-monitor.service"], check=True) finally: service_tmp.unlink(missing_ok=True) if args.any_storage: env = 'VDI_USB_MODE=any-storage\n' rule = ('# VDI USB: ignore USB disks locally; the monitor checks exclusive use\n' 'SUBSYSTEM=="block", ENV{ID_BUS}=="usb", ENV{ID_TYPE}=="disk", ' 'ENV{UDISKS_IGNORE}="1"\n') else: env = f'VDI_USB_VENDOR={args.vendor}\nVDI_USB_PRODUCT={args.product}\nVDI_USB_SERIAL={args.serial}\n' rule = (f'# VDI USB device: only the approved device\nSUBSYSTEM=="block", ' f'ATTRS{{idVendor}}=="{args.vendor}", ATTRS{{idProduct}}=="{args.product}", ' f'ATTRS{{serial}}=="{args.serial}", ENV{{UDISKS_IGNORE}}="1"\n') box_install = '''import pathlib,shutil,subprocess,sys stage,backup,env,rule=sys.argv[1:] backup=pathlib.Path(backup) files={ '/etc/systemd/system/vdi-usb-monitor.service':pathlib.Path(stage,'vdi-usb-monitor.service').read_bytes(), '/etc/vdi-usb-monitor/device.env':env.encode(), '/etc/udev/rules.d/99-vdi-usb-device-poc.rules':rule.encode(), } for name,data in files.items(): p=pathlib.Path(name) if p.exists():shutil.copy2(p,backup/p.name) p.parent.mkdir(parents=True,exist_ok=True) p.write_bytes(data) p.chmod(0o600 if p.suffix=='.env' else 0o644) binary=pathlib.Path('/usr/local/libexec/vdi-usb-monitor') if binary.exists():shutil.copy2(binary,backup/binary.name) binary.parent.mkdir(parents=True,exist_ok=True) shutil.copy2(pathlib.Path.home()/'vdi-build/usb-monitor/target/release/vdi-usb-monitor',binary) binary.chmod(0o755) subprocess.run(['udevadm','control','--reload-rules'],check=True) subprocess.run(['systemctl','daemon-reload'],check=True) subprocess.run(['systemctl','enable','vdi-usb-monitor.service'],check=True) print('Prepared USB service; not started; existing mounted data untouched') ''' # sudo changes HOME; use the explicitly selected build path instead. box_install = box_install.replace("pathlib.Path.home()/'vdi-build", f"pathlib.Path('/home/{args.box.split('@')[0]}')/'vdi-build") box_run(args.box, "sudo -n python3 - " + " ".join(map(shlex.quote, [stage, backup, env, rule])), box_install) unit = f'''[Unit] Description=VDI USB POC encrypted tunnel for VM {args.vm} After=network-online.target Wants=network-online.target [Service] Type=simple ExecStart=/usr/bin/ssh -NT -i {host_dir}/id_ed25519 -o IdentitiesOnly=yes -o BatchMode=yes -o StrictHostKeyChecking=yes -o UserKnownHostsFile={host_dir}/known_hosts -o ExitOnForwardFailure=yes -o ServerAliveInterval=10 -o ServerAliveCountMax=3 -L 127.0.0.1:40001:127.0.0.1:40001 -L 127.0.0.1:40002:127.0.0.1:40002 {args.box} Restart=on-failure RestartSec=3 NoNewPrivileges=yes PrivateTmp=yes [Install] WantedBy=multi-user.target ''' qmp_args = ("--any-storage" if args.any_storage else f"--vendor {args.vendor} --product {args.product}") qmp_unit = f'''[Unit] Description=VDI USB QMP endpoint for VM {args.vm} After=network-online.target vdi-usb-tunnel-{args.vm}.service Wants=network-online.target [Service] Type=simple ExecStart=/usr/local/libexec/vdi-usb-device-qmp.py watch --vm {args.vm} --port 40001 --status-port 40002 --interval 1 --remove-after 3 --bus ehci-2.0 {qmp_args} Restart=always RestartSec=3 NoNewPrivileges=yes PrivateTmp=yes [Install] WantedBy=multi-user.target ''' sftp = client.open_sftp() try: for path, content in [(f"{host_dir}/known_hosts", known), (f"/etc/systemd/system/vdi-usb-tunnel-{args.vm}.service", unit), (f"/etc/systemd/system/vdi-usb-qmp-{args.vm}.service", qmp_unit), ("/usr/local/libexec/vdi-usb-device-qmp.py", (repo / "tools/usb-device-qmp.py").read_text())]: with sftp.open(path, "w") as f: f.write(content) mode = 0o600 if "known_hosts" in path else (0o755 if path.endswith("usb-device-qmp.py") else 0o644) sftp.chmod(path, mode) finally: sftp.close() pve_run(client, f"systemctl daemon-reload && systemctl enable --now vdi-usb-tunnel-{args.vm}.service vdi-usb-qmp-{args.vm}.service") print(json.dumps({"box_backup": backup, "pve_directory": host_dir, "state": "prepared; unmount USB safely before starting export"}, indent=2)) finally: client.close() if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--pve", required=True) parser.add_argument("--box", required=True) parser.add_argument("--vm", type=int, required=True) parser.add_argument("--any-storage", action="store_true", help="accept any pure USB Mass Storage device") parser.add_argument("--vendor") parser.add_argument("--product") parser.add_argument("--serial") install(parser.parse_args())