pyinfra 3.0.dev0__py2.py3-none-any.whl → 3.0.2__py2.py3-none-any.whl

pyinfra/connectors/ssh.py

@@ -1,31 +1,22 @@
-"""
-Connect to hosts over SSH. This is the default connector and all targets default
-to this meaning you do not need to specify it - ie the following two commands
-are identical:
-
-.. code:: shell
-
-    pyinfra my-host.net ...
-    pyinfra @ssh/my-host.net ...
-"""
-
 from __future__ import annotations
 
 import shlex
-from distutils.spawn import find_executable
+from random import uniform
+from shutil import which
 from socket import error as socket_error, gaierror
+from time import sleep
 from typing import TYPE_CHECKING, Any, Iterable, Optional, Tuple
 
 import click
 from paramiko import AuthenticationException, BadHostKeyException, SFTPClient, SSHException
-from typing_extensions import Unpack
+from typing_extensions import TypedDict, Unpack
 
 from pyinfra import logger
 from pyinfra.api.command import QuoteString, StringCommand
 from pyinfra.api.exceptions import ConnectError
 from pyinfra.api.util import get_file_io, memoize
 
-from .base import BaseConnector, make_keys
+from .base import BaseConnector, DataMeta
 from .ssh_util import get_private_key, raise_connect_error
 from .sshuserclient import SSHClient
 from .util import (
@@ -41,94 +32,197 @@ if TYPE_CHECKING:
     from pyinfra.api.arguments import ConnectorArguments
 
 
-class DataKeys:
-    hostname = "SSH hostname"
-    port = "SSH port"
+class ConnectorData(TypedDict):
+    ssh_hostname: str
+    ssh_port: int
+    ssh_user: str
+    ssh_password: str
+    ssh_key: str
+    ssh_key_password: str
+
+    ssh_allow_agent: bool
+    ssh_look_for_keys: bool
+    ssh_forward_agent: bool
+
+    ssh_config_file: str
+    ssh_known_hosts_file: str
+    ssh_strict_host_key_checking: str
+
+    ssh_paramiko_connect_kwargs: dict
+
+    ssh_connect_retries: int
+    ssh_connect_retry_min_delay: float
+    ssh_connect_retry_max_delay: float
+
+
+connector_data_meta: dict[str, DataMeta] = {
+    "ssh_hostname": DataMeta("SSH hostname"),
+    "ssh_port": DataMeta("SSH port"),
+    "ssh_user": DataMeta("SSH user"),
+    "ssh_password": DataMeta("SSH password"),
+    "ssh_key": DataMeta("SSH key filename"),
+    "ssh_key_password": DataMeta("SSH key password"),
+    "ssh_allow_agent": DataMeta(
+        "Whether to use any active SSH agent",
+        True,
+    ),
+    "ssh_look_for_keys": DataMeta(
+        "Whether to look for private keys",
+        True,
+    ),
+    "ssh_forward_agent": DataMeta(
+        "Whether to enable SSH forward agent",
+        False,
+    ),
+    "ssh_config_file": DataMeta("SSH config filename"),
+    "ssh_known_hosts_file": DataMeta("SSH known_hosts filename"),
+    "ssh_strict_host_key_checking": DataMeta(
+        "SSH strict host key checking",
+        "accept-new",
+    ),
+    "ssh_paramiko_connect_kwargs": DataMeta(
+        "Override keyword arguments passed into Paramiko's ``SSHClient.connect``"
+    ),
+    "ssh_connect_retries": DataMeta("Number of tries to connect via ssh", 0),
+    "ssh_connect_retry_min_delay": DataMeta(
+        "Lower bound for random delay between retries",
+        0.1,
+    ),
+    "ssh_connect_retry_max_delay": DataMeta(
+        "Upper bound for random delay between retries",
+        0.5,
+    ),
+}
 
-    user = "User to SSH as"
-    password = "Password to use for authentication"
-    key = "Key file to use for authentication"
-    key_password = "Key file password"
 
-    allow_agent = "Allow using SSH agent"
-    look_for_keys = "Allow looking up users keys"
+class SSHConnector(BaseConnector):
+    """
+    Connect to hosts over SSH. This is the default connector and all targets default
+    to this meaning you do not need to specify it - ie the following two commands
+    are identical:
 
-    forward_agent = "Enable SSH forward agent"
-    config_file = "Custom SSH config file"
-    known_hosts_file = "Custom SSH known hosts file"
-    strict_host_key_checking = "Override strict host keys check setting"
+    .. code:: shell
 
-    paramiko_connect_kwargs = (
-        "Override keyword arguments passed into paramiko's `SSHClient.connect`"
-    )
+        pyinfra my-host.net ...
+        pyinfra @ssh/my-host.net ...
+    """
 
+    __examples_doc__ = """
+    An inventory file (``inventory.py``) containing a single SSH target with SSH
+    forward agent enabled:
 
-DATA_KEYS = make_keys("ssh", DataKeys)
+    .. code:: python
 
+        hosts = [
+            ("my-host.net", {"ssh_forward_agent": True}),
+        ]
+
+    Multiple hosts sharing the same SSH username:
+
+    .. code:: python
+
+        hosts = (
+            ["my-host-1.net", "my-host-2.net"],
+            {"ssh_user": "ssh-user"},
+        )
+
+    Multiple hosts with different SSH usernames:
+
+    .. code:: python
+
+        hosts = [
+            ("my-host-1.net", {"ssh_user": "ssh-user"}),
+            ("my-host-2.net", {"ssh_user": "other-user"}),
+        ]
+    """
 
-class SSHConnector(BaseConnector):
     handles_execution = True
+
+    data_cls = ConnectorData
+    data_meta = connector_data_meta
+    data: ConnectorData
+
     client: Optional[SSHClient] = None
 
-    def make_names_data(hostname):
-        yield "@ssh/{0}".format(hostname), {DATA_KEYS.hostname: hostname}, []
+    @staticmethod
+    def make_names_data(name):
+        yield "@ssh/{0}".format(name), {"ssh_hostname": name}, []
 
     def make_paramiko_kwargs(self) -> dict[str, Any]:
         kwargs = {
             "allow_agent": False,
             "look_for_keys": False,
-            "hostname": self.host.data.get(DATA_KEYS.hostname, self.host.name),
+            "hostname": self.data["ssh_hostname"] or self.host.name,
             # Overrides of SSH config via pyinfra host data
-            "_pyinfra_ssh_forward_agent": self.host.data.get(DATA_KEYS.forward_agent),
-            "_pyinfra_ssh_config_file": self.host.data.get(DATA_KEYS.config_file),
-            "_pyinfra_ssh_known_hosts_file": self.host.data.get(DATA_KEYS.known_hosts_file),
-            "_pyinfra_ssh_strict_host_key_checking": self.host.data.get(
-                DATA_KEYS.strict_host_key_checking
-            ),
-            "_pyinfra_ssh_paramiko_connect_kwargs": self.host.data.get(
-                DATA_KEYS.paramiko_connect_kwargs
-            ),
+            "_pyinfra_ssh_forward_agent": self.data["ssh_forward_agent"],
+            "_pyinfra_ssh_config_file": self.data["ssh_config_file"],
+            "_pyinfra_ssh_known_hosts_file": self.data["ssh_known_hosts_file"],
+            "_pyinfra_ssh_strict_host_key_checking": self.data["ssh_strict_host_key_checking"],
+            "_pyinfra_ssh_paramiko_connect_kwargs": self.data["ssh_paramiko_connect_kwargs"],
         }
 
         for key, value in (
-            ("username", self.host.data.get(DATA_KEYS.user)),
-            ("port", int(self.host.data.get(DATA_KEYS.port, 0))),
+            ("username", self.data["ssh_user"]),
+            ("port", int(self.data["ssh_port"] or 0)),
             ("timeout", self.state.config.CONNECT_TIMEOUT),
         ):
             if value:
                 kwargs[key] = value
 
         # Password auth (boo!)
-        ssh_password = self.host.data.get(DATA_KEYS.password)
+        ssh_password = self.data["ssh_password"]
         if ssh_password:
             kwargs["password"] = ssh_password
 
         # Key auth!
-        ssh_key = self.host.data.get(DATA_KEYS.key)
+        ssh_key = self.data["ssh_key"]
         if ssh_key:
             kwargs["pkey"] = get_private_key(
                 self.state,
                 key_filename=ssh_key,
-                key_password=self.host.data.get(DATA_KEYS.key_password),
+                key_password=self.data["ssh_key_password"],
             )
 
         # No key or password, so let's have paramiko look for SSH agents and user keys
         # unless disabled by the user.
         else:
-            kwargs["allow_agent"] = self.host.data.get(DATA_KEYS.allow_agent, True)
-            kwargs["look_for_keys"] = self.host.data.get(DATA_KEYS.look_for_keys, True)
+            kwargs["allow_agent"] = self.data["ssh_allow_agent"]
+            kwargs["look_for_keys"] = self.data["ssh_look_for_keys"]
 
         return kwargs
 
     def connect(self) -> None:
+        retries = self.data["ssh_connect_retries"]
+
+        try:
+            while True:
+                try:
+                    return self._connect()
+                except (SSHException, gaierror, socket_error, EOFError):
+                    if retries == 0:
+                        raise
+                    retries -= 1
+                    min_delay = self.data["ssh_connect_retry_min_delay"]
+                    max_delay = self.data["ssh_connect_retry_max_delay"]
+                    sleep(uniform(min_delay, max_delay))
+        except SSHException as e:
+            raise_connect_error(self.host, "SSH error", e)
+        except gaierror as e:
+            raise_connect_error(self.host, "Could not resolve hostname", e)
+        except socket_error as e:
+            raise_connect_error(self.host, "Could not connect", e)
+        except EOFError as e:
+            raise_connect_error(self.host, "EOF error", e)
+
+    def _connect(self) -> None:
         """
         Connect to a single host. Returns the SSH client if successful. Stateless by
         design so can be run in parallel.
         """
 
         kwargs = self.make_paramiko_kwargs()
-        logger.debug("Connecting to: %s (%r)", self.host.name, kwargs)
         hostname = kwargs.pop("hostname")
+        logger.debug("Connecting to: %s (%r)", hostname, kwargs)
 
         self.client = SSHClient()
 
@@ -143,7 +237,7 @@ class SSHConnector(BaseConnector):
                     continue
 
                 if key == "pkey" and value:
-                    auth_kwargs["key"] = self.host.data.get(DATA_KEYS.key)
+                    auth_kwargs["key"] = self.data["ssh_key"]
 
             auth_args = ", ".join(
                 "{0}={1}".format(key, value) for key, value in auth_kwargs.items()
@@ -170,18 +264,6 @@ class SSHConnector(BaseConnector):
                 f"Host key for {e.hostname} does not match.",
             )
 
-        except SSHException as e:
-            raise_connect_error(self.host, "SSH error", e)
-
-        except gaierror:
-            raise_connect_error(self.host, "Could not resolve hostname", hostname)
-
-        except socket_error as e:
-            raise_connect_error(self.host, "Could not connect", e)
-
-        except EOFError as e:
-            raise_connect_error(self.host, "EOF error", e)
-
     def run_shell_command(
         self,
         command: StringCommand,
@@ -303,7 +385,7 @@ class SSHConnector(BaseConnector):
 
         if _sudo or _su_user:
             # Get temp file location
-            temp_file = remote_temp_filename or self.state.get_temp_filename(remote_filename)
+            temp_file = remote_temp_filename or self.host.get_temp_filename(remote_filename)
 
             # Copy the file to the tempfile location and add read permissions
             command = StringCommand(
@@ -395,20 +477,14 @@ class SSHConnector(BaseConnector):
         # user connected, so upload to tmp and copy/chown w/sudo and/or su_user
         if _sudo or _doas or _su_user:
             # Get temp file location
-            temp_file = remote_temp_filename or self.state.get_temp_filename(remote_filename)
+            temp_file = remote_temp_filename or self.host.get_temp_filename(remote_filename)
             self._put_file(filename_or_io, temp_file)
 
             # Make sure our sudo/su user can access the file
-            if _su_user:
-                command = StringCommand("setfacl", "-m", "u:{0}:r".format(_su_user), temp_file)
-            elif _sudo_user:
-                command = StringCommand("setfacl", "-m", "u:{0}:r".format(_sudo_user), temp_file)
-            elif _doas_user:
-                command = StringCommand("setfacl", "-m", "u:{0}:r".format(_doas_user), temp_file)
-
-            if _su_user or _sudo_user or _doas_user:
+            other_user = _su_user or _sudo_user or _doas_user
+            if other_user:
                 status, output = self.run_shell_command(
-                    command,
+                    StringCommand("setfacl", "-m", f"u:{other_user}:r", temp_file),
                     print_output=print_output,
                     print_input=print_input,
                     **arguments,
@@ -459,15 +535,15 @@ class SSHConnector(BaseConnector):
         return True
 
     def check_can_rsync(self):
-        if self.host.data.get(DATA_KEYS.key_password):
+        if self.data["ssh_key_password"]:
             raise NotImplementedError(
                 "Rsync does not currently work with SSH keys needing passwords."
             )
 
-        if self.host.data.get(DATA_KEYS.password):
+        if self.data["ssh_password"]:
             raise NotImplementedError("Rsync does not currently work with SSH passwords.")
 
-        if not find_executable("rsync"):
+        if not which("rsync"):
             raise NotImplementedError("The `rsync` binary is not available on this system.")
 
     def rsync(
@@ -482,8 +558,8 @@ class SSHConnector(BaseConnector):
         _sudo = arguments.pop("_sudo", False)
         _sudo_user = arguments.pop("_sudo_user", False)
 
-        hostname = self.host.data.get(DATA_KEYS.hostname, self.host.name)
-        user = self.host.data.get(DATA_KEYS.user, "")
+        hostname = self.data["ssh_hostname"] or self.host.name
+        user = self.data["ssh_user"]
         if user:
             user = "{0}@".format(user)
 
@@ -491,27 +567,27 @@ class SSHConnector(BaseConnector):
         # To avoid asking for interactive input, specify BatchMode=yes
         ssh_flags.append("-o BatchMode=yes")
 
-        known_hosts_file = self.host.data.get(DATA_KEYS.known_hosts_file, "")
+        known_hosts_file = self.data["ssh_known_hosts_file"]
         if known_hosts_file:
             ssh_flags.append(
                 '-o \\"UserKnownHostsFile={0}\\"'.format(shlex.quote(known_hosts_file))
             )  # never trust users
 
-        strict_host_key_checking = self.host.data.get(DATA_KEYS.strict_host_key_checking, "")
+        strict_host_key_checking = self.data["ssh_strict_host_key_checking"]
         if strict_host_key_checking:
             ssh_flags.append(
                 '-o \\"StrictHostKeyChecking={0}\\"'.format(shlex.quote(strict_host_key_checking))
             )
 
-        ssh_config_file = self.host.data.get(DATA_KEYS.config_file, "")
+        ssh_config_file = self.data["ssh_config_file"]
         if ssh_config_file:
             ssh_flags.append("-F {0}".format(shlex.quote(ssh_config_file)))
 
-        port = self.host.data.get(DATA_KEYS.port)
+        port = self.data["ssh_port"]
         if port:
             ssh_flags.append("-p {0}".format(port))
 
-        ssh_key = self.host.data.get(DATA_KEYS.key)
+        ssh_key = self.data["ssh_key"]
         if ssh_key:
             ssh_flags.append("-i {0}".format(ssh_key))
 

pyinfra/connectors/sshuserclient/client.py

@@ -124,7 +124,7 @@ class SSHClient(ParamikoClient):
     original idea at http://bitprophet.org/blog/2012/11/05/gateway-solutions/.
     """
 
-    def connect(
+    def connect(  # type: ignore[override]
         self,
         hostname,
         _pyinfra_ssh_forward_agent=None,

pyinfra/connectors/terraform.py

@@ -1,37 +1,3 @@
-"""
-.. warning::
-    This connector is in alpha and may change in future releases.
-
-Generate one or more SSH hosts from a Terraform output variable. The variable
-must be a list of hostnames or IP addresses that ``pyinfra`` can connect to
-over SSH. Currently there is no support for specifying SSH user/pass/port/key
-from Terraform, these must be provided via ``pyinfra`` group data or ``--data``
-CLI flags.
-
-Output is fetched from a flattened JSON dictionary output from ``terraform output
--json``. For example the following object:
-
-.. code:: json
-
-    {
-      "server_group": {
-        "value": {
-          "server_group_node_ips": [
-            "1.2.3.4",
-            "1.2.3.5",
-            "1.2.3.6"
-          ]
-        }
-      }
-    }
-
-The IP list ``server_group_node_ips`` would be used like so:
-
-.. code:: python
-
-    pyinfra @terraform/server_group.value.server_group_node_ips ...
-"""
-
 import json
 
 from pyinfra import local, logger
@@ -61,22 +27,74 @@ def _flatten_dict(d: dict, parent_key: str = "", sep: str = "."):
 
 
 class TerraformInventoryConnector(BaseConnector):
+    """
+    Generate one or more SSH hosts from a Terraform output variable. The variable
+    must be a list of hostnames or dictionaries.
+
+    Output is fetched from a flattened JSON dictionary output from ``terraform output
+    -json``. For example the following object:
+
+    .. code:: json
+
+        {
+          "server_group": {
+            "value": {
+              "server_group_node_ips": [
+                "1.2.3.4",
+                "1.2.3.5",
+                "1.2.3.6"
+              ]
+            }
+          }
+        }
+
+    The IP list ``server_group_node_ips`` would be used like so:
+
+    .. code:: sh
+
+        pyinfra @terraform/server_group.value.server_group_node_ips ...
+
+    You can also specify dictionaries to include extra data with hosts:
+
+    .. code:: json
+
+        {
+          "server_group": {
+            "value": {
+              "server_group_node_ips": [
+                {
+                    "ssh_hostname": "1.2.3.4",
+                    "ssh_user": "ssh-user"
+                },
+                {
+                    "ssh_hostname": "1.2.3.5",
+                    "ssh_user": "ssh-user"
+                }
+              ]
+            }
+          }
+        }
+
+    """
+
     @staticmethod
-    def make_names_data(output_key=None):
+    def make_names_data(name=None):
         show_warning()
 
-        if not output_key:
-            raise InventoryError("No Terraform output key!")
+        if not name:
+            name = ""
 
         with progress_spinner({"fetch terraform output"}):
             tf_output_raw = local.shell("terraform output -json")
 
+        assert isinstance(tf_output_raw, str)
         tf_output = json.loads(tf_output_raw)
         tf_output = _flatten_dict(tf_output)
 
-        tf_output_value = tf_output.get(output_key)
+        tf_output_value = tf_output.get(name)
         if tf_output_value is None:
-            raise InventoryError(f"No Terraform output with key: `{output_key}`")
+            keys = "\n".join(f"   - {k}" for k in tf_output.keys())
+            raise InventoryError(f"No Terraform output with key: `{name}`, valid keys:\n{keys}")
 
         if not isinstance(tf_output_value, list):
             raise InventoryError(

pyinfra/connectors/util.py

@@ -40,7 +40,7 @@ def run_local_process(
     stdin=None,
     timeout: Optional[int] = None,
     print_output: bool = False,
-    print_prefix=None,
+    print_prefix: str = "",
 ) -> tuple[int, "CommandOutput"]:
     process = Popen(command, shell=True, stdout=PIPE, stderr=PIPE, stdin=PIPE)
 
@@ -201,7 +201,10 @@ def execute_command_with_sudo_retry(
     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":
-            command_arguments["_use_sudo_password"] = True  # ask for the password
+            # 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()
 
     return return_code, output
@@ -221,25 +224,6 @@ def write_stdin(stdin, buffer):
     buffer.close()
 
 
-def _get_sudo_password(host: "Host", use_sudo_password) -> str:
-    if not host.connector_data.get("sudo_askpass_path"):
-        _, output = host.run_shell_command(SUDO_ASKPASS_COMMAND)
-        host.connector_data["sudo_askpass_path"] = shlex.quote(output.stdout_lines[0])
-
-    if use_sudo_password is True:
-        sudo_password = host.connector_data.get("sudo_password")
-        if not sudo_password:
-            sudo_password = getpass("{0}sudo password: ".format(host.print_prefix))
-            host.connector_data["sudo_password"] = sudo_password
-        sudo_password = sudo_password
-    elif callable(use_sudo_password):
-        sudo_password = use_sudo_password()
-    else:
-        sudo_password = use_sudo_password
-
-    return shlex.quote(sudo_password)
-
-
 def remove_any_sudo_askpass_file(host) -> None:
     sudo_askpass_path = host.connector_data.get("sudo_askpass_path")
     if sudo_askpass_path:
@@ -273,17 +257,32 @@ def extract_control_arguments(arguments: "ConnectorArguments") -> "ConnectorArgu
     return control_arguments
 
 
+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)
+    host.connector_data["sudo_askpass_path"] = shlex.quote(output.stdout_lines[0])
+
+
 def make_unix_command_for_host(
     state: "State",
     host: "Host",
     command: StringCommand,
     **command_arguments,
 ) -> StringCommand:
-    use_sudo_password = command_arguments.pop("_use_sudo_password", None)
-    if use_sudo_password:
-        command_arguments["_sudo_password"] = _get_sudo_password(host, use_sudo_password)
-        command_arguments["_sudo_askpass_path"] = host.connector_data.get("sudo_askpass_path")
-
+    if not command_arguments.get("_sudo"):
+        # If no sudo, we've nothing to do here
+        return make_unix_command(command, **command_arguments)
+
+    # If the sudo password is not set in the direct arguments,
+    # set it from the connector data value.
+    if "_sudo_password" not in command_arguments or not command_arguments["_sudo_password"]:
+        command_arguments["_sudo_password"] = host.connector_data.get("prompted_sudo_password")
+
+    if command_arguments["_sudo_password"]:
+        # Ensure the askpass path is correctly set and passed through
+        _ensure_sudo_askpass_set_for_host(host)
+        command_arguments["_sudo_askpass_path"] = host.connector_data["sudo_askpass_path"]
     return make_unix_command(command, **command_arguments)
 
 
@@ -339,7 +338,7 @@ def make_unix_command(
             [
                 "env",
                 "SUDO_ASKPASS={0}".format(_sudo_askpass_path),
-                MaskString("{0}={1}".format(SUDO_ASKPASS_ENV_VAR, _sudo_password)),
+                MaskString("{0}={1}".format(SUDO_ASKPASS_ENV_VAR, shlex.quote(_sudo_password))),
             ],
         )