docker-ssh

#!/usr/bin/python3
#
# docker-ssh: push and pull docker images over ssh
#
# usage:
#   docker-ssh pull URL IMAGE [IMAGE ...]    
#
#   docker-ssh push URL IMAGE [IMAGE ...]
#   docker-ssh push URL [URL ...] . IMAGE [IMAGE ...]
#
# This script uses `docker save` `docker load` and a ssh session to transfer
# images from/to a docker daemon running on a remote host.
#
# If possible it calls the `docker save` command with the --exclude option so
# as not to transfer images that are already present in the destination daemon.
#

import argparse
import concurrent.futures
import contextlib
import io
import itertools
import json
import logging
import os
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import threading
import time

PROG = "docker-ssh"

log = logging.getLogger(PROG)

def die(msg):
    sys.stderr.write("%s: %s\n" % (PROG, msg))
    sys.exit(1)

def log_exec(cmd):
    log.debug("exec: %s", ((" ".join(map(repr, cmd))) if isinstance(cmd, list) else cmd))

def check_output(cmd, *k, **kw):
    log_exec(cmd)
    return subprocess.check_output(cmd, *k, **kw)

def os_read_all(fd):
    lst = []
    while True:
        try:
            data = os.read(fd, 65536)
            if not data:
                return b"".join(lst)
            lst.append(data)
        except InterruptedError:
            pass

@contextlib.contextmanager
def process(cmd, *k, **kw):
    proc = None
    try:
        log_exec(cmd)
        proc = subprocess.Popen(cmd, *k, **kw)
        yield proc
        code = proc.wait()
        if code:
            raise subprocess.CalledProcessError(code, cmd)
    finally:
        if proc is not None and proc.poll() is None:
            try:
                log.debug("kill process: %r", cmd)
                proc.kill()
            except OSError:
                pass

class KillerThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor):
    def __init__(self, *k):
        super().__init__(*k)

        self._lock  = threading.Lock()
        self._procs = set()

    def shutdown(self, wait=True):
        with self._lock:
            for p in self._procs:
                if p.returncode is None:
                    try:
                        p.kill()
                    except OSError:
                        pass

        return super().shutdown(wait)

    @contextlib.contextmanager
    def __call__(self, process_manager):
        with self._lock, process_manager as proc:
            assert isinstance(proc, subprocess.Popen)
            self._procs.add(proc)
            self._lock.release()
            try:
                yield proc
            finally:
                self._lock.acquire()
                self._procs.remove(proc)


class Pipe:
    def __init__(self, pv = False):
        if pv and shutil.which("pv"):
            self.rd, self._wr = os.pipe()
            self._rd, self.wr = os.pipe()
        else:
            self.rd, self.wr = os.pipe()

    def close(self):
        os.close(self.rd)
        os.close(self.wr)

    def __enter__(self):
        if hasattr(self, "_rd"):
            self._proc = process(["pv"], stdin=self._rd, stdout=self._wr)
            self._proc.__enter__()
            os.close(self._rd)
            os.close(self._wr)
        return self

    def __exit__(self, *k):
        if hasattr(self, "_proc"):
            self._proc.__exit__(*k)


class Session:

    def cmd(self, args=[]):
        raise NotImplementedError()

    def __enter__(self):
        return self

    def wait_ready(self):
        pass

    def __exit__(self, a, b, c):
        pass

    def can_save_exclude(self):
        return bool(re.search(
            b"--exclude.* Layers not to be included",
            check_output(self.cmd(["docker", "save", "--help"]))))

    def can_load_print_excludes(self):
        return b"--print-excludes" in check_output(
                self.cmd(["docker", "load", "--help"]))

    def save(self, images, *, stdout=None, exclude=()):

        cmd = ["docker", "save"]
        for img in exclude:
            cmd.extend(("--exclude", img))
        cmd.append("--")
        cmd.extend(images)

        return process(self.cmd(cmd), stdout=stdout)

    @contextlib.contextmanager
    def load(self, *, stdin=None, stdout=None, print_excludes=False):

        cmd = ["docker", "load"]
        if print_excludes:
            cmd.append("--print-excludes")

        with process(self.cmd(cmd), stdin=stdin, stdout=stdout) as proc:
            if stdout is subprocess.PIPE:
                proc.result = lambda: proc.communicate()[0].decode().split()
            yield proc

    def tag(self, img, tag):
        check_output(self.cmd(["docker", "tag", "--", img, tag]))


class LocalSession(Session):
    def cmd(self, args=[]):
        return list(args)

