[PATCH 08/12] qemu_system_units: render three ssh_config stanzas with vsock default

From: Daniel Gomez <[email protected]>

ssh qsu wrote a single TCP-over-SLIRP stanza and nothing for
AF_VSOCK. Users got the VSOCK path only by typing the raw
vsock/<cid> URL with an explicit -i <key>; the test harness and
any muscle-memory ssh <vm> still paid SLIRP overhead even
though the guest had sshd-vsock.socket listening (via
systemd-ssh-generator when CONFIG_VIRTIO_VSOCKETS=y) and
machined knew the CID.

Teach update_ssh_config_nixos.py an optional --vsock-cid and
--default-transport pair. When --vsock-cid is given the script
emits three stanzas: <host>-vsock via systemd-ssh-proxy,
<host>-tcp via the SLIRP hostfwd, and the bare <host> grouped
with whichever --default-transport selects. Without --vsock-cid
the single-stanza layout stays, so the libvirt-backed NIXOS
path is unchanged. The -vsock/-tcp suffixes name the socket
family's transport: vsock (AF_VSOCK) and tcp (AF_INET TCP).

Generated-by: Claude AI
Signed-off-by: Daniel Gomez <[email protected]>
---
 scripts/update_ssh_config_nixos.py | 239 +++++++++++++++++++++++--------------
 1 file changed, 152 insertions(+), 87 deletions(-)

diff --git a/scripts/update_ssh_config_nixos.py b/scripts/update_ssh_config_nixos.py
index 206c69cc..8aca6133 100755
--- a/scripts/update_ssh_config_nixos.py
+++ b/scripts/update_ssh_config_nixos.py
@@ -8,121 +8,186 @@ via native QEMU virtualization (not libvirt). It handles both adding
 and removing SSH config entries.
 
 Usage:
-    update_ssh_config_nixos.py update <hostname> <host> <port> <user> <ssh_config> <privkey> <tag>
-    update_ssh_config_nixos.py remove <hostname> '' '' '' <ssh_config> '' <tag>
+    update_ssh_config_nixos.py update <hostname> <host> <port> <user>
+        <ssh_config> <privkey> <tag>
+        [--vsock-cid <N>] [--default-transport vsock|tcp]
+    update_ssh_config_nixos.py remove <hostname> '' '' '' <ssh_config>
+        '' <tag>
+
+When --vsock-cid is given, emit three Host stanzas:
+    <hostname>-vsock    : ProxyCommand systemd-ssh-proxy vsock/<cid> 22
+    <hostname>-tcp      : TCP via <host>:<port>
+    <hostname>          : alias of whichever --default-transport picks
+                          (defaults to vsock when --vsock-cid is set).
+
+Without --vsock-cid, one <hostname> stanza is emitted over TCP
+(backward-compatible with the libvirt NIXOS caller).
 """
 
+import argparse
 import os
-import sys
 import re
-from pathlib import Path
+import sys
 
 
-def update_ssh_config(
-    action, hostname, host_ip, port, username, ssh_config_path, ssh_key_path, tag
-):
-    """Update or remove SSH config entries for NixOS VMs.
-
-    Args:
-        action: 'update' to add/update entry, 'remove' to remove entry
-        hostname: VM hostname
-        host_ip: Host IP (usually localhost for NixOS VMs)
-        port: SSH port number
-        username: SSH username
-        ssh_config_path: Path to SSH config file
-        ssh_key_path: Path to SSH private key
-        tag: Tag to identify entries (e.g., 'NixOS VM')
+SYSTEMD_SSH_PROXY = "/usr/lib/systemd/systemd-ssh-proxy"
+
+
+def render_vsock_stanza(host_tokens, user, key, cid):
+    return (
+        f"Host {host_tokens}\n"
+        f"    User {user}\n"
+        f"    IdentityFile {key}\n"
+        f"    ProxyCommand {SYSTEMD_SSH_PROXY} vsock/{cid} 22\n"
+        f"    ProxyUseFdpass yes\n"
+        f"    CheckHostIP no\n"
+        f"    StrictHostKeyChecking no\n"
+        f"    UserKnownHostsFile /dev/null\n"
+        f"    LogLevel ERROR\n"
+    )
+
+
+def render_tcp_stanza(host_tokens, host_ip, port, user, key):
+    return (
+        f"Host {host_tokens}\n"
+        f"    HostName {host_ip}\n"
+        f"    Port {port}\n"
+        f"    User {user}\n"
+        f"    IdentityFile {key}\n"
+        f"    StrictHostKeyChecking no\n"
+        f"    UserKnownHostsFile /dev/null\n"
+        f"    LogLevel ERROR\n"
+    )
+
+
+def build_entries(hostname, host_ip, port, user, key, vsock_cid, default_transport):
+    """Return a list of (alias, block_body) pairs to write.
+
+    Each block starts with its own `# kdevops-managed: {tag} - {alias}`
+    marker so remove-on-update can match and drop the old block via
+    regex.
     """
