facts.server: add AuthorizedKeys, make user_authorized_keys idempotent (#1670)

Author: wowi42

src/pyinfra/facts/server.py

@@ -810,6 +810,39 @@ def process(self, output):
         return users
 
 
+class AuthorizedKeys(FactBase[List[str]]):
+    """
+    Returns the SSH public keys listed in a user's ``~/.ssh/authorized_keys`` file as a
+    list of full key strings. Empty lines and lines starting with ``#`` are skipped; the
+    file's order is preserved.
+
+    .. code:: python
+
+        [
+            "ssh-ed25519 AAAAC3Nz... user@host",
+            "ssh-rsa AAAAB3Nz... other@host",
+        ]
+    """
+
+    default = list
+
+    @override
+    def command(self, user: str, path: Optional[str] = None) -> str:
+        # Tilde expansion resolves the user's home without another fact round-trip.
+        target = path if path is not None else "~{0}/.ssh/authorized_keys".format(user)
+        return "cat {0} 2>/dev/null || true".format(target)
+
+    @override
+    def process(self, output: Iterable[str]) -> List[str]:
+        keys: List[str] = []
+        for raw in output:
+            line = raw.strip()
+            if not line or line.startswith("#"):
+                continue
+            keys.append(line)
+        return keys
+
+
 class LinuxDistributionDict(TypedDict):
     name: Optional[str]
     major: Optional[int]

src/pyinfra/operations/server.py

@@ -17,6 +17,7 @@
 from pyinfra.connectors.util import remove_any_sudo_askpass_file
 from pyinfra.facts.files import Directory, FindInFile, Link
 from pyinfra.facts.server import (
+    AuthorizedKeys,
     Groups,
     Home,
     Hostname,
@@ -801,34 +802,49 @@ def read_any_pub_key_file(key):
 
     authorized_key_file = f"{authorized_key_directory}/{authorized_key_filename}"
 
-    if delete_keys:
-        # Create a whole new authorized_keys file
-        keys_file = StringIO(
-            "{0}\n".format(
-                "\n".join(public_keys),
-            ),
-        )
+    # Pull the currently installed keys once; individual files.line calls otherwise
+    # issue one FindInFile fact per key, which dominates the cost for users with many
+    # keys.
+    current_keys = host.get_fact(AuthorizedKeys, user=user, path=authorized_key_file)
 
-        # And ensure it exists
-        yield from files.put._inner(
-            src=keys_file,
-            dest=authorized_key_file,
-            user=user,
-            group=group or user,
-            mode=600,
-        )
+    if delete_keys:
+        if current_keys == public_keys:
+            # Still ensure the file and its ownership/mode stay correct.
+            yield from files.file._inner(
+                path=authorized_key_file,
+                user=user,
+                group=group or user,
+                mode=600,
+            )
+        else:
+            keys_file = StringIO(
+                "{0}\n".format(
+                    "\n".join(public_keys),
+                ),
+            )
+            yield from files.put._inner(
+                src=keys_file,
+                dest=authorized_key_file,
+                user=user,
+                group=group or user,
+                mode=600,
+            )
 
     else:
-        # Ensure authorized_keys exists
+        # Ensure authorized_keys exists with the right ownership and mode.
         yield from files.file._inner(
             path=authorized_key_file,
             user=user,
             group=group or user,
             mode=600,
         )
 
-        # And every public key is present
+        # Only append the keys that the fact says are missing; an empty fact result
+        # also covers the "file does not exist yet" case.
+        current_key_set = set(current_keys)
         for key in public_keys:
+            if key in current_key_set:
+                continue
             yield from files.line._inner(path=authorized_key_file, line=key, ensure_newline=True)
 
 

tests/facts/server.AuthorizedKeys/custom_path.json

@@ -0,0 +1,9 @@
+{
+    "arg": {
+        "user": "alice",
+        "path": "/etc/ssh/keys/alice"
+    },
+    "command": "cat /etc/ssh/keys/alice 2>/dev/null || true",
+    "output": [],
+    "fact": []
+}

tests/facts/server.AuthorizedKeys/default_path.json

@@ -0,0 +1,17 @@
+{
+    "arg": {
+        "user": "alice"
+    },
+    "command": "cat ~alice/.ssh/authorized_keys 2>/dev/null || true",
+    "output": [
+        "# some comment",
+        "",
+        "ssh-ed25519 AAAAC3Nz...key1 user@host",
+        "   ",
+        "ssh-rsa AAAAB3Nz...key2 other@host"
+    ],
+    "fact": [
+        "ssh-ed25519 AAAAC3Nz...key1 user@host",
+        "ssh-rsa AAAAB3Nz...key2 other@host"
+    ]
+}

tests/operations/server.user/key_files.json

@@ -43,6 +43,9 @@
             "interpolate_variables=False, path=homedir/.ssh/authorized_keys, pattern=^.*somekeydata.*$": ["somekeydata"],
             "interpolate_variables=False, path=homedir/.ssh/authorized_keys, pattern=^.*someotherkeydata.*$": []
         },
+        "server.AuthorizedKeys": {
+            "path=homedir/.ssh/authorized_keys, user=someuser": ["somekeydata"]
+        },
         "server.Groups": {}
     },
     "commands": [

tests/operations/server.user/keys.json

@@ -34,6 +34,9 @@
         "files.FindInFile": {
             "interpolate_variables=False, path=homedir/.ssh/authorized_keys, pattern=^.*abc.*$": []
         },
+        "server.AuthorizedKeys": {
+            "path=homedir/.ssh/authorized_keys, user=someuser": []
+        },
         "server.Groups": {}
     },
     "commands": [

tests/operations/server.user/keys_delete.json

@@ -45,7 +45,10 @@
             "path=homedir/.ssh/authorized_keys": null
         },
         "files.FileContents": {
-            "path=homedir/.ssh/authorized_keys": null 
+            "path=homedir/.ssh/authorized_keys": null
+        },
+        "server.AuthorizedKeys": {
+            "path=homedir/.ssh/authorized_keys, user=someuser": []
         },
         "server.Groups": {}
     },

