docker-upgrade

#!/usr/bin/python3

import argparse
import collections
import contextlib
import logging
import os
import shlex
import signal
import socket
import sys
import tempfile
import time
from   typing import BinaryIO, Optional

import docker
from   docker.errors import APIError, NotFound

import pkg_resources

log = logging.getLogger("docker-upgrade")

SQUID_CONTAINER    = "docker-upgrade-squid"
SQUID_SPOOL_VOLUME = "docker-upgrade-squid-spool"
SQUID_LOG_VOLUME   = "docker-upgrade-squid-log"

@contextlib.contextmanager
def run_squid_container():
    client = docker.DockerClient()
    try:
        log.info("starting squid container")
        spool_volume = client.volumes.create(SQUID_SPOOL_VOLUME)
        log_volume   = client.volumes.create(SQUID_LOG_VOLUME)
        container = client.containers.run("ubuntu/squid:edge",
                name = SQUID_CONTAINER,
                detach = True,
                mounts = [
                    docker.types.Mount(target="/var/spool/squid", source=SQUID_SPOOL_VOLUME),
                    docker.types.Mount(target="/var/log/squid", source=SQUID_LOG_VOLUME),
                    ])
        container.reload()
        ip = container.attrs["NetworkSettings"]["Networks"]["bridge"]["IPAddress"]
        for _ in range(50):
            with contextlib.suppress(ConnectionRefusedError):
                socket.create_connection((ip, 3128))
                break
            time.sleep(.1)
        else:
            raise RuntimeError("unable to connect to squid port 3128")
        yield container.id
    finally:
        signal.signal(signal.SIGINT, lambda *k: print("SIGINT"))
        log.info("stopping squid container")
        for func in (
                lambda: container.remove(force=True),
                lambda: spool_volume.remove(),
                lambda: log_volume.remove(),
                ):
            try:
                func()
            except UnboundLocalError:
                pass
            except Exception:
                log.exception("squid cleanup error")

def fmt_image_id(image_id: str):
    return image_id.strip("sha256:")[:12]

def fmt_image(image_id: str, tags: list[str]):
    short_id = fmt_image_id(image_id)
    return f"{short_id} [{' '.join(tags)}]" if tags else short_id

def load_upgrade_script() -> bytes:
    
    path = pkg_resources.resource_filename("docker_utils_aba", "docker-upgrade-script.sh")

    with open(path, "rb") as fp:
        return fp.read()

def write_to_stdin(client: docker.APIClient, container_id: str, data: bytes):
    sock = client.attach_socket(container_id, dict(stdin=True, stream=True))
    offset = 0
    while True:
        chunk = data[offset:]
        if not chunk:
            sock.close()
            return
        offset += sock._sock.send(chunk)

def make_fifo(path: str) -> BinaryIO:
    """Create a fifo at the given path and open it in non-blocking mode"""

    os.mkfifo(path, mode=0o600)
    fd = os.open(path, os.O_RDONLY | os.O_NONBLOCK)
    return os.fdopen(fd, "rb")

# return true if the image was upgraded
def do_upgrade(client: docker.APIClient, image_id: str, tags: list[str], *, force: bool = False,
        script: bytes, tmp_dir: str, squid_container: Optional[str]) -> bool:

    desc = fmt_image(image_id, tags)

    fifo_path = os.path.join(tmp_dir, image_id.replace(":", "_"))
    cid = None
    fifo = None
    try:
        if log.getEffectiveLevel() < logging.INFO:
            log.info("============== upgrading %s ==============", desc)

        # create the fifo and open it for reading
        fifo = make_fifo(fifo_path)

        # get the current config of the image
        orig_config = client.inspect_image(image_id)["Config"]

        # create and start the temporary container
        cmd = ["/bin/sh"]
        if log.getEffectiveLevel() <= logging.DEBUG:
            # trace shell commands with -vv
            cmd.append("-x")
        cid = client.create_container(image=image_id, user="0:0", entrypoint=[],
                command=cmd, stdin_open=True, host_config=client.create_host_config(
                    binds={fifo_path: {"bind": "/.docker-upgrade-fifo", "mode": "rw"}},
                    network_mode=("bridge" if squid_container is None
                        else f"container:{squid_container}"),
                ))["Id"]
        client.start(cid)

        # write the script to stdin
        write_to_stdin(client, cid, script)

        # wait until the script terminates
        returncode = client.wait(cid)["StatusCode"]

        # get the status line from the fifo
        status = fifo.read().strip().decode(errors="replace")

        # get the output
        output = client.logs(cid)

        os_supported = status != "not_supported"

        # dump the container output on:
        #   -vv (always)
        #   -v  (except for os_not_supported errors)
        if log.getEffectiveLevel() < logging.INFO:
            for line in output.decode(errors="replace").splitlines():
                if line:
                    log.info("output: %s", line)

        if returncode:
            if not os_supported:
                log.error("failed to upgrade %s (os not supported)", desc)
            else:
                log.error("failed to upgrade %s (%s, exit %s)", desc, status, returncode)

        elif not output and not force:
            log.info("image not upgraded: %s (%s)", desc, status)
        else:
            # commit the image
            log.debug("committing new image for %s", desc)
            new_id = client.commit(cid,
                conf = orig_config,
                message = "automatic upgrade, %s" % time.ctime(),
            )["Id"]

            # update the tags
            for tag in tags:
                repo, tag = docker.utils.parse_repository_tag(tag)
                assert tag is not None
                client.tag(image=new_id, repository=repo, tag=tag, force=True)

            log.info("image upgraded %s -> %s (%s)", desc, fmt_image_id(new_id), status)
            return True

    except APIError as e:
        log.error("failed to upgrade %s (%s)", desc, e)
    finally:
        if cid is not None:
            try:
                client.remove_container(cid, v=True, force=True)
            except APIError as e:
                log.warning("failed to remove temporary container %s (%s)", cid, e)
        if fifo is not None:
            fifo.close()
    return False

