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 import re
  24 import signal
  25 from tarfile import TarFile, TarInfo
  26 from StringIO import StringIO
  27 from shutil import copy, rmtree
  28
  29
  30 DEVNULL = open(os.devnull, 'wb')
  31
  32
  33 def _text_checksum(text):
  34     """Calculate a digest string unique to the text content"""
  35     return hashlib.sha1(text).hexdigest()
  36
  37 def _guess_docker_command():
  38     """ Guess a working docker command or raise exception if not found"""
  39     commands = [["docker"], ["sudo", "-n", "docker"]]
  40     for cmd in commands:
  41         try:
  42             if subprocess.call(cmd + ["images"],
  43                                stdout=DEVNULL, stderr=DEVNULL) == 0:
  44                 return cmd
  45         except OSError:
  46             pass
  47     commands_txt = "\n".join(["  " + " ".join(x) for x in commands])
  48     raise Exception("Cannot find working docker command. Tried:\n%s" % \
  49                     commands_txt)
  50
  51 def _copy_with_mkdir(src, root_dir, sub_path):
  52     """Copy src into root_dir, creating sub_path as needed."""
  53     dest_dir = os.path.normpath("%s/%s" % (root_dir, sub_path))
  54     try:
  55         os.makedirs(dest_dir)
  56     except OSError:
  57         # we can safely ignore already created directories
  58         pass
  59
  60     dest_file = "%s/%s" % (dest_dir, os.path.basename(src))
  61     copy(src, dest_file)
  62
  63
  64 def _get_so_libs(executable):
  65     """Return a list of libraries associated with an executable.
  66
  67     The paths may be symbolic links which would need to be resolved to
  68     ensure theright data is copied."""
  69
  70     libs = []
  71     ldd_re = re.compile(r"(/.*/)(\S*)")
  72     try:
  73         ldd_output = subprocess.check_output(["ldd", executable])
  74         for line in ldd_output.split("\n"):
  75             search = ldd_re.search(line)
  76             if search and len(search.groups()) == 2:
  77                 so_path = search.groups()[0]
  78                 so_lib = search.groups()[1]
  79                 libs.append("%s/%s" % (so_path, so_lib))
  80     except subprocess.CalledProcessError:
  81         print "%s had no associated libraries (static build?)" % (executable)
  82
  83     return libs
  84
  85 def _copy_binary_with_libs(src, dest_dir):
  86     """Copy a binary executable and all its dependant libraries.
  87
  88     This does rely on the host file-system being fairly multi-arch
  89     aware so the file don't clash with the guests layout."""
  90
  91     _copy_with_mkdir(src, dest_dir, "/usr/bin")
  92
  93     libs = _get_so_libs(src)
  94     if libs:
  95         for l in libs:
  96             so_path = os.path.dirname(l)
  97             _copy_with_mkdir(l , dest_dir, so_path)
  98
  99 class Docker(object):
 100     """ Running Docker commands """
 101     def __init__(self):
 102         self._command = _guess_docker_command()
 103         self._instances = []
 104         atexit.register(self._kill_instances)
 105         signal.signal(signal.SIGTERM, self._kill_instances)
 106         signal.signal(signal.SIGHUP, self._kill_instances)
 107
 108     def _do(self, cmd, quiet=True, infile=None, **kwargs):
 109         if quiet:
 110             kwargs["stdout"] = DEVNULL
 111         if infile:
 112             kwargs["stdin"] = infile
 113         return subprocess.call(self._command + cmd, **kwargs)
 114
 115     def _do_kill_instances(self, only_known, only_active=True):
 116         cmd = ["ps", "-q"]
 117         if not only_active:
 118             cmd.append("-a")
 119         for i in self._output(cmd).split():
 120             resp = self._output(["inspect", i])
 121             labels = json.loads(resp)[0]["Config"]["Labels"]
 122             active = json.loads(resp)[0]["State"]["Running"]
 123             if not labels:
 124                 continue
 125             instance_uuid = labels.get("com.qemu.instance.uuid", None)
 126             if not instance_uuid:
 127                 continue
 128             if only_known and instance_uuid not in self._instances:
 129                 continue
 130             print "Terminating", i
 131             if active:
 132                 self._do(["kill", i])
 133             self._do(["rm", i])
 134
 135     def clean(self):
 136         self._do_kill_instances(False, False)
 137         return 0
 138
 139     def _kill_instances(self, *args, **kwargs):
 140         return self._do_kill_instances(True)
 141
 142     def _output(self, cmd, **kwargs):
 143         return subprocess.check_output(self._command + cmd,
 144                                        stderr=subprocess.STDOUT,
 145                                        **kwargs)
 146
 147     def get_image_dockerfile_checksum(self, tag):
 148         resp = self._output(["inspect", tag])
 149         labels = json.loads(resp)[0]["Config"].get("Labels", {})
 150         return labels.get("com.qemu.dockerfile-checksum", "")
 151
 152     def build_image(self, tag, docker_dir, dockerfile, quiet=True, argv=None):
 153         if argv == None:
 154             argv = []
 155
 156         tmp_df = tempfile.NamedTemporaryFile(dir=docker_dir, suffix=".docker")
 157         tmp_df.write(dockerfile)
 158
 159         tmp_df.write("\n")
 160         tmp_df.write("LABEL com.qemu.dockerfile-checksum=%s" %
 161                      _text_checksum(dockerfile))
 162         tmp_df.flush()
 163
 164         self._do(["build", "-t", tag, "-f", tmp_df.name] + argv + \
 165                  [docker_dir],
 166                  quiet=quiet)
 167
 168     def update_image(self, tag, tarball, quiet=True):
 169         "Update a tagged image using "
 170
 171         self._do(["build", "-t", tag, "-"], quiet=quiet, infile=tarball)
 172
 173     def image_matches_dockerfile(self, tag, dockerfile):
 174         try:
 175             checksum = self.get_image_dockerfile_checksum(tag)
 176         except Exception:
 177             return False
 178         return checksum == _text_checksum(dockerfile)
 179
 180     def run(self, cmd, keep, quiet):
 181         label = uuid.uuid1().hex
 182         if not keep:
 183             self._instances.append(label)
 184         ret = self._do(["run", "--label",
 185                         "com.qemu.instance.uuid=" + label] + cmd,
 186                        quiet=quiet)
 187         if not keep:
 188             self._instances.remove(label)
 189         return ret
 190
 191     def command(self, cmd, argv, quiet):
 192         return self._do([cmd] + argv, quiet=quiet)
 193
 194 class SubCommand(object):
 195     """A SubCommand template base class"""
 196     name = None # Subcommand name
 197     def shared_args(self, parser):
 198         parser.add_argument("--quiet", action="store_true",
 199                             help="Run quietly unless an error occured")
 200
 201     def args(self, parser):
 202         """Setup argument parser"""
 203         pass
 204     def run(self, args, argv):
 205         """Run command.
 206         args: parsed argument by argument parser.
 207         argv: remaining arguments from sys.argv.
 208         """
 209         pass
 210
 211 class RunCommand(SubCommand):
 212     """Invoke docker run and take care of cleaning up"""
 213     name = "run"
 214     def args(self, parser):
 215         parser.add_argument("--keep", action="store_true",
 216                             help="Don't remove image when command completes")
 217     def run(self, args, argv):
 218         return Docker().run(argv, args.keep, quiet=args.quiet)
 219
 220 class BuildCommand(SubCommand):
 221     """ Build docker image out of a dockerfile. Arguments: <tag> <dockerfile>"""
 222     name = "build"
 223     def args(self, parser):
 224         parser.add_argument("--include-executable", "-e",
 225                             help="""Specify a binary that will be copied to the
 226                             container together with all its dependent
 227                             libraries""")
 228         parser.add_argument("tag",
 229                             help="Image Tag")
 230         parser.add_argument("dockerfile",
 231                             help="Dockerfile name")
 232
 233     def run(self, args, argv):
 234         dockerfile = open(args.dockerfile, "rb").read()
 235         tag = args.tag
 236
 237         dkr = Docker()
 238         if dkr.image_matches_dockerfile(tag, dockerfile):
 239             if not args.quiet:
 240                 print "Image is up to date."
 241         else:
 242             # Create a docker context directory for the build
 243             docker_dir = tempfile.mkdtemp(prefix="docker_build")
 244
 245             # Is there a .pre file to run in the build context?
 246             docker_pre = os.path.splitext(args.dockerfile)[0]+".pre"
 247             if os.path.exists(docker_pre):
 248                 stdout = DEVNULL if args.quiet else None
 249                 rc = subprocess.call(os.path.realpath(docker_pre),
 250                                      cwd=docker_dir, stdout=stdout)
 251                 if rc == 3:
 252                     print "Skip"
 253                     return 0
 254                 elif rc != 0:
 255                     print "%s exited with code %d" % (docker_pre, rc)
 256                     return 1
 257
 258             # Do we include a extra binary?
 259             if args.include_executable:
 260                 _copy_binary_with_libs(args.include_executable,
 261                                        docker_dir)
 262
 263             dkr.build_image(tag, docker_dir, dockerfile,
 264                             quiet=args.quiet, argv=argv)
 265
 266             rmtree(docker_dir)
 267
 268         return 0
 269
 270 class UpdateCommand(SubCommand):
 271     """ Update a docker image with new executables. Arguments: <tag> <executable>"""
 272     name = "update"
 273     def args(self, parser):
 274         parser.add_argument("tag",
 275                             help="Image Tag")
 276         parser.add_argument("executable",
 277                             help="Executable to copy")
 278
 279     def run(self, args, argv):
 280         # Create a temporary tarball with our whole build context and
 281         # dockerfile for the update
 282         tmp = tempfile.NamedTemporaryFile(suffix="dckr.tar.gz")
 283         tmp_tar = TarFile(fileobj=tmp, mode='w')
 284
 285         # Add the executable to the tarball
 286         bn = os.path.basename(args.executable)
 287         ff = "/usr/bin/%s" % bn
 288         tmp_tar.add(args.executable, arcname=ff)
 289
 290         # Add any associated libraries
 291         libs = _get_so_libs(args.executable)
 292         if libs:
 293             for l in libs:
 294                 tmp_tar.add(os.path.realpath(l), arcname=l)
 295
 296         # Create a Docker buildfile
 297         df = StringIO()
 298         df.write("FROM %s\n" % args.tag)
 299         df.write("ADD . /\n")
 300         df.seek(0)
 301
 302         df_tar = TarInfo(name="Dockerfile")
 303         df_tar.size = len(df.buf)
 304         tmp_tar.addfile(df_tar, fileobj=df)
 305
 306         tmp_tar.close()
 307
 308         # reset the file pointers
 309         tmp.flush()
 310         tmp.seek(0)
 311
 312         # Run the build with our tarball context
 313         dkr = Docker()
 314         dkr.update_image(args.tag, tmp, quiet=args.quiet)
 315
 316         return 0
 317
 318 class CleanCommand(SubCommand):
 319     """Clean up docker instances"""
 320     name = "clean"
 321     def run(self, args, argv):
 322         Docker().clean()
 323         return 0
 324
 325 class ImagesCommand(SubCommand):
 326     """Run "docker images" command"""
 327     name = "images"
 328     def run(self, args, argv):
 329         return Docker().command("images", argv, args.quiet)
 330
 331 def main():
 332     parser = argparse.ArgumentParser(description="A Docker helper",
 333             usage="%s <subcommand> ..." % os.path.basename(sys.argv[0]))
 334     subparsers = parser.add_subparsers(title="subcommands", help=None)
 335     for cls in SubCommand.__subclasses__():
 336         cmd = cls()
 337         subp = subparsers.add_parser(cmd.name, help=cmd.__doc__)
 338         cmd.shared_args(subp)
 339         cmd.args(subp)
 340         subp.set_defaults(cmdobj=cmd)
 341     args, argv = parser.parse_known_args()
 342     return args.cmdobj.run(args, argv)
 343
 344 if __name__ == "__main__":
 345     sys.exit(main())