tests/operations/server.user/keys_nohome.json

@@ -33,6 +33,9 @@
         "files.FindInFile": {
             "interpolate_variables=False, path=/root/.ssh/authorized_keys, pattern=^.*abc.*$": []
         },
+        "server.AuthorizedKeys": {
+            "path=/root/.ssh/authorized_keys, user=root": []
+        },
         "server.Groups": {}
     },
     "commands": [

tests/operations/server.user/keys_single.json

@@ -34,6 +34,9 @@
         "files.FindInFile": {
             "interpolate_variables=False, path=homedir/.ssh/authorized_keys, pattern=^.*abc.*$": []
         },
+        "server.AuthorizedKeys": {
+            "path=homedir/.ssh/authorized_keys, user=someuser": []
+        },
         "server.Groups": {}
     },
     "commands": [

tests/operations/server.user_authorized_keys/all_keys_present.json

@@ -0,0 +1,33 @@
+{
+    "args": ["someuser"],
+    "kwargs": {
+        "public_keys": ["ssh-ed25519 AAAAkey1 alice", "ssh-rsa AAAAkey2 bob"]
+    },
+    "facts": {
+        "server.Home": {
+            "user=someuser": "/home/someuser"
+        },
+        "files.Directory": {
+            "path=/home/someuser/.ssh": {
+                "user": "someuser",
+                "group": "someuser",
+                "mode": 700
+            }
+        },
+        "files.File": {
+            "path=/home/someuser/.ssh/authorized_keys": {
+                "user": "someuser",
+                "group": "someuser",
+                "mode": 600
+            }
+        },
+        "server.AuthorizedKeys": {
+            "path=/home/someuser/.ssh/authorized_keys, user=someuser": [
+                "ssh-ed25519 AAAAkey1 alice",
+                "ssh-rsa AAAAkey2 bob"
+            ]
+        }
+    },
+    "commands": [],
+    "noop_description": "file /home/someuser/.ssh/authorized_keys already exists"
+}