def stop_running_containers(client: docker.DockerClient, images: set[str]):
    for ctr in client.containers.list():
        ctr_image = ctr.attrs["Config"]["Image"]
        if ctr.status == "running" and ctr_image in images:
            log.info("stopping %s (image %s)", ctr.name, ctr_image)
            try:
                ctr.stop()
            except APIError as e:
                log.error("failed to stop %s (%s)", ctr.name, e)

def expand_tag(image: str) -> str:
    """parse an image name and append a ':latest' if the image does not have any tag"""

    repo, tag = docker.utils.parse_repository_tag(image)

    return f"{repo}:{'latest' if tag is None else tag}"


def main():
    parser = argparse.ArgumentParser(description="upgrade docker images",
            epilog="""This command is intended for security upgrades. For each image, it runs a
            container and attempts to perform a system upgrade. If the upgrade is effective, then a
            new image is committed, replacing the old one. It supports distributions based on debian
            (apt), alpine (apk) and redhat (yum/dnf).""")

    parser.add_argument("images", metavar="IMAGE", nargs="+",
            help="image name (with wildcard expansion)")
    parser.add_argument("-f", "--force", action="store_true",
            help="force committing a new image even if there is no reported upgrades")
    parser.add_argument("--ignore-unknown", action="store_true",
            help = "silently ignore unknown images listed in the command line")
    parser.add_argument("--ignore-inspect-errors", action="store_true",
            help = "do the upgrades even if there are image inspect errors (by default the script "
            "aborts as soon as any image fails to be inspected)")

    proxy_group = parser.add_mutually_exclusive_group()
    proxy_group.add_argument("--http-proxy", metavar="URL",
            help="""value for the `http_proxy` environment variable to be set in the upgrade
            container""")
    proxy_group.add_argument("--squid", action="store_true",
            help="""launch a squid container in the background to act as a HTTP proxy""")
    parser.add_argument("--https-proxy", metavar="URL",
            help="""value for the `https_proxy` environment variable to be set in the upgrade
            container""")

    parser.add_argument("--stop", action="store_true",
            help="stop running containers whose image was upgraded")

    parser.add_argument("-q", "--quiet",   action="store_true", help="decrease verbosity")
    parser.add_argument("-v", "--verbose", action="count", default=0,
            help="increase verbosity")

    args = parser.parse_args()
    log_level = (
            logging.DEBUG   if args.verbose >  1 else
            logging.INFO-1  if args.verbose == 1 else
            logging.INFO    if not args.quiet    else
            logging.WARNING)

    logging.basicConfig(level=log_level)
    
    have_error = False
    client = docker.APIClient()
    docker_client = docker.DockerClient()

    # prepare the list of images to be upgraded

    # key:    image id
    # values: repo:tag
    images : dict[str, list[str]] = collections.defaultdict(list)
    for image in args.images:
        expected_repo_tag = expand_tag(image)
        try:
            img = client.inspect_image(image)
        except APIError as e:
            if args.ignore_unknown and isinstance(e, NotFound):
                log.debug("skipping unknown image %r", image)
            else:
                log.error("failed to inspect image %r (%s)", image, e)
                if not args.ignore_inspect_errors:
                    have_error = True
            continue

        lst = images[img["Id"]]
        if expected_repo_tag in img["RepoTags"] and expected_repo_tag not in lst:
            lst.append(expected_repo_tag)

    http_proxy = f"http://127.0.0.1:3128/" if args.squid else args.http_proxy
    try:
        script = load_upgrade_script()
        log.debug("upgrade script loaded (%d bytes)", len(script))
        if http_proxy is not None:
            log.info("using http_proxy: %s", http_proxy)
            script = f"export http_proxy={shlex.quote(http_proxy)}\n".encode() + script
        if args.https_proxy is not None:
            log.info("using https_proxy: %s", args.https_proxy)
            script = f"export https_proxy={shlex.quote(args.https_proxy)}\n".encode() + script
    except Exception as e:
        log.error("failed to load the upgrade script (%s)", e)
        have_error = True
    if have_error:
        return 1

    with contextlib.ExitStack() as stack:
        stop_set = set()
        if args.stop:
            stack.callback(stop_running_containers, docker_client, stop_set)
        tmp_dir = stack.enter_context(tempfile.TemporaryDirectory())
        squid_container = stack.enter_context(run_squid_container()) if args.squid else None

        for image_id, tags in sorted(images.items(),
                # order by first tag (and put all untagged images unsorted at the end)
                key=lambda x: x[1][0] if x[1] else "~"):

            if do_upgrade(client, image_id, tags, force=args.force, script=script, tmp_dir=tmp_dir,
                    squid_container=squid_container):
                stop_set.update(tags)

sys.exit(main())