tower-of-babel-public/tasks/docker.py

import logging
import subprocess
import json
import tempfile
from pathlib import Path
from typing import Optional

from pyinfra import host
from pyinfra.api import deploy, FactBase, operation, FunctionCommand
from pyinfra.operations import files, apt, dnf, systemd, python, server
from pyinfra.facts.server import LinuxName
from .install_consul_server import template_and_install_systemd


DEFAULTS = {
    "docker_registry_image": "registry:2.8.1",
}


@deploy("install docker")
def install_docker():
    linux_name = host.get_fact(LinuxName)
    if linux_name == "Fedora":
        dnf.packages(["docker", "docker-compose"])
        systemd.service("docker", enabled=True, running=True)
    else:
        apt.packages(["docker.io", "docker-compose"])


class TailscaleIPs(FactBase):
    requires_command = "tailscale"
    command = "tailscale ip"

    def process(self, output):
        # TODO provide ipaddress for nicer formatting in conf tools
        for line in output:
            if ":" in line:
                continue
            return [line]


class DockerImage(FactBase):
    requires_command = "docker"

    def command(self, object_id):
        return f"docker image inspect {object_id} || true"

    def process(self, output):
        joined_out = "".join(output)
        return json.loads(joined_out)


class DockerManifestInspect(FactBase):
    requires_command = "docker"

    def command(self, object_id):
        return f"docker image inspect {object_id} || true"

    def process(self, output):
        if "no such manifest" in output:
            return None
        joined_out = "".join(output)
        return json.loads(joined_out)


log = logging.getLogger(__name__)


def docker_image_from_host_to_target(image_reference: str):
    assert image_reference

    username = host.data.ssh_user
    hostname = host.name

    log.warning(
        "hello, sending image %r to host %s@%s", image_reference, username, hostname
    )

    with tempfile.NamedTemporaryFile() as f:
        cmdline = f"docker save {image_reference} | gzip | pv > {f.name}"
        log.warning("exec %r", cmdline)
        subprocess.check_output(cmdline, shell=True)

        with subprocess.Popen(["croc", "send", f.name], stderr=subprocess.PIPE) as proc:
            transfer_code = None
            for line_in in proc.stderr:
                line = line_in.decode()
                log.warning("got stdin line: %r", line)
                TRANSFER_CODE_PHRASE = "Code is: "
                transfer_code_index = line.find(TRANSFER_CODE_PHRASE)
                if transfer_code_index == -1:
                    continue
                transfer_code = line[
                    transfer_code_index + len(TRANSFER_CODE_PHRASE) :
                ].strip()
                assert len(transfer_code) > 10
                log.warning("extracted transfer code: %r", transfer_code)
                break
            assert transfer_code

            target_path = Path(f.name).name
            send_cmdline = f"croc --yes {transfer_code}"
            server.shell(send_cmdline, _chdir="/tmp")
        server.shell(
            f"cat {target_path} | docker load", _chdir="/tmp", name="load image file"
        )
        server.shell(f"rm /tmp/{target_path}", name="remove image file after importing")


@operation
def docker_image(image_reference: str):
    images = host.get_fact(DockerImage, image_reference)
    if not images:
        name, *_version = image_reference.split(":")
        if name in host.data.manual_docker_images:
            # get it from my machine lmao
            log.warning(
                "this deploy script wants image %r, taking it from host system and sending it",
                image_reference,
            )
            yield FunctionCommand(
                docker_image_from_host_to_target, (image_reference,), {}
            )
        else:
            # take it from given image ref
            yield f"docker pull {image_reference}"


def template_and_install_compose(
    compose_template_path: str,
    env_dict: Optional[dict] = None,
    *,
    systemd_service: Optional[str] = None,
):
    env_dict = env_dict or {}
    compose_template = Path(compose_template_path)
    systemd_service = systemd_service or compose_template.name.split(".")[0]
    assert systemd_service != "compose"
    assert systemd_service.endswith(".service")

    systemd_service_name = systemd_service.split(".")[0]

    working_directory = f"/opt/{systemd_service_name}"

    files.template(
        compose_template_path,
        f"{working_directory}/compose.yaml",
        env_dict=env_dict,
        name=f"sending compose file {systemd_service_name}",
    )

    template_and_install_systemd(
        "files/compose.service.j2",
        env_dict={
            "service_name": systemd_service_name,
            "working_directory": working_directory,
        },
        service_name=systemd_service,
    )


@deploy("install docker registry", data_defaults=DEFAULTS)
def install_registry():
    install_docker()
    docker_image(host.data.docker_registry_image)
    template_and_install_compose(
        "files/registry/compose.yaml.j2",
        {
            "docker_registry_image": host.data.docker_registry_image,
        },
        systemd_service="registry.service",
    )