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