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