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