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