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, infile=None, **kwargs):
 116         if quiet:
 117             kwargs["stdout"] = DEVNULL
 118         if infile:
 119             kwargs["stdin"] = infile
 120         return subprocess.call(self._command + cmd, **kwargs)
 121
 122     def _do_kill_instances(self, only_known, only_active=True):
 123         cmd = ["ps", "-q"]
 124         if not only_active:
 125             cmd.append("-a")
 126         for i in self._output(cmd).split():
 127             resp = self._output(["inspect", i])
 128             labels = json.loads(resp)[0]["Config"]["Labels"]
 129             active = json.loads(resp)[0]["State"]["Running"]
 130             if not labels:
 131                 continue
 132             instance_uuid = labels.get("com.qemu.instance.uuid", None)
 133             if not instance_uuid:
 134                 continue
 135             if only_known and instance_uuid not in self._instances:
 136                 continue
 137             print "Terminating", i
 138             if active:
 139                 self._do(["kill", i])
 140             self._do(["rm", i])
 141
 142     def clean(self):
 143         self._do_kill_instances(False, False)
 144         return 0
 145
 146     def _kill_instances(self, *args, **kwargs):
 147         return self._do_kill_instances(True)
 148
 149     def _output(self, cmd, **kwargs):
 150         return subprocess.check_output(self._command + cmd,
 151                                        stderr=subprocess.STDOUT,
 152                                        **kwargs)
 153
 154     def get_image_dockerfile_checksum(self, tag):
 155         resp = self._output(["inspect", tag])
 156         labels = json.loads(resp)[0]["Config"].get("Labels", {})
 157         return labels.get("com.qemu.dockerfile-checksum", "")
 158
 159     def build_image(self, tag, docker_dir, dockerfile,
 160                     quiet=True, user=False, argv=None, extra_files_cksum=[]):
 161         if argv == None:
 162             argv = []
 163
 164         tmp_df = tempfile.NamedTemporaryFile(dir=docker_dir, suffix=".docker")
 165         tmp_df.write(dockerfile)
 166
 167         if user:
 168             uid = os.getuid()
 169             uname = getpwuid(uid).pw_name
 170             tmp_df.write("\n")
 171             tmp_df.write("RUN id %s 2>/dev/null || useradd -u %d -U %s" %
 172                          (uname, uid, uname))
 173
 174         tmp_df.write("\n")
 175         tmp_df.write("LABEL com.qemu.dockerfile-checksum=%s" %
 176                      _text_checksum("\n".join([dockerfile] +
 177                                     extra_files_cksum)))
 178         tmp_df.flush()
 179
 180         self._do(["build", "-t", tag, "-f", tmp_df.name] + argv + \
 181                  [docker_dir],
 182                  quiet=quiet)
 183
 184     def update_image(self, tag, tarball, quiet=True):
 185         "Update a tagged image using "
 186
 187         self._do(["build", "-t", tag, "-"], quiet=quiet, infile=tarball)
 188
 189     def image_matches_dockerfile(self, tag, dockerfile):
 190         try:
 191             checksum = self.get_image_dockerfile_checksum(tag)
 192         except Exception:
 193             return False
 194         return checksum == _text_checksum(dockerfile)
 195
 196     def run(self, cmd, keep, quiet):
 197         label = uuid.uuid1().hex
 198         if not keep:
 199             self._instances.append(label)
 200         ret = self._do(["run", "--label",
 201                         "com.qemu.instance.uuid=" + label] + cmd,
 202                        quiet=quiet)
 203         if not keep:
 204             self._instances.remove(label)
 205         return ret
 206
 207     def command(self, cmd, argv, quiet):
 208         return self._do([cmd] + argv, quiet=quiet)
 209
 210 class SubCommand(object):
 211     """A SubCommand template base class"""
 212     name = None # Subcommand name
 213     def shared_args(self, parser):
 214         parser.add_argument("--quiet", action="store_true",
 215                             help="Run quietly unless an error occured")
 216
 217     def args(self, parser):
 218         """Setup argument parser"""
 219         pass
 220     def run(self, args, argv):
 221         """Run command.
 222         args: parsed argument by argument parser.
 223         argv: remaining arguments from sys.argv.
 224         """
 225         pass
 226
 227 class RunCommand(SubCommand):
 228     """Invoke docker run and take care of cleaning up"""
 229     name = "run"
 230     def args(self, parser):
 231         parser.add_argument("--keep", action="store_true",
 232                             help="Don't remove image when command completes")
 233     def run(self, args, argv):
 234         return Docker().run(argv, args.keep, quiet=args.quiet)
 235
 236 class BuildCommand(SubCommand):
 237     """ Build docker image out of a dockerfile. Arguments: <tag> <dockerfile>"""
 238     name = "build"
 239     def args(self, parser):
 240         parser.add_argument("--include-executable", "-e",
 241                             help="""Specify a binary that will be copied to the
 242                             container together with all its dependent
 243                             libraries""")
 244         parser.add_argument("--extra-files", "-f", nargs='*',
 245                             help="""Specify files that will be copied in the
 246                             Docker image, fulfilling the ADD directive from the
 247                             Dockerfile""")
 248         parser.add_argument("--add-current-user", "-u", dest="user",
 249                             action="store_true",
 250                             help="Add the current user to image's passwd")
 251         parser.add_argument("tag",
 252                             help="Image Tag")
 253         parser.add_argument("dockerfile",
 254                             help="Dockerfile name")
 255
 256     def run(self, args, argv):
 257         dockerfile = open(args.dockerfile, "rb").read()
 258         tag = args.tag
 259
 260         dkr = Docker()
 261         if dkr.image_matches_dockerfile(tag, dockerfile):
 262             if not args.quiet:
 263                 print "Image is up to date."
 264         else:
 265             # Create a docker context directory for the build
 266             docker_dir = tempfile.mkdtemp(prefix="docker_build")
 267
 268             # Is there a .pre file to run in the build context?
 269             docker_pre = os.path.splitext(args.dockerfile)[0]+".pre"
 270             if os.path.exists(docker_pre):
 271                 stdout = DEVNULL if args.quiet else None
 272                 rc = subprocess.call(os.path.realpath(docker_pre),
 273                                      cwd=docker_dir, stdout=stdout)
 274                 if rc == 3:
 275                     print "Skip"
 276                     return 0
 277                 elif rc != 0:
 278                     print "%s exited with code %d" % (docker_pre, rc)
 279                     return 1
 280
 281             # Copy any extra files into the Docker context. These can be
 282             # included by the use of the ADD directive in the Dockerfile.
 283             cksum = []
 284             if args.include_executable:
 285                 # FIXME: there is no checksum of this executable and the linked
 286                 # libraries, once the image built any change of this executable
 287                 # or any library won't trigger another build.
 288                 _copy_binary_with_libs(args.include_executable, docker_dir)
 289             for filename in args.extra_files or []:
 290                 _copy_with_mkdir(filename, docker_dir)
 291                 cksum += [_file_checksum(filename)]
 292
 293             argv += ["--build-arg=" + k.lower() + "=" + v
 294                         for k, v in os.environ.iteritems()
 295                         if k.lower() in FILTERED_ENV_NAMES]
 296             dkr.build_image(tag, docker_dir, dockerfile,
 297                             quiet=args.quiet, user=args.user, argv=argv,
 298                             extra_files_cksum=cksum)
 299
 300             rmtree(docker_dir)
 301
 302         return 0
 303
 304 class UpdateCommand(SubCommand):
 305     """ Update a docker image with new executables. Arguments: <tag> <executable>"""
 306     name = "update"
 307     def args(self, parser):
 308         parser.add_argument("tag",
 309                             help="Image Tag")
 310         parser.add_argument("executable",
 311                             help="Executable to copy")
 312
 313     def run(self, args, argv):
 314         # Create a temporary tarball with our whole build context and
 315         # dockerfile for the update
 316         tmp = tempfile.NamedTemporaryFile(suffix="dckr.tar.gz")
 317         tmp_tar = TarFile(fileobj=tmp, mode='w')
 318
 319         # Add the executable to the tarball
 320         bn = os.path.basename(args.executable)
 321         ff = "/usr/bin/%s" % bn
 322         tmp_tar.add(args.executable, arcname=ff)
 323
 324         # Add any associated libraries
 325         libs = _get_so_libs(args.executable)
 326         if libs:
 327             for l in libs:
 328                 tmp_tar.add(os.path.realpath(l), arcname=l)
 329
 330         # Create a Docker buildfile
 331         df = StringIO()
 332         df.write("FROM %s\n" % args.tag)
 333         df.write("ADD . /\n")
 334         df.seek(0)
 335
 336         df_tar = TarInfo(name="Dockerfile")
 337         df_tar.size = len(df.buf)
 338         tmp_tar.addfile(df_tar, fileobj=df)
 339
 340         tmp_tar.close()
 341
 342         # reset the file pointers
 343         tmp.flush()
 344         tmp.seek(0)
 345
 346         # Run the build with our tarball context
 347         dkr = Docker()
 348         dkr.update_image(args.tag, tmp, quiet=args.quiet)
 349
 350         return 0
 351
 352 class CleanCommand(SubCommand):
 353     """Clean up docker instances"""
 354     name = "clean"
 355     def run(self, args, argv):
 356         Docker().clean()
 357         return 0
 358
 359 class ImagesCommand(SubCommand):
 360     """Run "docker images" command"""
 361     name = "images"
 362     def run(self, args, argv):
 363         return Docker().command("images", argv, args.quiet)
 364
 365 def main():
 366     parser = argparse.ArgumentParser(description="A Docker helper",
 367             usage="%s <subcommand> ..." % os.path.basename(sys.argv[0]))
 368     subparsers = parser.add_subparsers(title="subcommands", help=None)
 369     for cls in SubCommand.__subclasses__():
 370         cmd = cls()
 371         subp = subparsers.add_parser(cmd.name, help=cmd.__doc__)
 372         cmd.shared_args(subp)
 373         cmd.args(subp)
 374         subp.set_defaults(cmdobj=cmd)
 375     args, argv = parser.parse_known_args()
 376     return args.cmdobj.run(args, argv)
 377
 378 if __name__ == "__main__":
 379     sys.exit(main())