+    if vsock_cid is None:
+        # Legacy single-stanza form used by libvirt NIXOS path.
+        return [
+            (hostname, render_tcp_stanza(hostname, host_ip, port, user, key)),
+        ]
+
+    # Three-stanza form: default alias groups with either -vsock or -tcp.
+    if default_transport == "tcp":
+        default_alias = f"{hostname} {hostname}-tcp"
+        bare_stanza = render_tcp_stanza(default_alias, host_ip, port, user, key)
+        return [
+            (hostname, bare_stanza),
+            (f"{hostname}-vsock",
+             render_vsock_stanza(f"{hostname}-vsock", user, key, vsock_cid)),
+        ]
+
+    # default_transport == "vsock"
+    default_alias = f"{hostname} {hostname}-vsock"
+    bare_stanza = render_vsock_stanza(default_alias, user, key, vsock_cid)
+    return [
+        (hostname, bare_stanza),
+        (f"{hostname}-tcp",
+         render_tcp_stanza(f"{hostname}-tcp", host_ip, port, user, key)),
+    ]
+
+
+def managed_block_regex(tag, hostname):
+    """Match every `# kdevops-managed: {tag} - <alias>` block whose alias
+    is either the base hostname or starts with it followed by a dash.
+    This catches qsu, qsu-vsock, qsu-tcp in one sweep without disturbing
+    unrelated entries that share the tag prefix.
+    """
+    return re.compile(
+        rf"^# kdevops-managed: {re.escape(tag)} - "
+        rf"{re.escape(hostname)}(?:-[\w-]+)?\n"
+        r"Host [^\n]+\n"
+        r"(?:[ \t]+[^\n]+\n)*",
+        re.MULTILINE,
+    )
 
-    ssh_config_path = os.path.expanduser(ssh_config_path)
 
-    # Ensure SSH config directory exists
+def update_ssh_config(
+    action, hostname, host_ip, port, username, ssh_config_path, ssh_key_path,
+    tag, vsock_cid, default_transport,
+):
+    ssh_config_path = os.path.expanduser(ssh_config_path)
     os.makedirs(os.path.dirname(ssh_config_path), exist_ok=True)
 
-    # Read existing config
     config_content = ""
     if os.path.exists(ssh_config_path):
         with open(ssh_config_path, "r") as f:
             config_content = f.read()
 
-    # Pattern to match our managed entries
-    entry_pattern = re.compile(
-        rf"^# kdevops-managed: {re.escape(tag)} - {re.escape(hostname)}\n"
-        r"Host [^\n]+\n"
-        r"(?:[ \t]+[^\n]+\n)*",
-        re.MULTILINE,
-    )
+    entry_pattern = managed_block_regex(tag, hostname)
 
-    if action == "remove":
-        # Remove existing entry
-        config_content = entry_pattern.sub("", config_content)
-        print(f"Removed SSH config entry for {hostname}")
+    # Always drop prior managed blocks for this base hostname.
+    config_content = entry_pattern.sub("", config_content)
 
+    if action == "remove":
+        print(f"Removed SSH config entries for {hostname}")
     elif action == "update":
