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