class SshSession(Session):
    def __init__(self, user, host, port, url, *, identity):
        self.user = user
        self.host = host
        self.port = port
        self.proc = None
        self.url  = url
        self.identity = identity

    def cmd(self, args=[]):
        cmd = list(self.basecmd)
        cmd.append(self.host)
        if self.url is not None:
            cmd.extend(["/usr/bin/env", "DOCKER_HOST=" + self.url])
        assert isinstance(args, list)
        cmd.extend(args)
        return cmd

    def __enter__(self):
        assert self.proc is None
        self.tmpdir  = tempfile.TemporaryDirectory()
        self.tmpsock = os.path.join(self.tmpdir.__enter__(), "sock")

        self.basecmd = ["ssh", "-o", "ControlPath %s" % self.tmpsock, "-e", "none"]
        if self.user:
            self.basecmd.extend(["-l", self.user])
        if self.port:
            self.basecmd.extend(["-p", self.port])
        if self.identity:
            self.basecmd.extend(["-i", self.identity])

        cmd = self.basecmd + ["-MN", self.host]
        log_exec(cmd)
        self.proc = subprocess.Popen(cmd)
        return self

    def wait_ready(self):
        while not os.path.exists(self.tmpsock):
            time.sleep(0.1)
            if self.proc.poll() is not None:
                die("ssh connection error")

    def __exit__(self, a, b, c):
        try:
            self.proc.terminate()
            self.proc.wait()
        finally:
            self.tmpdir.cleanup()
            self.tmpdir.__exit__(a, b, c)


class UrlSession(Session):
    def __init__(self, url):
        self.url = url

    def cmd(self, args=[]):
        return ["/usr/bin/env", "DOCKER_HOST=" + self.url] + args

class MultiSession(Session):
    def __init__(self, sessions):
        self.sessions = list(sessions)
        self._stack = contextlib.ExitStack()

    def cmd(self, args=[]):
        raise NotImplementedError()

    def __enter__(self):
        try:
            for s in self.sessions:
                self._stack.enter_context(s)
        except:
            self._stack.close()
            raise

    def __exit__(self, *exc_info):
        self._stack.close()

    def _executor(self):
        return KillerThreadPoolExecutor(8)

    def can_load_print_excludes(self):
        with self._executor() as executor:
            return all(executor.map(
                lambda s: s.can_load_print_excludes(), self.sessions))

    @contextlib.contextmanager
    def load(self, *, stdin=None, stdout=None, print_excludes=False):

        if not isinstance(stdin, int):
            raise NotImplementedError()

        if print_excludes:
            with self._executor() as executor:

                def pe_thread():

                    def load_thread(session):
                        with executor(session.load(print_excludes=True,
                                stdin=subprocess.PIPE, stdout=subprocess.PIPE)) as proc:
                            proc.stdin.write(rawtar)
                            exclude = proc.result()
                            log.debug("(%s) exclude: %s", session.label, exclude)
                            return exclude

                    rawtar = os_read_all(fd)
                    os.close(fd)

                    exclude = None
                    for ex in executor.map(load_thread, self.sessions):
                        if exclude is None:
                            exclude = set(ex)
                        else:
                            exclude.intersection_update(ex)
                    
                    return list(exclude or ())


                fd = os.dup(stdin)
                future_result = executor.submit(pe_thread)

                class Proc:
                    result = lambda: future_result.result()
                yield Proc


        else:
            # not print_excludes
            with contextlib.ExitStack() as stack:
                tee_wr = []
                for session in self.sessions:
                    rd, wr = os.pipe()
                    stack.enter_context(session.load(stdin=rd))
                    os.close(rd)
                    tee_wr.append(wr)

                stack.enter_context(process(
                        ["tee"] + ["/dev/fd/%d" % fd for fd in tee_wr[1:]],
                        stdin=stdin, stdout=tee_wr[0], pass_fds=tee_wr))

                for fd in tee_wr:
                    os.close(fd)

                yield




def prepare_translation(src, images, add_pfx, rm_pfx):

    with src.save(images, exclude=["all"], stdout=subprocess.PIPE) as proc:
        rawtar = io.BytesIO(proc.stdout.read())
    tf = tarfile.open(fileobj=rawtar)

    tag_dict = {}
    image_ids = []
    for entry in json.loads(tf.extractfile("manifest.json").read().decode()):
        img  = "sha256:" + re.match(r"([0-9a-f]{64})\.json\Z", entry["Config"]).group(1)
        image_ids.append(img)

        for old_tag in entry["RepoTags"] or ():
            if rm_pfx:
                if not old_tag.startswith(rm_pfx):
                    die("tag %r does not start with prefix %r" % (old_tag, rm_pfx))
                new_tag = add_pfx + old_tag[len(rm_pfx):]
            else:
                new_tag = add_pfx + old_tag

            tag_dict[new_tag] = img
            log.info("tag %r translated to %r", old_tag, new_tag)

    return image_ids, tag_dict