-        # Remove existing entry first
-        config_content = entry_pattern.sub("", config_content)
-
-        # Create new entry
-        new_entry = f"""# kdevops-managed: {tag} - {hostname}
-Host {hostname}
-    HostName {host_ip}
-    Port {port}
-    User {username}
-    IdentityFile {ssh_key_path}
-    StrictHostKeyChecking no
-    UserKnownHostsFile /dev/null
-    LogLevel ERROR
+        blocks = []
+        for alias, body in build_entries(
+            hostname, host_ip, port, username, ssh_key_path,
+            vsock_cid, default_transport,
+        ):
+            blocks.append(
+                f"# kdevops-managed: {tag} - {alias}\n{body}"
+            )
+        new_entries = "\n".join(blocks)
+        config_content = config_content.rstrip() + "\n\n" + new_entries + "\n"
+        if vsock_cid is None:
+            print(f"Updated SSH config entry for {hostname} (port {port})")
+        else:
+            print(
+                f"Updated SSH config entries for {hostname}, "
+                f"{hostname}-vsock, {hostname}-tcp "
+                f"(default={default_transport}, vsock CID={vsock_cid}, "
+                f"TCP port={port})"
+            )
 
-"""
-
-        # Add new entry at the end
-        config_content = config_content.rstrip() + "\n\n" + new_entry
-        print(f"Updated SSH config entry for {hostname} (port {port})")
-
-    # Write updated config
     with open(ssh_config_path, "w") as f:
         f.write(config_content)
 
 
 def main():
-    """Main function to handle command line arguments."""
-    if len(sys.argv) < 8:
-        print(
-            "Usage: update_ssh_config_nixos.py <action> <hostname> <host> <port> <user> <ssh_config> <privkey> <tag>"
-        )
-        print("  action: 'update' or 'remove'")
-        print("  hostname: VM hostname")
-        print("  host: Host IP (use 'localhost' for local VMs)")
-        print("  port: SSH port number")
-        print("  user: SSH username")
-        print("  ssh_config: Path to SSH config file")
-        print("  privkey: Path to SSH private key")
-        print("  tag: Tag to identify entries (e.g., 'NixOS VM')")
-        sys.exit(1)
-
-    action = sys.argv[1]
-    hostname = sys.argv[2]
-    host_ip = sys.argv[3] if sys.argv[3] else "localhost"
-    port = sys.argv[4] if sys.argv[4] else "22"
-    username = sys.argv[5] if sys.argv[5] else "kdevops"
-    ssh_config_path = sys.argv[6]
-    ssh_key_path = sys.argv[7] if len(sys.argv) > 7 else ""
-    tag = sys.argv[8] if len(sys.argv) > 8 else "NixOS VM"
-
-    if action not in ["update", "remove"]:
-        print(f"Error: Invalid action '{action}'. Use 'update' or 'remove'.")
-        sys.exit(1)
+    parser = argparse.ArgumentParser(
+        description="Manage SSH config entries for NixOS-on-QEMU VMs.",
+    )
+    parser.add_argument("action", choices=["update", "remove"])
+    parser.add_argument("hostname")
+    parser.add_argument("host_ip", nargs="?", default="localhost")
+    parser.add_argument("port", nargs="?", default="22")
+    parser.add_argument("username", nargs="?", default="kdevops")
+    parser.add_argument("ssh_config_path")
+    parser.add_argument("ssh_key_path", nargs="?", default="")
+    parser.add_argument("tag", nargs="?", default="NixOS VM")
+    parser.add_argument(
+        "--vsock-cid", type=int, default=None,
+        help="VSOCK CID of the VM. When set, emit qsu-style three-stanza "
+             "layout (host, host-vsock, host-tcp).",
+    )
+    parser.add_argument(
+        "--default-transport", choices=["vsock", "tcp"], default="vsock",
+        help="Which transport the bare <hostname> alias groups with. Only "
+             "meaningful when --vsock-cid is set.",
+    )
+    args = parser.parse_args()
 
     try:
         update_ssh_config(
-            action,
-            hostname,
-            host_ip,
-            port,
-            username,
-            ssh_config_path,
-            ssh_key_path,
-            tag,
+            args.action,
+            args.hostname,
+            args.host_ip or "localhost",
+            args.port or "22",
+            args.username or "kdevops",
+            args.ssh_config_path,
+            args.ssh_key_path,
+            args.tag,
+            args.vsock_cid,
+            args.default_transport,
         )
     except Exception as e:
         print(f"Error: {e}")

-- 
2.53.0