From dc51ca580384b630b5e930e27f024995a43eb1a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Thalheim?= <joerg@thalheim.io>
Date: Tue, 19 Sep 2023 20:30:18 +0200
Subject: [PATCH] factor out deployment address parsing into a function

---
 pkgs/clan-cli/clan_cli/machines/update.py | 36 +++++------------------
 pkgs/clan-cli/clan_cli/ssh/__init__.py    | 30 +++++++++++++++++++
 2 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/pkgs/clan-cli/clan_cli/machines/update.py b/pkgs/clan-cli/clan_cli/machines/update.py
index c969a86..20d4228 100644
--- a/pkgs/clan-cli/clan_cli/machines/update.py
+++ b/pkgs/clan-cli/clan_cli/machines/update.py
@@ -2,13 +2,12 @@ import argparse
 import json
 import os
 import subprocess
-from typing import Optional
 
 from ..dirs import get_clan_flake_toplevel
 from ..nix import nix_command, nix_eval
 from ..secrets.generate import generate_secrets
 from ..secrets.upload import upload_secrets
-from ..ssh import Host, HostGroup, HostKeyCheck
+from ..ssh import Host, HostGroup, HostKeyCheck, parse_deployment_address
 
 
 def deploy_nixos(hosts: HostGroup) -> None:
@@ -78,11 +77,12 @@ def deploy_nixos(hosts: HostGroup) -> None:
 # FIXME: we want some kind of inventory here.
 def update(args: argparse.Namespace) -> None:
     clan_dir = get_clan_flake_toplevel().as_posix()
-    host = json.loads(
+    machine = args.machine
+    address = json.loads(
         subprocess.run(
             nix_eval(
                 [
-                    f'{clan_dir}#nixosConfigurations."{args.machine}".config.clan.networking.deploymentAddress'
+                    f'{clan_dir}#nixosConfigurations."{machine}".config.clan.networking.deploymentAddress'
                 ]
             ),
             stdout=subprocess.PIPE,
@@ -90,31 +90,9 @@ def update(args: argparse.Namespace) -> None:
             text=True,
         ).stdout
     )
-    parts = host.split("@")
-    user: Optional[str] = None
-    if len(parts) > 1:
-        user = parts[0]
-        hostname = parts[1]
-    else:
-        hostname = parts[0]
-    maybe_port = hostname.split(":")
-    port = None
-    if len(maybe_port) > 1:
-        hostname = maybe_port[0]
-        port = int(maybe_port[1])
-    print(f"deploying {host}")
-    deploy_nixos(
-        HostGroup(
-            [
-                Host(
-                    host=hostname,
-                    port=port,
-                    user=user,
-                    meta=dict(flake_attr=args.machine),
-                )
-            ]
-        )
-    )
+    host = parse_deployment_address(machine, address)
+    print(f"deploying {machine}")
+    deploy_nixos(HostGroup([host]))
 
 
 def register_update_parser(parser: argparse.ArgumentParser) -> None:
diff --git a/pkgs/clan-cli/clan_cli/ssh/__init__.py b/pkgs/clan-cli/clan_cli/ssh/__init__.py
index 5f1f8fe..735db17 100644
--- a/pkgs/clan-cli/clan_cli/ssh/__init__.py
+++ b/pkgs/clan-cli/clan_cli/ssh/__init__.py
@@ -756,6 +756,36 @@ class HostGroup:
         return HostGroup(list(filter(pred, self.hosts)))
 
 
+def parse_deployment_address(machine_name: str, host: str) -> Host:
+    parts = host.split("@")
+    user: Optional[str] = None
+    if len(parts) > 1:
+        user = parts[0]
+        hostname = parts[1]
+    else:
+        hostname = parts[0]
+    maybe_options = hostname.split("?")
+    options: Dict[str, str] = {}
+    if len(maybe_options) > 1:
+        hostname = maybe_options[0]
+        for option in maybe_options[1].split("&"):
+            k, v = option.split("=")
+            options[k] = v
+    maybe_port = hostname.split(":")
+    port = None
+    if len(maybe_port) > 1:
+        hostname = maybe_port[0]
+        port = int(maybe_port[1])
+    return Host(
+        hostname,
+        user=user,
+        port=port,
+        command_prefix=machine_name,
+        meta=dict(flake_attr=machine_name),
+        ssh_options=options,
+    )
+
+
 @overload
 def run(
     cmd: Union[List[str], str],