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

pyinfra/connectors/docker.py

@@ -0,0 +1,381 @@
+from __future__ import annotations
+
+import json
+import os
+from tempfile import mkstemp
+from typing import TYPE_CHECKING
+
+import click
+from typing_extensions import TypedDict, 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
+from pyinfra.progress import progress_spinner
+
+from .base import BaseConnector, DataMeta
+from .local import LocalConnector
+from .util import CommandOutput, 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
+
+
+class ConnectorData(TypedDict):
+    docker_identifier: str
+
+
+connector_data_meta: dict[str, DataMeta] = {
+    "docker_identifier": DataMeta("ID of container or image to start from"),
+}
+
+
+class DockerConnector(BaseConnector):
+    """
+    The Docker connector allows you to use pyinfra to create new Docker images or modify running
+    Docker containers.
+
+    .. note::
+
+        The Docker connector allows pyinfra to target Docker containers as inventory and is
+        unrelated to the :doc:`../operations/docker` & :doc:`../facts/docker`.
+
+    You can pass either an image name or existing container ID:
+
+    + Image - will create a new container from the image, execute operations against it, save into \
+        a new Docker image and remove the container
+    + Existing container ID - will execute operations against the running container, leaving it \
+        running
+
+    .. code:: shell
+
+        # A Docker base image must be provided
+        pyinfra @docker/alpine:3.8 ...
+
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @docker/alpine:3.8,@docker/ubuntu:bionic ...
+
+        # Execute against a running container
+        pyinfra @docker/2beb8c15a1b1 ...
+
+    The Docker connector is great for testing pyinfra operations locally, rather than connecting to
+    a remote host over SSH each time. This gives you a fast, local-first devloop to iterate on when
+    writing deploys, operations or facts.
+    """
+
+    # enable the use of other docker cli compatible tools like podman
+    docker_cmd = "docker"
+
+    handles_execution = True
+
+    data_cls = ConnectorData
+    data_meta = connector_data_meta
+    data: ConnectorData
+
+    local: LocalConnector
+
+    container_id: str
+    no_stop: bool = False
+
+    def __init__(self, state: "State", host: "Host"):
+        super().__init__(state, host)
+        self.local = LocalConnector(state, host)
+
+    @override
+    @staticmethod
+    def make_names_data(name=None):
+        if not name:
+            raise InventoryError("No docker base ID provided!")
+
+        yield (
+            f"@docker/{name}",
+            {"docker_identifier": name},
+            ["@docker"],
+        )
+
+    # 2 helper functions
+    def _find_start_docker_container(self, container_id) -> tuple[str, bool]:
+        docker_info = local.shell(f"{self.docker_cmd} container inspect {container_id}")
+        assert isinstance(docker_info, str)
+        docker_info = json.loads(docker_info)[0]
+        if docker_info["State"]["Running"] is False:
+            logger.info(f"Starting stopped container: {container_id}")
+            local.shell(f"{self.docker_cmd} container start {container_id}")
+            return container_id, False
+        return container_id, True
+
+    def _start_docker_image(self, image_name):
+        try:
+            return local.shell(
+                f"{self.docker_cmd} run -d {image_name} tail -f /dev/null",
+                splitlines=True,
+            )[-1]  # last line is the container ID
+        except PyinfraError as e:
+            raise ConnectError(e.args[0])
+
+    @override
+    def connect(self) -> None:
+        self.local.connect()
+
+        docker_identifier = self.data["docker_identifier"]
+        with progress_spinner({f"prepare {self.docker_cmd} container"}):
+            try:
+                self.container_id, was_running = self._find_start_docker_container(
+                    docker_identifier
+                )
+                if was_running:
+                    self.no_stop = True
+            except PyinfraError:
+                self.container_id = self._start_docker_image(docker_identifier)
+
+    @override
+    def disconnect(self) -> None:
+        container_id = self.container_id
+
+        if self.no_stop:
+            logger.info(
+                "{0}{1} build complete, container left running: {2}".format(
+                    self.host.print_prefix,
+                    self.docker_cmd,
+                    click.style(container_id, bold=True),
+                ),
+            )
+            return
+
+        with progress_spinner({f"{self.docker_cmd} commit"}):
+            image_id = local.shell(f"{self.docker_cmd} commit {container_id}", splitlines=True)[-1][
+                7:19
+            ]  # last line is the image ID, get sha256:[XXXXXXXXXX]...
+
+        with progress_spinner({f"{self.docker_cmd} rm"}):
+            local.shell(
+                f"{self.docker_cmd} rm -f {container_id}",
+            )
+
+        logger.info(
+            "{0}{1} build complete, image ID: {2}".format(
+                self.host.print_prefix,
+                self.docker_cmd,
+                click.style(image_id, bold=True),
+            ),
+        )
+
+    @override
+    def run_shell_command(
+        self,
+        command: StringCommand,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ) -> tuple[bool, CommandOutput]:
+        local_arguments = extract_control_arguments(arguments)
+
+        container_id = self.container_id
+
+        command = make_unix_command_for_host(self.state, self.host, command, **arguments)
+        command = StringCommand(QuoteString(command))
+
+        docker_flags = "-it" if local_arguments.get("_get_pty") else "-i"
+        docker_command = StringCommand(
+            self.docker_cmd,
+            "exec",
+            docker_flags,
+            container_id,
+            "sh",
+            "-c",
+            command,
+        )
+
+        return self.local.run_shell_command(
+            docker_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=False,
+        print_input=False,
+        **kwargs,  # ignored (sudo/etc)
+    ) -> bool:
+        """
+        Upload a file/IO object to the target container by copying it to a
+        temporary location and then uploading it into the container using ``docker cp``.
+        """
+
+        fd, 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)
+
+            docker_command = StringCommand(
+                self.docker_cmd,
+                "cp",
+                temp_filename,
+                f"{self.container_id}:{remote_filename}",
+            )
+
+            status, output = self.local.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+        finally:
+            os.close(fd)
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file uploaded to container: {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=False,
+        print_input=False,
+        **kwargs,  # ignored (sudo/etc)
+    ) -> bool:
+        """
+        Download a file from the target container by copying it to a temporary
+        location and then reading that into our final file/IO object.
+        """
+
+        fd, temp_filename = mkstemp()
+
+        try:
+            docker_command = StringCommand(
+                self.docker_cmd,
+                "cp",
+                f"{self.container_id}:{remote_filename}",
+                temp_filename,
+            )
+
+            status, output = self.local.run_shell_command(
+                docker_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()
+                    file_io.write(data)
+        finally:
+            os.close(fd)
+            os.remove(temp_filename)
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file downloaded from container: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status
+
+
+class PodmanConnector(DockerConnector):
+    """
+    The Podman connector allows you to use pyinfra to create new Podman images or modify running
+    Podman containers.
+
+    .. note::
+
+        The Podman connector allows pyinfra to target Podman containers as inventory and is
+        unrelated to the :doc:`../operations/docker` & :doc:`../facts/docker`.
+
+    You can pass either an image name or existing container ID:
+
+    + Image - will create a new container from the image, execute operations against it, save into \
+        a new Podman image and remove the container
+    + Existing container ID - will execute operations against the running container, leaving it \
+        running
+
+    .. code:: shell
+
+        # A Podman base image must be provided
+        pyinfra @podman/alpine:3.8 ...
+
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @podman/alpine:3.8,@podman/ubuntu:bionic ...
+
+        # Execute against a running container
+        pyinfra @podman/2beb8c15a1b1 ...
+
+    The Podman connector is great for testing pyinfra operations locally, rather than connecting to
+    a remote host over SSH each time. This gives you a fast, local-first devloop to iterate on when
+    writing deploys, operations or facts.
+    """
+
+    docker_cmd = "podman"
+
+    @override
+    @staticmethod
+    def make_names_data(name=None):
+        if not name:
+            raise InventoryError("No podman base ID provided!")
+
+        yield (
+            f"@podman/{name}",
+            {"docker_identifier": name},
+            ["@podman"],
+        )
+
+    # Duplicate function definition to swap the docstring.
+    @override
+    def put_file(
+        self,
+        filename_or_io,
+        remote_filename,
+        remote_temp_filename=None,  # ignored
+        print_output=False,
+        print_input=False,
+        **kwargs,  # ignored (sudo/etc)
+    ) -> bool:
+        """
+        Upload a file/IO object to the target container by copying it to a
+        temporary location and then uploading it into the container using ``podman cp``.
+        """
+        return super().put_file(
+            filename_or_io,
+            remote_filename,
+            remote_temp_filename,  # ignored
+            print_output,
+            print_input,
+            **kwargs,  # ignored (sudo/etc)
+        )

pyinfra/connectors/dockerssh.py

@@ -0,0 +1,297 @@
+import os
+from tempfile import mkstemp
+from typing import TYPE_CHECKING
+
+import click
+from typing_extensions import Unpack, override
+
+from pyinfra import 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 .ssh import SSHConnector
+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 @dockerssh connector is in beta!")
+
+
+class DockerSSHConnector(BaseConnector):
+    """
+    **Note**: this connector is in beta!
+
+    The ``@dockerssh`` connector allows you to run commands on Docker containers \
+    on a remote machine.
+
+    .. code:: shell
+
+        # A Docker base image must be provided
+        pyinfra @dockerssh/remotehost:alpine:3.8 ...
+
+        # pyinfra can run on multiple Docker images in parallel
+        pyinfra @dockerssh/remotehost:alpine:3.8,@dockerssh/remotehost:ubuntu:bionic ...
+    """
+
+    handles_execution = True
+
+    ssh: SSHConnector
+
+    def __init__(self, state: "State", host: "Host"):
+        super().__init__(state, host)
+        self.ssh = SSHConnector(state, host)
+
+    @override
+    @staticmethod
+    def make_names_data(name):
+        try:
+            hostname, image = name.split(":", 1)
+        except (AttributeError, ValueError):  # failure to parse the name
+            raise InventoryError("No ssh host or docker base image provided!")
+
+        if not image:
+            raise InventoryError("No docker base image provided!")
+
+        show_warning()
+
+        yield (
+            "@dockerssh/{0}:{1}".format(hostname, image),
+            {"ssh_hostname": hostname, "docker_image": image},
+            ["@dockerssh"],
+        )
+
+    @override
+    def connect(self) -> None:
+        self.ssh.connect()
+
+        if "docker_container_id" in self.host.host_data:  # user can provide a docker_container_id
+            return
+
+        try:
+            with progress_spinner({"docker run"}):
+                # last line is the container ID
+                status, output = self.ssh.run_shell_command(
+                    StringCommand(
+                        "docker",
+                        "run",
+                        "-d",
+                        self.host.data.docker_image,
+                        "tail",
+                        "-f",
+                        "/dev/null",
+                    ),
+                )
+                if not status:
+                    raise IOError(output.stderr)
+                container_id = output.stdout_lines[-1]
+
+        except PyinfraError as e:
+            raise ConnectError(e.args[0])
+
+        self.host.host_data["docker_container_id"] = container_id
+
+    @override
+    def disconnect(self) -> None:
+        container_id = self.host.host_data["docker_container_id"][:12]
+
+        with progress_spinner({"docker commit"}):
+            _, output = self.ssh.run_shell_command(StringCommand("docker", "commit", container_id))
+
+            # Last line is the image ID, get sha256:[XXXXXXXXXX]...
+            image_id = output.stdout_lines[-1][7:19]
+
+        with progress_spinner({"docker rm"}):
+            self.ssh.run_shell_command(
+                StringCommand("docker", "rm", "-f", container_id),
+            )
+
+        logger.info(
+            "{0}docker build complete, image ID: {1}".format(
+                self.host.print_prefix,
+                click.style(image_id, bold=True),
+            ),
+        )
+
+    @override
+    def run_shell_command(
+        self,
+        command,
+        print_output: bool = False,
+        print_input: bool = False,
+        **arguments: Unpack["ConnectorArguments"],
+    ):
+        local_arguments = extract_control_arguments(arguments)
+
+        container_id = self.host.host_data["docker_container_id"]
+
+        command = make_unix_command_for_host(self.state, self.host, command, **arguments)
+        command = QuoteString(command)
+
+        docker_flags = "-it" if local_arguments.get("_get_pty") else "-i"
+        docker_command = StringCommand(
+            "docker",
+            "exec",
+            docker_flags,
+            container_id,
+            "sh",
+            "-c",
+            command,
+        )
+
+        return self.ssh.run_shell_command(
+            docker_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,
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        """
+        Upload a file/IO object to the target Docker container by copying it to a
+        temporary location and then uploading it into the container using ``docker cp``.
+        """
+
+        fd, local_temp_filename = mkstemp()
+        remote_temp_filename = remote_temp_filename or self.host.get_temp_filename(
+            local_temp_filename
+        )
+
+        # 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(local_temp_filename, "wb") as temp_f:
+                data = file_io.read()
+
+                if isinstance(data, str):
+                    data = data.encode()
+
+                temp_f.write(data)
+
+        # upload file to remote server
+        ssh_status = self.ssh.put_file(local_temp_filename, remote_temp_filename)
+        if not ssh_status:
+            raise IOError("Failed to copy file over ssh")
+
+        try:
+            docker_id = self.host.host_data["docker_container_id"]
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                remote_temp_filename,
+                f"{docker_id}:{remote_filename}",
+            )
+
+            status, output = self.ssh.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+        finally:
+            os.close(fd)
+            os.remove(local_temp_filename)
+            self.remote_remove(
+                local_temp_filename,
+                print_output=print_output,
+                print_input=print_input,
+            )
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file uploaded to container: {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,
+        print_output: bool = False,
+        print_input: bool = False,
+        **kwargs,  # ignored (sudo/etc)
+    ):
+        """
+        Download a file from the target Docker container by copying it to a temporary
+        location and then reading that into our final file/IO object.
+        """
+
+        remote_temp_filename = remote_temp_filename or self.host.get_temp_filename(remote_filename)
+
+        try:
+            docker_id = self.host.host_data["docker_container_id"]
+            docker_command = StringCommand(
+                "docker",
+                "cp",
+                f"{docker_id}:{remote_filename}",
+                remote_temp_filename,
+            )
+
+            status, output = self.ssh.run_shell_command(
+                docker_command,
+                print_output=print_output,
+                print_input=print_input,
+            )
+
+            ssh_status = self.ssh.get_file(remote_temp_filename, filename_or_io)
+        finally:
+            self.remote_remove(
+                remote_temp_filename,
+                print_output=print_output,
+                print_input=print_input,
+            )
+
+        if not ssh_status:
+            raise IOError("failed to copy file over ssh")
+
+        if not status:
+            raise IOError(output.stderr)
+
+        if print_output:
+            click.echo(
+                "{0}file downloaded from container: {1}".format(
+                    self.host.print_prefix,
+                    remote_filename,
+                ),
+                err=True,
+            )
+
+        return status
+
+    def remote_remove(self, filename, print_output: bool = False, print_input: bool = False):
+        """
+        Deletes a file on a remote machine over ssh.
+        """
+        remove_status, output = self.ssh.run_shell_command(
+            StringCommand("rm", "-f", filename),
+            print_output=print_output,
+            print_input=print_input,
+        )
+
+        if not remove_status:
+            raise IOError(output.stderr)