pyinfra 3.5.1__py3-none-any.whl → 3.6.1__py3-none-any.whl

pyinfra/connectors/sshuserclient/client.py

@@ -28,15 +28,17 @@ HOST_KEYS_LOCK = BoundedSemaphore()
 class StrictPolicy(MissingHostKeyPolicy):
     @override
     def missing_host_key(self, client, hostname, key):
-        logger.error("No host key for {0} found in known_hosts".format(hostname))
-        raise SSHException(
-            "StrictPolicy: No host key for {0} found in known_hosts".format(hostname),
-        )
+        logger.error("No host key for %s found in known_hosts", hostname)
+        raise SSHException(f"StrictPolicy: No host key for {hostname} found in known_hosts")
 
 
 def append_hostkey(client, hostname, key):
     """Append hostname to the clients host_keys_file"""
 
+    if client._host_keys_filename is None:
+        logger.warning("No host keys filename, not saving key for: %s", hostname)
+        return
+
     with HOST_KEYS_LOCK:
         # The paramiko client saves host keys incorrectly whereas the host keys object does
         # this correctly, so use that with the client filename variable.
@@ -67,30 +69,28 @@ class AcceptNewPolicy(MissingHostKeyPolicy):
         )
 
         append_hostkey(client, hostname, key)
-        logger.warning("Added host key for {0} to known_hosts".format(hostname))
+        logger.warning("Added host key for %s to known_hosts", hostname)
 
 
 class AskPolicy(MissingHostKeyPolicy):
     @override
     def missing_host_key(self, client, hostname, key):
         should_continue = input(
-            "No host key for {0} found in known_hosts, do you want to continue [y/n] ".format(
-                hostname,
-            ),
+            f"No host key for {hostname} found in known_hosts, do you want to continue [y/n] ",
         )
         if should_continue.lower() != "y":
             raise SSHException(
-                "AskPolicy: No host key for {0} found in known_hosts".format(hostname),
+                f"AskPolicy: No host key for {hostname} found in known_hosts",
             )
         append_hostkey(client, hostname, key)
-        logger.warning("Added host key for {0} to known_hosts".format(hostname))
+        logger.warning("Added host key for %s to known_hosts", hostname)
         return
 
 
 class WarningPolicy(MissingHostKeyPolicy):
     @override
     def missing_host_key(self, client, hostname, key):
-        logger.warning("No host key for {0} found in known_hosts".format(hostname))
+        logger.warning("No host key for %s found in known_hosts", hostname)
 
 
 def get_missing_host_key_policy(policy):
@@ -102,7 +102,7 @@ def get_missing_host_key_policy(policy):
         return StrictPolicy()
     if policy == "accept-new":
         return AcceptNewPolicy()
-    raise SSHException("Invalid value StrictHostKeyChecking={}".format(policy))
+    raise SSHException(f"Invalid value StrictHostKeyChecking={policy}")
 
 
 @memoize
@@ -120,17 +120,24 @@ def get_ssh_config(user_config_file=None):
 
 
 @memoize
-def get_host_keys(filename):
+def get_host_keys(filenames):
+    """
+    Load host keys from one or more files.
+
+    Args:
+        filenames: A tuple of filenames to load host keys from.
+    """
     with HOST_KEYS_LOCK:
         host_keys = HostKeys()
 
-        try:
-            host_keys.load(filename)
-        # When paramiko encounters a bad host keys line it sometimes bails the
-        # entire load incorrectly.
-        # See: https://github.com/paramiko/paramiko/pull/1990
-        except Exception as e:
-            logger.warning("Failed to load host keys from {0}: {1}".format(filename, e))
+        for filename in filenames:
+            try:
+                host_keys.load(filename)
+            # When paramiko encounters a bad host keys line it sometimes bails the
+            # entire load incorrectly.
+            # See: https://github.com/paramiko/paramiko/pull/1990
+            except Exception as e:
+                logger.warning("Failed to load host keys from %s: %s", filename, e)
 
         return host_keys
 
@@ -157,7 +164,7 @@ class SSHClient(ParamikoClient):
             config,
             forward_agent,
             missing_host_key_policy,
