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