pyinfra 0.11.dev3__py3-none-any.whl → 3.5.1__py3-none-any.whl

pyinfra/connectors/base.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import abc
+from io import IOBase
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    Iterator,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_type_hints,
+)
+
+from typing_extensions import TypedDict, Unpack
+
+from pyinfra.api.exceptions import ConnectorDataTypeError
+from pyinfra.api.util import raise_if_bad_type
+
+if TYPE_CHECKING:
+    from pyinfra.api.arguments import ConnectorArguments
+    from pyinfra.api.command import StringCommand
+    from pyinfra.api.host import Host, HostData
+    from pyinfra.api.state import State
+
+    from .util import CommandOutput
+
+
+T = TypeVar("T")
+default_sentinel = object()
+
+
+def host_to_connector_data(
+    connector_data: Type[T],
+    connector_data_meta: dict[str, DataMeta],
+    host_data: "HostData",
+) -> T:
+    data: T = cast(T, {})
+    for key, type_ in get_type_hints(connector_data).items():
+        value = host_data.get(key, default_sentinel)
+        if value is default_sentinel:
+            value = connector_data_meta[key].default
+        else:
+            raise_if_bad_type(
+                value,
+                type_,
+                ConnectorDataTypeError,
+                f"Invalid connector data `{key}`:",
+            )
+
+        data[key] = value  # type: ignore
+    return data
+
+
+class DataMeta:
+    description: str
+    default: Any
+
+    def __init__(self, description, default=None) -> None:
+        self.description = description
+        self.default = default
+
+
+class ConnectorData(TypedDict, total=False):
+    pass
+
+
+class BaseConnector(abc.ABC):
+    state: "State"
+    host: "Host"
+
+    handles_execution = False
+
+    data_cls: Type = ConnectorData
+    data_meta: dict[str, DataMeta] = {}
+
+    def __init__(self, state: "State", host: "Host"):
+        self.state = state
+        self.host = host
+        self.data = host_to_connector_data(self.data_cls, self.data_meta, host.data)
+
+    @staticmethod
+    @abc.abstractmethod
+    def make_names_data(name: str) -> Iterator[tuple[str, dict, list[str]]]:
+        """
+        Generate inventory targets. This is a staticmethod because each yield will become a new host
+        object with a new (ie not this) instance of the connector.
+        """
+
+    def connect(self) -> None:
+        """
+        Connect this connector instance. Should raise ConnectError exceptions to indicate failure.
+        """
+
+    def disconnect(self) -> None:
+        """
+        Disconnect this connector instance.
+        """
+
+    @abc.abstractmethod
+    def run_shell_command(
+        self,
+        command: "StringCommand",
+        print_output: bool,
+        print_input: bool,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> tuple[bool, "CommandOutput"]:
+        """
+        Execute a command.
+
+        Args:
+            command (StringCommand): actual command to execute
+            print_output (bool): whether to print command output
+            print_input (bool): whether to print command input
+            arguments: (ConnectorArguments): connector global arguments
+
+        Returns:
+            tuple: (bool, CommandOutput)
+            Bool indicating success and CommandOutput with stdout/stderr lines.
+        """
+
+    @abc.abstractmethod
+    def put_file(
+        self,
+        filename_or_io: Union[str, IOBase],
+        remote_filename: str,
+        remote_temp_filename: Optional[str] = None,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> bool:
+        """
+        Upload a local file or IO object by copying it to a temporary directory
+        and then writing it to the upload location.
+
+        Returns:
+            bool: indicating success or failure.
+        """
+
+    @abc.abstractmethod
+    def get_file(
+        self,
+        remote_filename: str,
+        filename_or_io: Union[str, IOBase],
+        remote_temp_filename: Optional[str] = None,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> bool:
+        """
+        Download a local file by copying it to a temporary location and then writing
+        it to our filename or IO object.
+
+        Returns:
+            bool: indicating success or failure.
+        """
+
+    def check_can_rsync(self) -> None:
+        raise NotImplementedError("This connector does not support rsync")
+
+    def rsync(
+        self,
+        src: str,
+        dest: str,
+        flags: Iterable[str],
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> bool:
+        raise NotImplementedError("This connector does not support rsync")

pyinfra/connectors/chroot.py

@@ -0,0 +1,212 @@
+import os
+from tempfile import mkstemp
+from typing import TYPE_CHECKING, Optional
+
+import click
+from typing_extensions import Unpack, override
+
+from pyinfra import local, logger
+from pyinfra.api import QuoteString, StringCommand
+from pyinfra.api.exceptions import ConnectError, InventoryError, PyinfraError
+from pyinfra.api.util import get_file_io, memoize
+from pyinfra.progress import progress_spinner
+
+from .base import BaseConnector
+from .local import LocalConnector
+from .util import extract_control_arguments, make_unix_command_for_host
+
+if TYPE_CHECKING:
+    from pyinfra.api.arguments import ConnectorArguments
+    from pyinfra.api.host import Host
+    from pyinfra.api.state import State
+
+
+@memoize
+def show_warning() -> None:
+    logger.warning("The @chroot connector is in beta!")
+
+
+class ChrootConnector(BaseConnector):
+    """
+    The chroot connector allows you to execute operations within another root.
+    """
+
+    handles_execution = True
+
+    local: LocalConnector
+
+    def __init__(self, state: "State", host: "Host"):
+        super().__init__(state, host)
+        self.local = LocalConnector(state, host)
+
+    @override
+    @staticmethod
+    def make_names_data(name: Optional[str] = None):
+        if not name:
+            raise InventoryError("No directory provided!")
+
+        show_warning()
+
+        yield (
+            "@chroot/{0}".format(name),
+            {
+                "chroot_directory": "/{0}".format(name.lstrip("/")),
+            },
+            ["@chroot"],
+        )
+
+    @override
+    def connect(self) -> None:
+        self.local.connect()
+
+        chroot_directory = self.host.data.chroot_directory
+
+        try:
+            with progress_spinner({"chroot run"}):
+                local.shell(
+                    "chroot {0} ls".format(chroot_directory),
+                    splitlines=True,
+                )
+        except PyinfraError as e:
+            raise ConnectError(e.args[0])
+
+        self.host.connector_data["chroot_directory"] = chroot_directory
+
+    @override
+    def run_shell_command(
+        self,
+        command,
+        print_output: bool = False,
+        print_input: bool = False,
+        **command_arguments: Unpack["ConnectorArguments"],
+    ):
+        local_arguments = extract_control_arguments(command_arguments)
+
+        chroot_directory = self.host.connector_data["chroot_directory"]
+
+        command = make_unix_command_for_host(self.state, self.host, command, **command_arguments)
+        command = QuoteString(command)
+
+        logger.debug("--> Running chroot command on (%s): %s", chroot_directory, command)
+
+        chroot_command = StringCommand(
+            "chroot",
+            chroot_directory,
+            "sh",
+            "-c",
+            command,
+        )
+
+        return self.local.run_shell_command(
+            chroot_command,
+            print_output=print_output,
+            print_input=print_input,
+            **local_arguments,
+        )
+
+    @override
+    def put_file(
+        self,
+        filename_or_io,
+        remote_filename,
+        remote_temp_filename=None,  # ignored
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        _, temp_filename = mkstemp()
+
+        try:
+            # Load our file or IO object and write it to the temporary file
+            with get_file_io(filename_or_io) as file_io:
+                with open(temp_filename, "wb") as temp_f:
+                    data = file_io.read()
+
+                    if isinstance(data, str):
+                        data = data.encode()
+
+                    temp_f.write(data)
+
+            chroot_directory = self.host.connector_data["chroot_directory"]
+            chroot_command = StringCommand(
+                "cp",
+                temp_filename,
+                f"{chroot_directory}/{remote_filename}",
+            )
+
+            status, output = self.local.run_shell_command(
+                chroot_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+        finally:
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file uploaded to chroot: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status
+
+    @override
+    def get_file(
+        self,
+        remote_filename,
+        filename_or_io,
+        remote_temp_filename=None,  # ignored
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        _, temp_filename = mkstemp()
+
+        try:
+            chroot_directory = self.host.connector_data["chroot_directory"]
+            chroot_command = StringCommand(
+                "cp",
+                f"{chroot_directory}/{remote_filename}",
+                temp_filename,
+            )
+
+            status, output = self.local.run_shell_command(
+                chroot_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+
+            # Load the temporary file and write it to our file or IO object
+            with open(temp_filename, "rb") as temp_f:
+                with get_file_io(filename_or_io, "wb") as file_io:
+                    data = temp_f.read()
+                    data_bytes: bytes
+
+                    if isinstance(data, str):
+                        data_bytes = data.encode()
+                    else:
+                        data_bytes = data
+
+                    file_io.write(data_bytes)
+        finally:
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file downloaded from chroot: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status