-            host_keys_file,
+            host_keys_files,
             keep_alive,
         ) = self.parse_config(
             hostname,
@@ -169,11 +176,13 @@ class SSHClient(ParamikoClient):
         config.update(kwargs)
 
         if _pyinfra_ssh_known_hosts_file:
-            host_keys_file = _pyinfra_ssh_known_hosts_file
+            host_keys_files = (path.expanduser(_pyinfra_ssh_known_hosts_file),)
 
         # Overwrite paramiko empty defaults with @memoize-d host keys object
-        self._host_keys = get_host_keys(host_keys_file)
-        self._host_keys_filename = host_keys_file
+        self._host_keys = get_host_keys(host_keys_files)
+        # Use the first file for writing new host keys
+        if len(host_keys_files) > 0:
+            self._host_keys_filename = host_keys_files[0]
 
         if _pyinfra_ssh_paramiko_connect_kwargs:
             config.update(_pyinfra_ssh_paramiko_connect_kwargs)
@@ -217,11 +226,18 @@ class SSHClient(ParamikoClient):
         keep_alive = 0
         forward_agent = False
         missing_host_key_policy = get_missing_host_key_policy(strict_host_key_checking)
-        host_keys_file = path.expanduser("~/.ssh/known_hosts")  # OpenSSH default
+        host_keys_files = (path.expanduser("~/.ssh/known_hosts"),)
 
         ssh_config = get_ssh_config(ssh_config_file)
         if not ssh_config:
-            return hostname, cfg, forward_agent, missing_host_key_policy, host_keys_file, keep_alive
+            return (
+                hostname,
+                cfg,
+                forward_agent,
+                missing_host_key_policy,
+                host_keys_files,
+                keep_alive,
+            )
 
         host_config = ssh_config.lookup(hostname)
         forward_agent = host_config.get("forwardagent") == "yes"
@@ -233,7 +249,10 @@ class SSHClient(ParamikoClient):
             )
 
         if "userknownhostsfile" in host_config:
-            host_keys_file = path.expanduser(host_config["userknownhostsfile"])
+            # OpenSSH supports multiple space-separated known hosts files
+            host_keys_files = tuple(
+                path.expanduser(f) for f in host_config["userknownhostsfile"].split()
+            )
 
         if "hostname" in host_config:
             hostname = host_config["hostname"]
@@ -275,7 +294,7 @@ class SSHClient(ParamikoClient):
                 sock = c.gateway(hostname, cfg["port"], target, target_config["port"])
             cfg["sock"] = sock
 
-        return hostname, cfg, forward_agent, missing_host_key_policy, host_keys_file, keep_alive
+        return hostname, cfg, forward_agent, missing_host_key_policy, host_keys_files, keep_alive
 
     @staticmethod
     def derive_shorthand(ssh_config, host_string):

pyinfra/connectors/util.py

@@ -22,17 +22,17 @@ if TYPE_CHECKING:
 
 
 SUDO_ASKPASS_ENV_VAR = "PYINFRA_SUDO_PASSWORD"
+
+
 SUDO_ASKPASS_COMMAND = r"""
-temp=$(mktemp "${{TMPDIR:=/tmp}}/pyinfra-sudo-askpass-XXXXXXXXXXXX")
+temp=$(mktemp "${{TMPDIR:={0}}}/pyinfra-sudo-askpass-XXXXXXXXXXXX")
 cat >"$temp"<<'__EOF__'
 #!/bin/sh
-printf '%s\n' "${0}"
+printf '%s\n' "${1}"
 __EOF__
 chmod 755 "$temp"
 echo "$temp"
-""".format(
-    SUDO_ASKPASS_ENV_VAR,
-)
+"""
 
 
 def run_local_process(
@@ -44,11 +44,14 @@ def run_local_process(
 ) -> tuple[int, "CommandOutput"]:
     process = Popen(command, shell=True, stdout=PIPE, stderr=PIPE, stdin=PIPE)
 
-    if stdin:
-        write_stdin(stdin, process.stdin)
-
     assert process.stdout is not None
     assert process.stderr is not None
+    assert process.stdin is not None
+
+    # Write any stdin and then close it
+    if stdin:
+        write_stdin(stdin, process.stdin)
+    process.stdin.close()
 
     combined_output = read_output_buffers(
         process.stdout,
@@ -264,7 +267,9 @@ def extract_control_arguments(arguments: "ConnectorArguments") -> "ConnectorArgu
 def _ensure_sudo_askpass_set_for_host(host: "Host"):
     if host.connector_data.get("sudo_askpass_path"):
         return
-    _, output = host.run_shell_command(SUDO_ASKPASS_COMMAND)
+    _, output = host.run_shell_command(
+        SUDO_ASKPASS_COMMAND.format(host.get_temp_dir_config(), SUDO_ASKPASS_ENV_VAR)
+    )
     host.connector_data["sudo_askpass_path"] = shlex.quote(output.stdout_lines[0])
 
 
@@ -318,6 +323,8 @@ def make_unix_command(
     _retries=0,
     _retry_delay=0,
     _retry_until=None,
+    # Temp dir config (ignored in command generation, used for temp file path generation)
+    _temp_dir=None,
 ) -> StringCommand:
     """
     Builds a shell command with various kwargs.

