vdi-cloud/tools/deploy-usb-device-poc.py

#!/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())