1# Copyright 2017 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import logging
6import os
7import sys
8
9import common
10from autotest_lib.client.bin import utils
11from autotest_lib.client.common_lib import error
12from autotest_lib.site_utils.lxc import constants
13from autotest_lib.site_utils.lxc import lxc
14from autotest_lib.site_utils.lxc import utils as lxc_utils
15from autotest_lib.site_utils.lxc.container import Container
16
17
18class BaseImage(object):
19    """A class that manages a base container.
20
21    Instantiating this class will cause it to search for a base container under
22    the given path and name.  If one is found, the class adopts it.  If not, the
23    setup() method needs to be called, to download and install a new base
24    container.
25
26    The actual base container can be obtained by calling the get() method.
27
28    Calling cleanup() will delete the base container along with all of its
29    associated snapshot clones.
30    """
31
32    def __init__(self, container_path, base_name):
33        """Creates a new BaseImage.
34
35        If a valid base container already exists on this machine, the BaseImage
36        adopts it.  Otherwise, setup needs to be called to download a base and
37        install a base container.
38
39        @param container_path: The LXC path for the base container.
40        @param base_name: The base container name.
41        """
42        self.container_path = container_path
43        self.base_name = base_name
44        try:
45            base_container = Container.create_from_existing_dir(
46                container_path, base_name)
47            base_container.refresh_status()
48            self.base_container = base_container
49        except error.ContainerError as e:
50            self.base_container = None
51            self.base_container_error = e
52
53    def setup(self, name=None, force_delete=False):
54        """Download and setup the base container.
55
56        @param name: Name of the base container, defaults to the name passed to
57                     the constructor.  If a different name is provided, that
58                     name overrides the name originally passed to the
59                     constructor.
60        @param force_delete: True to force to delete existing base container.
61                             This action will destroy all running test
62                             containers. Default is set to False.
63        """
64        if name is not None:
65            self.base_name = name
66
67        if not self.container_path:
68            raise error.ContainerError(
69                'You must set a valid directory to store containers in '
70                'global config "AUTOSERV/ container_path".')
71
72        if not os.path.exists(self.container_path):
73            os.makedirs(self.container_path)
74
75        if self.base_container and not force_delete:
76            logging.error(
77                'Base container already exists. Set force_delete to True '
78                'to force to re-stage base container. Note that this '
79                'action will destroy all running test containers')
80            # Set proper file permission. base container in moblab may have
81            # owner of not being root. Force to update the folder's owner.
82            self._set_root_owner()
83            return
84
85        # Destroy existing base container if exists.
86        if self.base_container:
87            self.cleanup()
88
89        try:
90            self._download_and_install_base_container()
91            self._set_root_owner()
92        except:
93            # Clean up if something went wrong.
94            base_path = os.path.join(self.container_path, self.base_name)
95            if lxc_utils.path_exists(base_path):
96                exc_info = sys.exc_info()
97                container = Container.create_from_existing_dir(
98                    self.container_path, self.base_name)
99                # Attempt destroy.  Log but otherwise ignore errors.
100                try:
101                    container.destroy()
102                except error.CmdError as e:
103                    logging.error(e)
104                # Raise the cached exception with original backtrace.
105                raise exc_info[0], exc_info[1], exc_info[2]
106            else:
107                raise
108        else:
109            self.base_container = Container.create_from_existing_dir(
110                self.container_path, self.base_name)
111
112    def cleanup(self):
113        """Destroys the base container.
114
115        This operation will also destroy all snapshot clones of the base
116        container.
117        """
118        # Find and delete clones first.
119        for clone in self._find_clones():
120            clone.destroy()
121        base = Container.create_from_existing_dir(self.container_path,
122                                                  self.base_name)
123        base.destroy()
124
125    def get(self):
126        """Returns the base container.
127
128        @raise ContainerError: If the base image is invalid or missing.
129        """
130        if self.base_container is None:
131            raise self.base_container_error
132        else:
133            return self.base_container
134
135    def _download_and_install_base_container(self):
136        """Downloads the base image, untars and configures it."""
137        base_path = os.path.join(self.container_path, self.base_name)
138        tar_path = os.path.join(self.container_path,
139                                '%s.tar.xz' % self.base_name)
140
141        # Force cleanup of any previously downloaded/installed base containers.
142        # This ensures a clean setup of the new base container.
143        #
144        # TODO(kenobi): Add a check to ensure that the base container doesn't
145        # get deleted while snapshot clones exist (otherwise running tests might
146        # get disrupted).
147        path_to_cleanup = [tar_path, base_path]
148        for path in path_to_cleanup:
149            if os.path.exists(path):
150                utils.run('sudo rm -rf "%s"' % path)
151        container_url = constants.CONTAINER_BASE_URL_FMT % self.base_name
152        lxc.download_extract(container_url, tar_path, self.container_path)
153        # Remove the downloaded container tar file.
154        utils.run('sudo rm "%s"' % tar_path)
155
156        # Update container config with container_path from global config.
157        config_path = os.path.join(base_path, 'config')
158        rootfs_path = os.path.join(base_path, 'rootfs')
159        utils.run(('sudo sed '
160                   '-i "s|\(lxc\.rootfs[[:space:]]*=\).*$|\\1 {rootfs}|" '
161                   '"{config}"').format(rootfs=rootfs_path,
162                                        config=config_path))
163
164    def _set_root_owner(self):
165        """Changes the container group and owner to root.
166
167        This is necessary because we currently run privileged containers.
168        """
169        # TODO(dshi): Change root to current user when test container can be
170        # unprivileged container.
171        base_path = os.path.join(self.container_path, self.base_name)
172        utils.run('sudo chown -R root "%s"' % base_path)
173        utils.run('sudo chgrp -R root "%s"' % base_path)
174
175    def _find_clones(self):
176        """Finds snapshot clones of the current base container."""
177        snapshot_file = os.path.join(self.container_path,
178                                     self.base_name,
179                                     'lxc_snapshots')
180        if not lxc_utils.path_exists(snapshot_file):
181            return
182        cmd = 'sudo cat %s' % snapshot_file
183        clone_info = [line.strip()
184                      for line in utils.run(cmd).stdout.splitlines()]
185        # lxc_snapshots contains pairs of lines (lxc_path, container_name).
186        for i in range(0, len(clone_info), 2):
187            lxc_path = clone_info[i]
188            name = clone_info[i+1]
189            yield Container.create_from_existing_dir(lxc_path, name)
190