pyinfra/facts/crontab.py

@@ -52,9 +52,11 @@ class CrontabFile:
 
         name_comment = "# pyinfra-name={0}".format(name)
         for cmd in self.commands:
+            if "command" not in cmd:
+                continue
             if cmd.get("command") == command:
                 return cmd
-            if cmd.get("comments") and name_comment in cmd["comments"]:
+            if name_comment in cmd.get("comments", []):
                 return cmd
         return None
 
@@ -125,15 +127,15 @@ class Crontab(FactBase[CrontabFile]):
         # or CrontabFile.to_json()
         [
             {
-               command: "/path/to/command",
-               minute: "*",
-               hour: "*",
-               month: "*",
-               day_of_month: "*",
-               day_of_week: "*",
+                "command": "/path/to/command",
+                "minute": "*",
+                "hour": "*",
+                "month": "*",
+                "day_of_month": "*",
+                "day_of_week": "*",
             },
             {
-                "command": "echo another command
+                "command": "echo another command",
                 "special_time": "@daily",
             }
         ]

pyinfra/facts/files.py

@@ -547,6 +547,8 @@ class FindFilesBase(FactBase):
             command.append("-regex")
             command.append(maybe_quote(regex))
 
+        command.extend(args)
+
         command.append("||")
         command.append("true")
 
@@ -664,14 +666,21 @@ class Block(FactBase):
 
 class FileContents(FactBase):
     """
-    Returns the contents of a file as a list of lines. Works with both sha1sum and sha1. Returns
-    ``None`` if the file doest not exist.
+    Returns the contents of a file as a list of lines, or ``None`` if the file does not exist.
     """
 
     @override
     def command(self, path):
-        return make_formatted_string_command("cat {0}", QuoteString(path))
+        self.missing_flag = "{0}{1}".format(MISSING, path)
+        return make_formatted_string_command(
+            "( test -e {0} && cat {0} ) || echo {1}",
+            QuoteString(path),
+            QuoteString(self.missing_flag),
+        )
 
     @override
     def process(self, output):
+        # If output is the missing flag, the file doesn't exist
+        if output and output[0] == self.missing_flag:
+            return None
         return output

pyinfra/facts/flatpak.py

@@ -74,4 +74,4 @@ class FlatpakPackages(FlatpakBaseFact):
 
     @override
     def process(self, output):
-        return [flatpak for flatpak in output[1:]]
+        return [flatpak for i, flatpak in enumerate(output) if i > 1 or flatpak != "Application ID"]

pyinfra/facts/npm.py

@@ -30,7 +30,7 @@ class NpmPackages(FactBase):
     @override
     def command(self, directory=None):
         if directory:
-            return ("cd {0} && npm list -g --depth=0").format(directory)
+            return ("! test -d {0} || (cd {0} && npm list -g --depth=0)").format(directory)
         return "npm list -g --depth=0"
 
     @override

pyinfra/facts/server.py

@@ -52,12 +52,28 @@ class Path(FactBase):
 
 class TmpDir(FactBase):
     """
-    Returns the temporary directory of the current server, if configured.
+    Returns the temporary directory of the current server.
+
+    According to POSIX standards, checks environment variables in this order:
+    1. TMPDIR (if set and accessible)
+    2. TMP (if set and accessible)
+    3. TEMP (if set and accessible)
+    4. Falls back to empty string
     """
 
     @override
     def command(self):