def do_transfer(src, dst, images, add_pfx, rm_pfx):

    def abort_if_pfx(role):
        if add_pfx or rm_pfx:
            if role == "send":
                die("--add-prefix/--remove-prefix requires the sending daemon to support 'docker save --excludes=all'")
            else:
                die("--add-prefix/--remove-prefix requires the receiving daemon to support 'docker load --print-excludes'")


    tag_dict = {}
    exclude = ()

    if not src.can_save_exclude():
        abort_if_pfx("send")
        log.warning("source does not support 'docker save --exclude' ==> fallback to full image export")

    elif not dst.can_load_print_excludes():
        abort_if_pfx("recv")
        log.warning("destination does not support 'docker load --print-exclude' ==> fallback to full image export")

    else:
        if add_pfx or rm_pfx:
            images, tag_dict = prepare_translation(src, images, add_pfx, rm_pfx)

        with    Pipe() as pipe,                                     \
                src.save(images, exclude=["all"], stdout=pipe.wr),  \
                dst.load(print_excludes=True,     stdin =pipe.rd, stdout=subprocess.PIPE) as proc:
            pipe.close()
            exclude = proc.result()
            log.debug("(%s) exclude: %s", dst.label, exclude)


    with    Pipe(pv=True) as pipe,                                     \
            src.save(images, exclude=exclude, stdout=pipe.wr),  \
            dst.load(stdin=pipe.rd):
        pipe.close()

    # apply translation
    for tag, img in tag_dict.items():
        dst.tag(img, tag)


def create_session(url, identity):

    if url == "local:":
        session = LocalSession()

    elif re.match(r"[a-z]+://", url):
        session = UrlSession(url)

    else:
        mo = re.match(r"(?:([^:@/\s]+)@)?([^:@/\s]+)(?::([^:@/\s]+))?(?::([a-z]+://.*))?$", url, re.I)
        if not mo:
            die("remote host must match: [user@]host[:port][:docker_url] or docker_url")

        user, host, port, remote_url = mo.groups()
        session = SshSession(user, host, port, remote_url, identity=identity)

    session.label = url
    return session

def main():
    def loglevel(value):
        v = value.upper()
        if v not in ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"):
            raise ValueError()
        return getattr(logging, v)

    parser = argparse.ArgumentParser(prog=PROG,
        	    description="pull and push docker images over ssh")
    parser.add_argument("--log-level", default="WARNING", type=loglevel,
            help="log level (default WARNING)")

    sub = parser.add_subparsers(dest="command", help="sub-command help")

    def add_push_pull_cmd(name, help):
        extra = "" if name!="push" else  "\n\t%(prog)s URL [URL ...] . IMAGE [IMAGE ...]"
        p = sub.add_parser(name, help=help, usage=
                "\n\t%(prog)s URL IMAGE [IMAGE ...]" + extra)
        p.add_argument("--add-prefix", default="", metavar="PREFIX",
                help="prefix to be prepended to the image name (in the destination daemon)")
        p.add_argument("--remove-prefix", default="", metavar="PREFIX",
                help="prefix to be removed from the image name (in the destination daemon)")
        p.add_argument("--local", metavar="URL", default="local:",
                help="use an arbitrary docker daemon instead of the default local daemon")
        p.add_argument("-i", "--identity",
                help="select the private key used for ssh authentication")
        p.add_argument("remote", metavar="URL",
                help="""url identifying the remote node: either a docker url:
                unix://, tcp://, ... or a ssh account:
                [USER@]HOST[:PORT][:DOCKER_URL]""")
        p.add_argument("images", metavar="IMAGE", nargs="+",
                help="docker image (or repo) to be transferred")
    
    add_push_pull_cmd("push", help="push images to the remote host")
    add_push_pull_cmd("pull", help="pull images from the remote host")

    args = parser.parse_args()

    logging.basicConfig (level=args.log_level, format="%(levelname)-7s:%(name)s:%(message)s")

    if not args.command:
        parser.print_help()
        sys.exit(1)

    mksession = lambda *k: create_session(*k, identity=args.identity)

    local  = mksession(args.local)

    if "." not in args.images:
        remote = mksession(args.remote)
    else:
        # multiple remotes
        if args.command == "pull":
            die("pull does not accept multiple remotes")
        i = args.images.index(".")

        remote = MultiSession(itertools.chain(
            (mksession(args.remote),),
            map(mksession, args.images[:i])))
        remote.label = "multi:"
        args.images[:i+1] = ()

    if args.command == "pull":
        src, dst = remote, local
    elif args.command == "push":
        src, dst = local, remote
    else:
        raise NotImplementedError(args.command)
 
    with src, dst:
        src.wait_ready()
        dst.wait_ready()

        do_transfer(src, dst, args.images, args.add_prefix, args.remove_prefix)

main()

# vim:sw=4:et:sts=4:nosta: