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