-        return "echo $TMPDIR"
+        return """
+if [ -n "$TMPDIR" ] && [ -d "$TMPDIR" ] && [ -w "$TMPDIR" ]; then
+    echo "$TMPDIR"
+elif [ -n "$TMP" ] && [ -d "$TMP" ] && [ -w "$TMP" ]; then
+    echo "$TMP"
+elif [ -n "$TEMP" ] && [ -d "$TEMP" ] && [ -w "$TEMP" ]; then
+    echo "$TEMP"
+else
+    echo ""
+fi
+        """.strip()
 
 
 class Hostname(FactBase):

pyinfra/operations/apk.py

@@ -1,5 +1,5 @@
 """
-Manage apk packages.
+Manage apk packages. (Alpine Linux)
 """
 
 from __future__ import annotations
@@ -64,6 +64,7 @@ def packages(
 
     .. code:: python
 
+        from pyinfra.operations import apk
         # Update package list and install packages
         apk.packages(
             name="Install Asterisk and Vim",

pyinfra/operations/apt.py

@@ -4,7 +4,7 @@ Manage apt packages and repositories.
 
 from __future__ import annotations
 
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 from urllib.parse import urlparse
 
 from pyinfra import host
@@ -67,6 +67,7 @@ def key(src: str | None = None, keyserver: str | None = None, keyid: str | list[
 
     .. code:: python
 
+        from pyinfra.operations import apt
         # Note: If using URL, wget is assumed to be installed.
         apt.key(
             name="Add the Docker apt gpg key",
@@ -309,12 +310,19 @@ def update(cache_time: int | None = None):
         # Ubuntu provides this handy file
         cache_info = host.get_fact(File, path=APT_UPDATE_FILENAME)
 
-        # Time on files is not tz-aware, and will be the same tz as the server's time,
-        # so we can safely remove the tzinfo from the Date fact before comparison.
-        host_cache_time = host.get_fact(Date).replace(tzinfo=None) - timedelta(seconds=cache_time)
-        if cache_info and cache_info["mtime"] and cache_info["mtime"] > host_cache_time:
-            host.noop("apt is already up to date")
-            return
+        if cache_info and cache_info["mtime"]:
+            # The fact Date contains the date of the server in its timezone.
+            # cache_info["mtime"] ignores the timezone and consider the timestamp as UTC.
+            # So let's do the same here for the server current Date : ignore the
+            # timezone and consider it as UTC to have correct comparison with
+            # cache_info["mtime].
+            host_utc_current_time = datetime.fromtimestamp(
+                host.get_fact(Date).timestamp(), timezone.utc
+            ).replace(tzinfo=None)
+            host_cache_time = host_utc_current_time - timedelta(seconds=cache_time)
+            if cache_info["mtime"] > host_cache_time:
+                host.noop("apt is already up to date")
+                return
 
     yield "apt-get update"
 

pyinfra/operations/brew.py

@@ -61,6 +61,7 @@ def packages(
 
     .. code:: python
 
+        from pyinfra.operations import brew
         # Update package list and install packages
         brew.packages(
             name='Install Vim and vimpager',

pyinfra/operations/crontab.py

@@ -54,6 +54,7 @@ def crontab(
 
     .. code:: python
 
+        from pyinfra.operations import crontab
         # simple example for a crontab
         crontab.crontab(
             name="Backup /etc weekly",
@@ -84,7 +85,9 @@ def crontab(
         ctb = ctb0
     name_comment = "# pyinfra-name={0}".format(cron_name)
 
-    existing_crontab = ctb.get_command(command=command, name=cron_name)
+    existing_crontab = ctb.get_command(
+        command=command if cron_name is None else None, name=cron_name
+    )
     existing_crontab_command = existing_crontab["command"] if existing_crontab else command
     existing_crontab_match = existing_crontab["command"] if existing_crontab else command
 

pyinfra/operations/dnf.py

@@ -25,9 +25,12 @@ def key(src: str):
 
     .. code:: python
 
+        from pyinfra import host
+        from pyinfra.operations import dnf
+        from pyinfra.facts.server import LinuxDistribution
         linux_id = host.get_fact(LinuxDistribution)["release_meta"].get("ID")
         dnf.key(
-            name="Add the Docker CentOS gpg key",
+            name="Add the Docker gpg key",
             src=f"https://download.docker.com/linux/{linux_id}/gpg",
         )
 

pyinfra/operations/docker.py

@@ -8,9 +8,15 @@ from __future__ import annotations
 
 from pyinfra import host
 from pyinfra.api import operation
-from pyinfra.facts.docker import DockerContainer, DockerNetwork, DockerPlugin, DockerVolume
+from pyinfra.facts.docker import (
+    DockerContainer,
+    DockerImage,
+    DockerNetwork,
+    DockerPlugin,
+    DockerVolume,
+)
 
-from .util.docker import ContainerSpec, handle_docker
+from .util.docker import ContainerSpec, handle_docker, parse_image_reference
 
 
 @operation()
@@ -21,10 +27,13 @@ def container(
     networks: list[str] | None = None,
     volumes: list[str] | None = None,
     env_vars: list[str] | None = None,
+    labels: list[str] | None = None,
     pull_always: bool = False,
     present: bool = True,
     force: bool = False,
     start: bool = True,
+    restart_policy: str | None = None,
+    auto_remove: bool = False,
 ):
     """
     Manage Docker containers
@@ -35,15 +44,19 @@ def container(
     + ports: port list to expose
     + volumes: volume list to map on container
     + env_vars: environment variable list to inject on container
+    + labels: Label list to attach to the container
     + pull_always: force image pull
     + force: remove a container with same name and create a new one
     + present: whether the container should be up and running
     + start: start or stop the container
+    + restart_policy: restart policy to apply when a container exits
+    + auto_remove: automatically remove the container and its associated anonymous volumes when it exits
 
     **Examples:**
 
     .. code:: python
 
+        from pyinfra.operations import docker
         # Run a container
         docker.container(
             name="Deploy Nginx container",
@@ -55,6 +68,8 @@ def container(
             networks=["proxy", "services"],
             volumes=["nginx_data:/usr/share/nginx/html"],
             pull_always=True,
+            restart_policy="unless-stopped",
+            auto_remove=True,
         )
 
         # Stop a container
@@ -78,7 +93,10 @@ def container(
         networks or list(),
         volumes or list(),
         env_vars or list(),
+        labels or list(),
         pull_always,
+        restart_policy,
+        auto_remove,
     )
     existent_container = host.get_fact(DockerContainer, object_id=container)
 
@@ -127,13 +145,14 @@ def container(
         )
 
 
-@operation(is_idempotent=False)
-def image(image, present=True):
+@operation()
+def image(image: str, present: bool = True, force: bool = False):
     """
     Manage Docker images
 
     + image: Image and tag ex: nginx:alpine
     + present: whether the Docker image should exist
+    + force: always pull the image if present is True
 
     **Examples:**
 
@@ -153,20 +172,55 @@ def image(image, present=True):
             present=False,
         )
     """
-
+    image_info = parse_image_reference(image)
     if present:
-        yield handle_docker(
-            resource="image",
-            command="pull",
-            image=image,
-        )
-
+        if force:
+            # always pull the image if force is True
+            yield handle_docker(
+                resource="image",
+                command="pull",
+                image=image,
+            )
+            return
+        else:
+            existent_image = host.get_fact(DockerImage, object_id=image)
+            if image_info.digest:
+                # If a digest is specified, we must ensure the exact image is present
+                if existent_image:
+                    host.noop(f"Image with digest {image_info.digest} already exists!")
+                else:
+                    yield handle_docker(
+                        resource="image",
+                        command="pull",
+                        image=image,
+                    )
+            elif image_info.tag == "latest" or not image_info.tag:
+                # If the tag is 'latest' or not specified, always pull to ensure freshness
+                yield handle_docker(
+                    resource="image",
+                    command="pull",
+                    image=image,
+                )
+            else:
+                # For other tags, check if the image exists
+                if existent_image:
+                    host.noop(f"Image with tag {image_info.tag} already exists!")
+                else:
+                    yield handle_docker(
+                        resource="image",
+                        command="pull",
+                        image=image,
+                    )
     else:
-        yield handle_docker(
-            resource="image",
-            command="remove",
-            image=image,
-        )
+        existent_image = host.get_fact(DockerImage, object_id=image)
+        if existent_image:
+            yield handle_docker(
+                resource="image",
+                command="remove",
+                image=image,
+            )
+        else:
+            host.noop("There is no {0} image!".format(image))
 
 
 @operation()