pyinfra 3.1.1__py2.py3-none-any.whl → 3.3__py2.py3-none-any.whl

pyinfra/connectors/ssh.py

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple
 
 import click
 from paramiko import AuthenticationException, BadHostKeyException, SFTPClient, SSHException
-from typing_extensions import TypedDict, Unpack
+from typing_extensions import TypedDict, Unpack, override
 
 from pyinfra import logger
 from pyinfra.api.command import QuoteString, StringCommand
@@ -191,6 +191,7 @@ class SSHConnector(BaseConnector):
 
         return kwargs
 
+    @override
     def connect(self) -> None:
         retries = self.data["ssh_connect_retries"]
 
@@ -264,9 +265,11 @@ class SSHConnector(BaseConnector):
                 f"Host key for {e.hostname} does not match.",
             )
 
+    @override
     def disconnect(self) -> None:
         self.get_sftp_connection.cache.clear()
 
+    @override
     def run_shell_command(
         self,
         command: StringCommand,
@@ -368,6 +371,7 @@ class SSHConnector(BaseConnector):
             sftp = self.get_sftp_connection()
             sftp.getfo(remote_filename, file_io)
 
+    @override
     def get_file(
         self,
         remote_filename: str,
@@ -454,6 +458,7 @@ class SSHConnector(BaseConnector):
         if last_e is not None:
             raise last_e
 
+    @override
     def put_file(
         self,
         filename_or_io,
@@ -537,7 +542,8 @@ class SSHConnector(BaseConnector):
 
         return True
 
-    def check_can_rsync(self):
+    @override
+    def check_can_rsync(self) -> None:
         if self.data["ssh_key_password"]:
             raise NotImplementedError(
                 "Rsync does not currently work with SSH keys needing passwords."
@@ -549,6 +555,7 @@ class SSHConnector(BaseConnector):
         if not which("rsync"):
             raise NotImplementedError("The `rsync` binary is not available on this system.")
 
+    @override
     def rsync(
         self,
         src: str,

pyinfra/connectors/sshuserclient/client.py

@@ -14,6 +14,8 @@ from paramiko import (
     SSHException,
 )
 from paramiko.agent import AgentRequestHandler
+from paramiko.hostkeys import HostKeyEntry
+from typing_extensions import override
 
 from pyinfra import logger
 from pyinfra.api.util import memoize
@@ -24,6 +26,7 @@ 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(
@@ -31,7 +34,30 @@ class StrictPolicy(MissingHostKeyPolicy):
         )
 
 
+def append_hostkey(client, hostname, key):
+    """Append hostname to the clients host_keys_file"""
+
+    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.
+        # See: https://github.com/paramiko/paramiko/pull/1989
+        host_key_entry = HostKeyEntry([hostname], key)
+        if host_key_entry is None:
+            raise SSHException(
+                "Append Hostkey: Failed to parse host {0}, could not append to hostfile".format(
+                    hostname
+                ),
+            )
+        with open(client._host_keys_filename, "a") as host_keys_file:
+            hk_entry = host_key_entry.to_line()
+            if hk_entry is None:
+                raise SSHException(f"Append Hostkey: Failed to append hostkey ({host_key_entry})")
+
+            host_keys_file.write(hk_entry)
+
+
 class AcceptNewPolicy(MissingHostKeyPolicy):
+    @override
     def missing_host_key(self, client, hostname, key):
         logger.warning(
             (
@@ -40,16 +66,12 @@ class AcceptNewPolicy(MissingHostKeyPolicy):
             ),
         )
 
-        with HOST_KEYS_LOCK:
-            host_keys = client.get_host_keys()
-            host_keys.add(hostname, key.get_name(), key)
-            # The paramiko client saves host keys incorrectly whereas the host keys object does
-            # this correctly, so use that with the client filename variable.
-            # See: https://github.com/paramiko/paramiko/pull/1989
-            host_keys.save(client._host_keys_filename)
+        append_hostkey(client, hostname, key)
+        logger.warning("Added host key for {0} to known_hosts".format(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(
@@ -60,18 +82,13 @@ class AskPolicy(MissingHostKeyPolicy):
             raise SSHException(
                 "AskPolicy: No host key for {0} found in known_hosts".format(hostname),
             )
-        with HOST_KEYS_LOCK:
-            host_keys = client.get_host_keys()
-            host_keys.add(hostname, key.get_name(), key)
-            # The paramiko client saves host keys incorrectly whereas the host keys object does
-            # this correctly, so use that with the client filename variable.
-            # See: https://github.com/paramiko/paramiko/pull/1989
-            host_keys.save(client._host_keys_filename)
+        append_hostkey(client, hostname, key)
         logger.warning("Added host key for {0} to known_hosts".format(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))
 
@@ -124,6 +141,7 @@ class SSHClient(ParamikoClient):
     original idea at http://bitprophet.org/blog/2012/11/05/gateway-solutions/.
     """
 
+    @override
     def connect(  # type: ignore[override]
         self,
         hostname,
@@ -165,6 +183,13 @@ class SSHClient(ParamikoClient):
         if _pyinfra_ssh_forward_agent is not None:
             forward_agent = _pyinfra_ssh_forward_agent
 
+        keep_alive = config.get("keep_alive")
+
+        if keep_alive:
+            transport = self.get_transport()
+            assert transport is not None, "No transport"
+            transport.set_keepalive(keep_alive)
+
         if forward_agent:
             transport = self.get_transport()
             assert transport is not None, "No transport"
@@ -222,6 +247,9 @@ class SSHClient(ParamikoClient):
         if "port" in host_config:
             cfg["port"] = int(host_config["port"])
 
+        if "serveraliveinterval" in host_config:
+            cfg["keep_alive"] = int(host_config["serveraliveinterval"])
+
         if "proxycommand" in host_config:
             cfg["sock"] = ProxyCommand(host_config["proxycommand"])
 

pyinfra/connectors/sshuserclient/config.py

@@ -10,6 +10,7 @@ from os import environ, path
 import paramiko.config
 from gevent.subprocess import CalledProcessError, check_call
 from paramiko import SSHConfig as ParamikoSSHConfig
+from typing_extensions import override
 
 from pyinfra import logger
 
@@ -95,6 +96,7 @@ class SSHConfig(ParamikoSSHConfig):
     https://github.com/paramiko/paramiko/pull/1194
     """
 
+    @override
     def parse(self, file_obj):
         file_obj = _expand_include_statements(file_obj)
         return super().parse(file_obj)

pyinfra/connectors/terraform.py

@@ -9,7 +9,7 @@ from .base import BaseConnector
 
 
 @memoize
-def show_warning():
+def show_warning() -> None:
     logger.warning("The @terraform connector is in beta!")
 
 

pyinfra/connectors/util.py

@@ -198,14 +198,18 @@ def execute_command_with_sudo_retry(
 ) -> tuple[int, CommandOutput]:
     return_code, output = execute_command()
 
+    # If we failed look for a sudo password prompt line and re-submit using the sudo password. Look
+    # at all lines here in case anything else gets printed, eg in:
+    # https://github.com/pyinfra-dev/pyinfra/issues/1292
     if return_code != 0 and output and output.combined_lines:
-        last_line = output.combined_lines[-1].line
-        if last_line.strip() == "sudo: a password is required":
-            # If we need a password, ask the user for it and attach to the host
-            # internal connector data for use when executing future commands.
-            sudo_password = getpass("{0}sudo password: ".format(host.print_prefix))
-            host.connector_data["prompted_sudo_password"] = sudo_password
-            return_code, output = execute_command()
+        for line in reversed(output.combined_lines):
+            if line.line.strip() == "sudo: a password is required":
+                # If we need a password, ask the user for it and attach to the host
+                # internal connector data for use when executing future commands.
+                sudo_password = getpass("{0}sudo password: ".format(host.print_prefix))
+                host.connector_data["prompted_sudo_password"] = sudo_password
+                return_code, output = execute_command()
+                break
 
     return return_code, output
 
@@ -232,7 +236,7 @@ def remove_any_sudo_askpass_file(host) -> None:
 
 
 @memoize
-def _show_use_su_login_warning():
+def _show_use_su_login_warning() -> None:
     logger.warning(
         (
             "Using `use_su_login` may not work: "
@@ -304,7 +308,7 @@ def make_unix_command(
     _sudo=False,
     _sudo_user=None,
     _use_sudo_login=False,
-    _sudo_password=False,
+    _sudo_password="",
     _sudo_askpass_path=None,
     _preserve_sudo_env=False,
     # Doas config

pyinfra/context.py

@@ -10,6 +10,7 @@ from types import ModuleType
 from typing import TYPE_CHECKING
 
 from gevent.local import local
+from typing_extensions import override
 
 if TYPE_CHECKING:
     from pyinfra.api.config import Config
@@ -26,22 +27,25 @@ class ContextObject:
     _container_cls = container
     _base_cls: ModuleType
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._container = self._container_cls()
         self._container.module = None
 
     def _get_module(self):
         return self._container.module
 
+    @override
     def __repr__(self):
         return "ContextObject({0}):{1}".format(
             self._base_cls.__name__,
             repr(self._get_module()),
         )
 
+    @override
     def __str__(self):
         return str(self._get_module())
 
+    @override
     def __dir__(self):
         return dir(self._base_cls)
 
@@ -50,6 +54,7 @@ class ContextObject:
             return getattr(self._base_cls, key)
         return getattr(self._get_module(), key)
 
+    @override
     def __setattr__(self, key, value):
         if key in ("_container", "_base_cls"):
             return super().__setattr__(key, value)
@@ -65,9 +70,11 @@ class ContextObject:
     def __len__(self):
         return len(self._get_module())
 
+    @override
     def __eq__(self, other):
         return self._get_module() == other
 
+    @override
     def __hash__(self):
         return hash(self._get_module())
 
@@ -89,7 +96,7 @@ class ContextManager:
     def set_base(self, module):
         self.context._base_cls = module
 
-    def reset(self):
+    def reset(self) -> None:
         self.context._container.module = None
 
     def isset(self):

pyinfra/facts/apk.py

@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+from typing_extensions import override
+
 from pyinfra.api import FactBase
 
 from .util.packaging import parse_packages
 
-APK_REGEX = r"^([a-zA-Z0-9\-_]+)-([0-9\.]+\-?[a-z0-9]*)\s"
+# Source: https://superuser.com/a/1472405
+# Modified to return version and release inside a single group and removed extra capturing groups
+APK_REGEX = r"(.+)-([^-]+-r[^-]+) \S+ \{\S+\} \(.+?\)"
 
 
 class ApkPackages(FactBase):
@@ -18,13 +22,16 @@ class ApkPackages(FactBase):
         }
     """
 
+    @override
     def command(self) -> str:
         return "apk list --installed"
 
+    @override
     def requires_command(self) -> str:
         return "apk"
 
     default = dict
 
+    @override
     def process(self, output):
         return parse_packages(APK_REGEX, output)

pyinfra/facts/apt.py

@@ -2,12 +2,36 @@ from __future__ import annotations
 
 import re
 
+from typing_extensions import TypedDict, override
+
 from pyinfra.api import FactBase
 
 from .gpg import GpgFactBase
 from .util import make_cat_files_command
 
 
+def noninteractive_apt(command: str, force=False):
+    args = ["DEBIAN_FRONTEND=noninteractive apt-get -y"]
+
+    if force:
+        args.append("--force-yes")
+
+    args.extend(
+        (
+            '-o Dpkg::Options::="--force-confdef"',
+            '-o Dpkg::Options::="--force-confold"',
+            command,
+        ),
+    )
+
+    return " ".join(args)
+
+
+APT_CHANGES_RE = re.compile(
+    r"^(\d+) upgraded, (\d+) newly installed, (\d+) to remove and (\d+) not upgraded.$"
+)
+
+
 def parse_apt_repo(name):
     regex = r"^(deb(?:-src)?)(?:\s+\[([^\]]+)\])?\s+([^\s]+)\s+([^\s]+)\s+([a-z-\s\d]*)$"
 
@@ -52,17 +76,20 @@ class AptSources(FactBase):
         ]
     """
 
+    @override
     def command(self) -> str:
         return make_cat_files_command(
             "/etc/apt/sources.list",
             "/etc/apt/sources.list.d/*.list",
         )
 
+    @override
     def requires_command(self) -> str:
         return "apt"  # if apt installed, above should exist
 
     default = list
 
+    @override
     def process(self, output):
         repos = []
 
@@ -89,8 +116,49 @@ class AptKeys(GpgFactBase):
     """
 
     # This requires both apt-key *and* apt-key itself requires gpg
+    @override
     def command(self) -> str:
         return "! command -v gpg || apt-key list --with-colons"
 
+    @override
     def requires_command(self) -> str:
         return "apt-key"
+
+
+class AptSimulationDict(TypedDict):
+    upgraded: int
+    newly_installed: int
+    removed: int
+    not_upgraded: int
+
+
+class SimulateOperationWillChange(FactBase[AptSimulationDict]):
+    """
+    Simulate an 'apt-get' operation and try to detect if any changes would be performed.
+    """
+
+    @override
+    def command(self, command: str) -> str:
+        # LC_ALL=C: Ensure the output is in english, as we want to parse it
+        return "LC_ALL=C " + noninteractive_apt(f"{command} --dry-run")
+
+    @override
+    def requires_command(self, command: str) -> str:
+        return "apt-get"
+
+    @override
+    def process(self, output) -> AptSimulationDict:
+        # We are looking for a line similar to
+        # "3 upgraded, 0 newly installed, 0 to remove and 0 not upgraded."
+        for line in output:
+            result = APT_CHANGES_RE.match(line)
+            if result is not None:
+                return {
+                    "upgraded": int(result[1]),
+                    "newly_installed": int(result[2]),
+                    "removed": int(result[3]),
+                    "not_upgraded": int(result[4]),
+                }
+
+        # We did not find the line we expected:
+        raise Exception("Did not find proposed changes in output")

pyinfra/facts/brew.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import re
 
+from typing_extensions import override
+
 from pyinfra import logger
 from pyinfra.api import FactBase
 
@@ -37,9 +39,11 @@ class BrewVersion(FactBase):
 
     """
 
+    @override
     def command(self) -> str:
         return "brew --version"
 
+    @override
     def requires_command(self) -> str:
         return "brew"
 
@@ -47,6 +51,7 @@ class BrewVersion(FactBase):
     def default():
         return [0, 0, 0]
 
+    @override
     def process(self, output):
         out = list(output)[0]
         m = VERSION_MATCHER.match(out)
@@ -67,14 +72,17 @@ class BrewPackages(FactBase):
         }
     """
 
+    @override
     def command(self) -> str:
         return "brew list --versions"
 
+    @override
     def requires_command(self) -> str:
         return "brew"
 
     default = dict
 
+    @override
     def process(self, output):
         return parse_packages(BREW_REGEX, output)
 
@@ -90,12 +98,14 @@ class BrewCasks(BrewPackages):
         }
     """
 
+    @override
     def command(self) -> str:
         return (
             r'if brew --version | grep -q -e "Homebrew\ +(1\.|2\.[0-5]).*" 1>/dev/null;'
             r"then brew cask list --versions; else brew list --cask --versions; fi"
         )
 
+    @override
     def requires_command(self) -> str:
         return "brew"
 
@@ -105,13 +115,16 @@ class BrewTaps(FactBase):
     Returns a list of brew taps.
     """
 
+    @override
     def command(self) -> str:
         return "brew tap"
 
+    @override
     def requires_command(self) -> str:
         return "brew"
 
     default = list
 
+    @override
     def process(self, output):
         return output

pyinfra/facts/bsdinit.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing_extensions import override
+
 from .sysvinit import InitdStatus
 
 
@@ -9,6 +11,7 @@ class RcdStatus(InitdStatus):
     BSD init scripts are well behaved and as such their output can be trusted.
     """
 
+    @override
     def command(self) -> str:
         return """
             for SERVICE in `find /etc/rc.d /usr/local/etc/rc.d -type f`; do

pyinfra/facts/cargo.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing_extensions import override
+
 from pyinfra.api import FactBase
 
 from .util.packaging import parse_packages
@@ -22,11 +24,14 @@ class CargoPackages(FactBase):
 
     default = dict
 
+    @override
     def command(self) -> str:
         return "cargo install --list"
 
+    @override
     def requires_command(self) -> str:
         return "cargo"
 
+    @override
     def process(self, output):
         return parse_packages(CARGO_REGEX, output)

pyinfra/facts/choco.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing_extensions import override
+
 from pyinfra.api import FactBase
 
 from .util.packaging import parse_packages
@@ -18,6 +20,7 @@ class ChocoPackages(FactBase):
         }
     """
 
+    @override
     def command(self) -> str:
         return "choco list"
 
@@ -25,6 +28,7 @@ class ChocoPackages(FactBase):
 
     default = dict
 
+    @override
     def process(self, output):
         return parse_packages(CHOCO_REGEX, output)
 
@@ -34,8 +38,10 @@ class ChocoVersion(FactBase):
     Returns the choco (Chocolatey) version.
     """
 
+    @override
     def command(self) -> str:
         return "choco --version"
 
+    @override
     def process(self, output):
         return "".join(output).replace("\n", "")