156 lines
7.7 KiB
Python
156 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Install a one-device usbredir POC 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).
|
|
This prepares services only. Safely unmount the selected USB volume, start the
|
|
box service, then run usb-device-qmp.py attach on the PVE node.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shlex
|
|
import subprocess
|
|
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
|
|
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("POC 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'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()
|
|
subprocess.run(["scp", "-q", str(repo / "terminal/usb-monitor/etc/vdi-usb-monitor.service"),
|
|
f"{args.box}:{stage}/vdi-usb-monitor.service"], check=True)
|
|
env = f'VDI_USB_VENDOR={args.vendor}\nVDI_USB_PRODUCT={args.product}\nVDI_USB_SERIAL={args.serial}\n'
|
|
rule = (f'# VDI USB device POC: 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 {args.box}
|
|
Restart=on-failure
|
|
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),
|
|
("/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)
|
|
sftp.chmod(path, 0o600 if "known_hosts" in path else 0o644)
|
|
finally:
|
|
sftp.close()
|
|
pve_run(client, f"systemctl daemon-reload && systemctl enable --now vdi-usb-tunnel-{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("--vendor", required=True)
|
|
parser.add_argument("--product", required=True)
|
|
parser.add_argument("--serial", required=True)
|
|
install(parser.parse_args())
|