tests/docker/docker.py

   1 #!/usr/bin/env python2
   2 #
   3 # Docker controlling module
   4 #
   5 # Copyright (c) 2016 Red Hat Inc.
   6 #
   7 # Authors:
   8 #  Fam Zheng <famz@redhat.com>
   9 #
  10 # This work is licensed under the terms of the GNU GPL, version 2
  11 # or (at your option) any later version. See the COPYING file in
  12 # the top-level directory.
  13
  14 import os
  15 import sys
  16 import subprocess
  17 import json
  18 import hashlib
  19 import atexit
  20 import uuid
  21 import argparse
  22 import tempfile
  23 from shutil import copy
  24
  25 def _text_checksum(text):
  26     """Calculate a digest string unique to the text content"""
  27     return hashlib.sha1(text).hexdigest()
  28
  29 def _guess_docker_command():
  30     """ Guess a working docker command or raise exception if not found"""
  31     commands = [["docker"], ["sudo", "-n", "docker"]]
  32     for cmd in commands:
  33         if subprocess.call(cmd + ["images"],
  34                            stdout=subprocess.PIPE,
  35                            stderr=subprocess.PIPE) == 0:
  36             return cmd
  37     commands_txt = "\n".join(["  " + " ".join(x) for x in commands])
  38     raise Exception("Cannot find working docker command. Tried:\n%s" % \
  39                     commands_txt)
  40
  41 class Docker(object):
  42     """ Running Docker commands """
  43     def __init__(self):
  44         self._command = _guess_docker_command()
  45         self._instances = []
  46         atexit.register(self._kill_instances)
  47
  48     def _do(self, cmd, quiet=True, **kwargs):
  49         if quiet:
  50             kwargs["stdout"] = subprocess.PIPE
  51         return subprocess.call(self._command + cmd, **kwargs)
  52
  53     def _do_kill_instances(self, only_known, only_active=True):
  54         cmd = ["ps", "-q"]
  55         if not only_active:
  56             cmd.append("-a")
  57         for i in self._output(cmd).split():
  58             resp = self._output(["inspect", i])
  59             labels = json.loads(resp)[0]["Config"]["Labels"]
  60             active = json.loads(resp)[0]["State"]["Running"]
  61             if not labels:
  62                 continue
  63             instance_uuid = labels.get("com.qemu.instance.uuid", None)
  64             if not instance_uuid:
  65                 continue
  66             if only_known and instance_uuid not in self._instances:
  67                 continue
  68             print "Terminating", i
  69             if active:
  70                 self._do(["kill", i])
  71             self._do(["rm", i])
  72
  73     def clean(self):
  74         self._do_kill_instances(False, False)
  75         return 0
  76
  77     def _kill_instances(self):
  78         return self._do_kill_instances(True)
  79
  80     def _output(self, cmd, **kwargs):
  81         return subprocess.check_output(self._command + cmd,
  82                                        stderr=subprocess.STDOUT,
  83                                        **kwargs)
  84
  85     def get_image_dockerfile_checksum(self, tag):
  86         resp = self._output(["inspect", tag])
  87         labels = json.loads(resp)[0]["Config"].get("Labels", {})
  88         return labels.get("com.qemu.dockerfile-checksum", "")
  89
  90     def build_image(self, tag, dockerfile, df_path, quiet=True, argv=None):
  91         if argv == None:
  92             argv = []
  93         tmp_dir = tempfile.mkdtemp(prefix="docker_build")
  94
  95         tmp_df = tempfile.NamedTemporaryFile(dir=tmp_dir, suffix=".docker")
  96         tmp_df.write(dockerfile)
  97
  98         tmp_df.write("\n")
  99         tmp_df.write("LABEL com.qemu.dockerfile-checksum=%s" %
 100                      _text_checksum(dockerfile))
 101         tmp_df.flush()
 102         self._do(["build", "-t", tag, "-f", tmp_df.name] + argv + \
 103                  [tmp_dir],
 104                  quiet=quiet)
 105
 106     def image_matches_dockerfile(self, tag, dockerfile):
 107         try:
 108             checksum = self.get_image_dockerfile_checksum(tag)
 109         except Exception:
 110             return False
 111         return checksum == _text_checksum(dockerfile)
 112
 113     def run(self, cmd, keep, quiet):
 114         label = uuid.uuid1().hex
 115         if not keep:
 116             self._instances.append(label)
 117         ret = self._do(["run", "--label",
 118                         "com.qemu.instance.uuid=" + label] + cmd,
 119                        quiet=quiet)
 120         if not keep:
 121             self._instances.remove(label)
 122         return ret
 123
 124 class SubCommand(object):
 125     """A SubCommand template base class"""
 126     name = None # Subcommand name
 127     def shared_args(self, parser):
 128         parser.add_argument("--quiet", action="store_true",
 129                             help="Run quietly unless an error occured")
 130
 131     def args(self, parser):
 132         """Setup argument parser"""
 133         pass
 134     def run(self, args, argv):
 135         """Run command.
 136         args: parsed argument by argument parser.
 137         argv: remaining arguments from sys.argv.
 138         """
 139         pass
 140
 141 class RunCommand(SubCommand):
 142     """Invoke docker run and take care of cleaning up"""
 143     name = "run"
 144     def args(self, parser):
 145         parser.add_argument("--keep", action="store_true",
 146                             help="Don't remove image when command completes")
 147     def run(self, args, argv):
 148         return Docker().run(argv, args.keep, quiet=args.quiet)
 149
 150 class BuildCommand(SubCommand):
 151     """ Build docker image out of a dockerfile. Arguments: <tag> <dockerfile>"""
 152     name = "build"
 153     def args(self, parser):
 154         parser.add_argument("tag",
 155                             help="Image Tag")
 156         parser.add_argument("dockerfile",
 157                             help="Dockerfile name")
 158
 159     def run(self, args, argv):
 160         dockerfile = open(args.dockerfile, "rb").read()
 161         tag = args.tag
 162
 163         dkr = Docker()
 164         if dkr.image_matches_dockerfile(tag, dockerfile):
 165             if not args.quiet:
 166                 print "Image is up to date."
 167             return 0
 168
 169         dkr.build_image(tag, dockerfile, args.dockerfile,
 170                         quiet=args.quiet, argv=argv)
 171         return 0
 172
 173 class CleanCommand(SubCommand):
 174     """Clean up docker instances"""
 175     name = "clean"
 176     def run(self, args, argv):
 177         Docker().clean()
 178         return 0
 179
 180 def main():
 181     parser = argparse.ArgumentParser(description="A Docker helper",
 182             usage="%s <subcommand> ..." % os.path.basename(sys.argv[0]))
 183     subparsers = parser.add_subparsers(title="subcommands", help=None)
 184     for cls in SubCommand.__subclasses__():
 185         cmd = cls()
 186         subp = subparsers.add_parser(cmd.name, help=cmd.__doc__)
 187         cmd.shared_args(subp)
 188         cmd.args(subp)
 189         subp.set_defaults(cmdobj=cmd)
 190     args, argv = parser.parse_known_args()
 191     return args.cmdobj.run(args, argv)
 192
 193 if __name__ == "__main__":
 194     sys.exit(main())