pyinfra 2.9.2__py2.py3-none-any.whl → 3.0__py2.py3-none-any.whl

pyinfra/connectors/ssh_util.py

@@ -0,0 +1,114 @@
+from getpass import getpass
+from os import path
+from typing import TYPE_CHECKING, Type, Union
+
+from paramiko import (
+    DSSKey,
+    ECDSAKey,
+    Ed25519Key,
+    PasswordRequiredException,
+    PKey,
+    RSAKey,
+    SSHException,
+)
+
+import pyinfra
+from pyinfra.api.exceptions import ConnectError, PyinfraError
+
+if TYPE_CHECKING:
+    from pyinfra.api.host import Host
+    from pyinfra.api.state import State
+
+
+def raise_connect_error(host: "Host", message, data):
+    message = "{0} ({1})".format(message, data)
+    raise ConnectError(message)
+
+
+def _load_private_key_file(filename: str, key_filename: str, key_password: str):
+    exception: Union[PyinfraError, SSHException] = PyinfraError("Invalid key: {0}".format(filename))
+
+    key_cls: Union[Type[RSAKey], Type[DSSKey], Type[ECDSAKey], Type[Ed25519Key]]
+
+    for key_cls in (RSAKey, DSSKey, ECDSAKey, Ed25519Key):
+        try:
+            return key_cls.from_private_key_file(
+                filename=filename,
+            )
+
+        except PasswordRequiredException:
+            if not key_password:
+                # If password is not provided, but we're in CLI mode, ask for it. I'm not a
+                # huge fan of having CLI specific code in here, but it doesn't really fit
+                # anywhere else without duplicating lots of key related code into cli.py.
+                if pyinfra.is_cli:
+                    key_password = getpass(
+                        "Enter password for private key: {0}: ".format(
+                            key_filename,
+                        ),
+                    )
+
+                # API mode and no password? We can't continue!
+                else:
+                    raise PyinfraError(
+                        "Private key file ({0}) is encrypted, set ssh_key_password to "
+                        "use this key".format(key_filename),
+                    )
+
+            try:
+                return key_cls.from_private_key_file(
+                    filename=filename,
+                    password=key_password,
+                )
+            except SSHException as e:  # key does not match key_cls type
+                exception = e
+        except SSHException as e:  # key does not match key_cls type
+            exception = e
+    raise exception
+
+
+def get_private_key(state: "State", key_filename: str, key_password: str) -> PKey:
+    if key_filename in state.private_keys:
+        return state.private_keys[key_filename]
+
+    ssh_key_filenames = [
+        # Global from executed directory
+        path.expanduser(key_filename),
+    ]
+
+    if state.cwd:
+        # Relative to the CWD
+        path.join(state.cwd, key_filename)
+
+    key = None
+    key_file_exists = False
+
+    for filename in ssh_key_filenames:
+        if not path.isfile(filename):
+            continue
+
+        key_file_exists = True
+
+        try:
+            key = _load_private_key_file(filename, key_filename, key_password)
+            break
+        except SSHException:
+            pass
+
+    # No break, so no key found
+    if not key:
+        if not key_file_exists:
+            raise PyinfraError("No such private key file: {0}".format(key_filename))
+        raise PyinfraError("Invalid private key file: {0}".format(key_filename))
+
+    # Load any certificate, names from OpenSSH:
+    # https://github.com/openssh/openssh-portable/blob/049297de975b92adcc2db77e3fb7046c0e3c695d/ssh-keygen.c#L2453  # noqa: E501
+    for certificate_filename in (
+        "{0}-cert.pub".format(key_filename),
+        "{0}.pub".format(key_filename),
+    ):
+        if path.isfile(certificate_filename):
+            key.load_certificate(certificate_filename)
+
+    state.private_keys[key_filename] = key
+    return key

pyinfra/connectors/sshuserclient/client.py

@@ -166,12 +166,14 @@ class SSHClient(ParamikoClient):
             forward_agent = _pyinfra_ssh_forward_agent
 
         if forward_agent:
-            # Enable SSH forwarding
-            session = self.get_transport().open_session()
+            transport = self.get_transport()
+            assert transport is not None, "No transport"
+            session = transport.open_session()
             AgentRequestHandler(session)
 
     def gateway(self, hostname, host_port, target, target_port):
         transport = self.get_transport()
+        assert transport is not None, "No transport"
         return transport.open_channel(
             "direct-tcpip",
             (target, target_port),
@@ -185,7 +187,7 @@ class SSHClient(ParamikoClient):
         ssh_config_file=None,
         strict_host_key_checking=None,
     ):
-        cfg = {"port": 22}
+        cfg: dict = {"port": 22}
         cfg.update(initial_cfg or {})
 
         forward_agent = False

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
@@ -39,6 +5,8 @@ from pyinfra.api.exceptions import InventoryError
 from pyinfra.api.util import memoize
 from pyinfra.progress import progress_spinner
 
+from .base import BaseConnector
+
 
 @memoize
 def show_warning():
@@ -58,44 +26,97 @@ def _flatten_dict(d: dict, parent_key: str = "", sep: str = "."):
     return dict(_flatten_dict_gen(d, parent_key, sep))
 
 
-def make_names_data(output_key=None):
-    show_warning()
+class TerraformInventoryConnector(BaseConnector):
+    """
+    Generate one or more SSH hosts from a Terraform output variable. The variable
+    must be a list of hostnames or dictionaries.
 
-    if not output_key:
-        raise InventoryError("No Terraform output key!")
+    Output is fetched from a flattened JSON dictionary output from ``terraform output
+    -json``. For example the following object:
 
-    with progress_spinner({"fetch terraform output"}):
-        tf_output_raw = local.shell("terraform output -json")
+    .. code:: json
 
-    tf_output = json.loads(tf_output_raw)
-    tf_output = _flatten_dict(tf_output)
+        {
+          "server_group": {
+            "value": {
+              "server_group_node_ips": [
+                "1.2.3.4",
+                "1.2.3.5",
+                "1.2.3.6"
+              ]
+            }
+          }
+        }
 
-    tf_output_value = tf_output.get(output_key)
-    if tf_output_value is None:
-        keys = "\n".join(f"   - {k}" for k in tf_output.keys())
-        raise InventoryError(f"No Terraform output with key: `{output_key}`, valid keys:\n{keys}")
+    The IP list ``server_group_node_ips`` would be used like so:
 
-    if not isinstance(tf_output_value, list):
-        raise InventoryError(
-            "Invalid Terraform output type, should be `list`, got "
-            f"`{type(tf_output_value).__name__}`",
-        )
+    .. code:: sh
 
-    for ssh_target in tf_output_value:
-        if isinstance(ssh_target, dict):
-            name = ssh_target.pop("name", ssh_target.get("ssh_hostname"))
-            if name is None:
-                raise InventoryError(
-                    "Invalid Terraform list item, missing `name` or `ssh_hostname` keys",
-                )
-            yield f"@terraform/{name}", ssh_target, ["@terraform"]
+        pyinfra @terraform/server_group.value.server_group_node_ips ...
 
-        elif isinstance(ssh_target, str):
-            data = {"ssh_hostname": ssh_target}
-            yield f"@terraform/{ssh_target}", data, ["@terraform"]
+    You can also specify dictionaries to include extra data with hosts:
 
-        else:
+    .. 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(name=None):
+        show_warning()
+
+        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(name)
+        if tf_output_value is None:
+            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(
-                "Invalid Terraform list item, should be `dict` or `str` got "
-                f"`{type(ssh_target).__name__}`",
+                "Invalid Terraform output type, should be `list`, got "
+                f"`{type(tf_output_value).__name__}`",
             )
+
+        for ssh_target in tf_output_value:
+            if isinstance(ssh_target, dict):
+                name = ssh_target.pop("name", ssh_target.get("ssh_hostname"))
+                if name is None:
+                    raise InventoryError(
+                        "Invalid Terraform list item, missing `name` or `ssh_hostname` keys",
+                    )
+                yield f"@terraform/{name}", ssh_target, ["@terraform"]
+
+            elif isinstance(ssh_target, str):
+                data = {"ssh_hostname": ssh_target}
+                yield f"@terraform/{ssh_target}", data, ["@terraform"]
+
+            else:
+                raise InventoryError(
+                    "Invalid Terraform list item, should be `dict` or `str` got "
+                    f"`{type(ssh_target).__name__}`",
+                )