From f0c26851e595cc62ac4222548f13bec9e32e2441 Mon Sep 17 00:00:00 2001
From: Adam Goldsmith <adam@adamgoldsmith.name>
Date: Wed, 19 Apr 2023 12:10:20 -0400
Subject: [PATCH] pass_typer: Try more candidate globs for user and host
 variants

---
 .config/i3/pass_typer.py | 112 ++++++++++++++++++++++++++++++---------
 1 file changed, 88 insertions(+), 24 deletions(-)

diff --git a/.config/i3/pass_typer.py b/.config/i3/pass_typer.py
index 70555ce..45f9fde 100755
--- a/.config/i3/pass_typer.py
+++ b/.config/i3/pass_typer.py
@@ -1,11 +1,14 @@
 #!/usr/bin/env python3
 
+import dataclasses
+import getpass
 import os
 import re
 import subprocess
+from collections.abc import Iterator
 from functools import partial
-from itertools import chain
 from pathlib import Path
+from typing import Optional
 
 PASSWORD_STORE = Path(
     os.environ.get("PASSWORD_STORE_DIR", Path("~/.password-store").expanduser())
@@ -14,18 +17,79 @@ PASSWORD_STORE = Path(
 cmd = partial(subprocess.run, capture_output=True, encoding="ascii")
 
 
-def alias(server_name):
-    match server_name.rpartition("@"):
-        case [user, at, ("salt" | "salt.sawtooth.claremontmakerspace.org")]:
-            return f"cms/{user}{at}cms-net-svcs"
-        case [user, at, host] if host.lower().startswith("cms-wall-display") or host.startswith("iPad"):
-            return f"cms/{user}{at}ipads"
-        case [user, at, ("octopi-taz-6" | "octopi-lulzbot-mini")]:
-            return f"cms/{user}{at}octopi"
-        case [user, at, server] if re.match(".*-pidgey(.vpn)?", server):
-            return "tsrc/pidgey"
-        case _:
-            return f"**/{server_name}"
+@dataclasses.dataclass
+class ServerMatch:
+    folder: Optional[str]
+    user: Optional[str]
+    host: str
+
+    @classmethod
+    def from_destination(cls, server_name: str) -> 'ServerMatch':
+        user, _, host = server_name.rpartition("@")
+        return cls(None, user if user != "" else None, host).alias()
+
+    def alias(self) -> 'ServerMatch':
+        # TODO: this still feels a bit hacky
+        host = self.host
+        if host.endswith(".sawtooth.claremontmakerspace.org"):
+            self.folder = "cms"
+            host = self.host.removesuffix(".sawtooth.claremontmakerspace.org")
+
+        if host == "salt":
+            self.folder = "cms"
+            self.host = "cms-net-svcs"
+        elif host.lower().startswith("cms-wall-display") or host.startswith("iPad"):
+            self.folder = "cms"
+            self.host = "ipads"
+        elif host in ("octopi-taz-6", "octopi-lulzbot-mini"):
+            self.folder = "cms"
+            self.host = "octopi"
+        elif re.match(".*-pidgey(.vpn)?", host):
+            self.folder = "tsrc"
+            self.host = "pidgey"
+
+        return self
+
+    def to_globs(self) -> Iterator[Path]:
+        path = Path("servers")
+        if self.folder is not None:
+            path = path / Path(self.folder)
+        else:
+            path = path / "**"
+
+        user_candidates = [None, getpass.getuser()]
+        if self.user is None:
+            user_candidates.append("*")
+        else:
+            user_candidates.insert(0, self.user)
+
+        # try host without each subdomain level
+        host_parts = self.host.split(".")
+        host_candidates = [
+            ".".join(host_parts[:i]) for i in range(len(host_parts), 0, -1)
+        ]
+        host_candidates.append(self.host + ".*")
+
+        # generate all combinations of user and host
+        for user in user_candidates:
+            for host in host_candidates:
+                if user is None:
+                    yield path / f"{host}.gpg"
+                else:
+                    yield path / f"{user}@{host}.gpg"
+
+    def to_paths(self) -> Iterator[Path]:
+        print("Candidate globs:")
+        for glob in self.to_globs():
+            print(glob)
+            yield from PASSWORD_STORE.glob(str(glob))
+
+    def __str__(self) -> str:
+        folder = self.folder if self.folder is not None else '**'
+        if self.user is not None and self.user != "":
+            return f"{folder}/{self.user}@{self.host}"
+        else:
+            return f"{folder}/{self.host}"
 
 
 def notify(summary: str, body: str) -> None:
@@ -44,15 +108,15 @@ def get_password(password_name: str) -> str:
     return password
 
 
-def select_and_type(server_name: Path) -> None:
-    path = ("servers" / server_name).with_name(server_name.name + ".gpg")
-    files = chain(
-        PASSWORD_STORE.glob(str(path)),
-        PASSWORD_STORE.glob(str(path.with_name("*@" + path.name))),
-    )
+def select_and_type(server_match: ServerMatch) -> None:
+    paths = server_match.to_paths()
+
     file_list = list(
-        dict.fromkeys(str(f.relative_to(PASSWORD_STORE).with_suffix("")) for f in files)
+        dict.fromkeys(str(f.relative_to(PASSWORD_STORE).with_suffix("")) for f in paths)
     )
+    if len(file_list) == 0:
+        notify(f"No matches found for '{server_match}'", "")
+        return
 
     selected = rofi_select(file_list)
     if selected:
@@ -63,13 +127,13 @@ def select_and_type(server_name: Path) -> None:
 
 window_name = cmd(["xdotool", "getactivewindow", "getwindowname"]).stdout.strip()
 
-ssh_match = re.search(":([^ ]+=[^ ]* )*(mosh|ssh) (?P<server>.*)", window_name)
+ssh_match = re.search(":([^ ]+=[^ ]* )*(mosh|ssh) (?P<destination>.*)", window_name)
 if ssh_match:
-    server_name = alias(ssh_match.group("server"))
+    server_match = ServerMatch.from_destination(ssh_match.group("destination"))
 
-    notify(f"Matched server '{server_name}'", f"Window name: {window_name}")
+    notify(f"Matched server '{server_match}'", f"Window name: {window_name}")
 
-    select_and_type(Path(server_name))
+    select_and_type(server_match)
 
 else:
     notify("Window name did not match any rules", f"Window name: {window_name}")