1#
2# Copyright (C) 2017 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import logging
18import os
19import re
20import shutil
21import tempfile
22import zipfile
23
24from host_controller import common
25from vts.runners.host import utils
26
27
28class BuildProvider(object):
29    """The base class for build provider.
30
31    Attributes:
32        _IMAGE_FILE_EXTENSIONS: a list of strings which are common image file
33                                extensions.
34        _BASIC_IMAGE_FILE_NAMES: a list of strings which are the image names in
35                                 an artifact zip.
36        _CONFIG_FILE_EXTENSION: string, the config file extension.
37        _additional_files: a dict containing additionally fetched files that
38                           custom features may need. The key is the path
39                           relative to temporary directory and the value is the
40                           full path.
41        _configs: dict where the key is config type and value is the config file
42                  path.
43        _device_images: dict where the key is image file name and value is the
44                        path.
45        _test_suites: dict where the key is test suite type and value is the
46                      test suite package file path.
47        _host_controller_package: dict where the key is a host controller
48                                  package and the value is the path
49                                  to a package file.
50        _tmp_dirpath: string, the temp dir path created to keep artifacts.
51        _last_fetched_artifact_type: string, stores the type of the last
52                                     artifact fetched.
53    """
54    _CONFIG_FILE_EXTENSION = ".zip"
55    _IMAGE_FILE_EXTENSIONS = [".img", ".bin"]
56    _BASIC_IMAGE_FILE_NAMES = ["boot.img", "system.img", "vendor.img"]
57
58    def __init__(self):
59        self._additional_files = {}
60        self._device_images = {}
61        self._test_suites = {}
62        self._host_controller_package = {}
63        self._configs = {}
64        self._last_fetched_artifact_type = None
65        tempdir_base = os.path.join(os.getcwd(), "tmp")
66        if not os.path.exists(tempdir_base):
67            os.mkdir(tempdir_base)
68        self._tmp_dirpath = tempfile.mkdtemp(dir=tempdir_base)
69
70    def __del__(self):
71        """Deletes the temp dir if still set."""
72        if self._tmp_dirpath:
73            shutil.rmtree(self._tmp_dirpath)
74            self._tmp_dirpath = None
75
76    @property
77    def tmp_dirpath(self):
78        return self._tmp_dirpath
79
80    def CreateNewTmpDir(self):
81        return tempfile.mkdtemp(dir=self._tmp_dirpath)
82
83    def SetDeviceImage(self, name, path):
84        """Sets device image `path` for the specified `name`."""
85        self._device_images[name] = path
86        self._last_fetched_artifact_type = common._ARTIFACT_TYPE_DEVICE
87
88    def _IsFullDeviceImage(self, namelist):
89        """Returns true if given namelist list has all common device images."""
90        for image_file in self._BASIC_IMAGE_FILE_NAMES:
91            if image_file not in namelist:
92                return False
93        return True
94
95    def _IsImageFile(self, file_path):
96        """Returns whether a file is an image.
97
98        Args:
99            file_path: string, the file path.
100
101        Returns:
102            boolean, whether the file is an image.
103        """
104        return any(file_path.endswith(ext)
105                   for ext in self._IMAGE_FILE_EXTENSIONS)
106
107    def SetDeviceImageZip(self, path, full_device_images=False):
108        """Sets device image(s) using files in a given zip file.
109
110        It extracts image files inside the given zip file and selects
111        known Android image files.
112
113        Args:
114            path: string, the path to a zip file.
115        """
116        dest_path = path + ".dir"
117        fetch_type = None
118        with zipfile.ZipFile(path, 'r') as zip_ref:
119            if full_device_images or self._IsFullDeviceImage(zip_ref.namelist()):
120                self.SetDeviceImage(common.FULL_ZIPFILE, path)
121                dir_key = common.FULL_ZIPFILE_DIR
122                fetch_type = common._ARTIFACT_TYPE_DEVICE
123            else:
124                self.SetDeviceImage("gsi-zipfile", path)
125                dir_key = common.GSI_ZIPFILE_DIR  # "gsi-zipfile-dir"
126                fetch_type = common._ARTIFACT_TYPE_GSI
127            if os.path.exists(dest_path):
128                shutil.rmtree(dest_path)
129                logging.info("%s %s deleted", dir_key, dest_path)
130            zip_ref.extractall(dest_path)
131            self.SetFetchedDirectory(dest_path)
132            self.SetDeviceImage(dir_key, dest_path)
133
134        self._last_fetched_artifact_type = fetch_type
135
136    def GetDeviceImage(self, name=None):
137        """Returns device image info."""
138        if name is None:
139            return self._device_images
140        return self._device_images[name]
141
142    def RemoveDeviceImage(self, name):
143        """Removes certain device image info.
144
145        Args:
146            name: string, the name of the device image file
147                  that needs to be removed.
148        """
149        if name in self._device_images:
150            self._device_images.pop(name)
151
152    def SetTestSuitePackage(self, test_suite, path):
153        """Sets test suite package `path` for the specified `type`.
154
155        Args:
156            test_suite: string, test suite type such as 'vts' or 'cts', etc.
157            path: string, the path of a file. if a file is a zip file,
158                  it's unziped and its main binary is set.
159        """
160        if re.match("[vcgs]ts", test_suite):
161            suite_name = "android-%s" % test_suite
162            tradefed_name = "%s-tradefed" % test_suite
163            dest_path = os.path.join(self.tmp_dirpath, suite_name)
164            if os.path.exists(dest_path):
165                shutil.rmtree(dest_path)
166                logging.info("test suite %s deleted", dest_path)
167            with zipfile.ZipFile(path, 'r') as zip_ref:
168                zip_ref.extractall(dest_path)
169                bin_path = os.path.join(dest_path, suite_name,
170                                        "tools", tradefed_name)
171                os.chmod(bin_path, 0766)
172                path = bin_path
173        else:
174            logging.info("unsupported zip file %s", path)
175        self._test_suites[test_suite] = path
176        self._last_fetched_artifact_type = common._ARTIFACT_TYPE_TEST_SUITE
177
178    def GetTestSuitePackage(self, type=None):
179        """Returns test suite package info."""
180        if type is None:
181            return self._test_suites
182        return self._test_suites[type]
183
184    def SetHostControllerPackage(self, package_type, path):
185        """Sets host controller package `path` for the specified `type`.
186
187        Args:
188            package_type: string, host controller type such as 'vtslab'.
189            path: string, the path of a package file.
190        """
191        self._host_controller_package[package_type] = path
192        self._last_fetched_artifact_type = common._ARTIFACT_TYPE_INFRA
193
194    def GetHostControllerPackage(self, package_type=None):
195        """Returns host controller package info.
196
197        Args:
198            package_type: string, key value to self._host_controller_package
199                          dict.
200
201        Returns:
202            the whole dict if package_type is None, otherwise a string which is
203            the path to the fetched host controller package.
204        """
205        if package_type is None:
206            return self._host_controller_package
207        return self._host_controller_package[package_type]
208
209    def SetConfigPackage(self, config_type, path):
210        """Sets test suite package `path` for the specified `type`.
211
212        All valid config files have .zip extension.
213
214        Args:
215            config_type: string, config type such as 'prod' or 'test'.
216            path: string, the path of a config file.
217        """
218        if path.endswith(self._CONFIG_FILE_EXTENSION):
219            dest_path = os.path.join(
220                self.tmp_dirpath, os.path.basename(path) + ".dir")
221            with zipfile.ZipFile(path, 'r') as zip_ref:
222                zip_ref.extractall(dest_path)
223                path = dest_path
224        else:
225            logging.info("unsupported config package file %s", path)
226        self._configs[config_type] = path
227        self._last_fetched_artifact_type = common._ARTIFACT_TYPE_INFRA
228
229    def GetConfigPackage(self, config_type=None):
230        """Returns config package info."""
231        if config_type is None:
232            return self._configs
233        return self._configs[config_type]
234
235    def SetAdditionalFile(self, rel_path, full_path):
236        """Sets the key and value of additionally fetched files.
237
238        Args:
239            rel_path: the file path relative to temporary directory.
240            abs_path: the file path that this process can access.
241        """
242        self._additional_files[rel_path] = full_path
243        self._last_fetched_artifact_type = common._ARTIFACT_TYPE_INFRA
244
245    def GetAdditionalFile(self, rel_path=None):
246        """Returns the paths to fetched files."""
247        if rel_path is None:
248            return self._additional_files
249        return self._additional_files[rel_path]
250
251    def SetFetchedDirectory(self,
252                            dir_path,
253                            root_path=None,
254                            full_device_images=False):
255        """Adds every file in a directory to one of the dictionaries.
256
257        This method follows symlink to file, but skips symlink to directory.
258
259        Args:
260            dir_path: string, the directory to find files in.
261            root_path: string, the temporary directory that dir_path is in.
262                       The default value is dir_path.
263        """
264        for dir_name, file_name in utils.iterate_files(dir_path):
265            full_path = os.path.join(dir_name, file_name)
266            self.SetFetchedFile(full_path, (root_path
267                                            if root_path else dir_path),
268                                full_device_images)
269
270    def SetFetchedFile(self,
271                       file_path,
272                       root_dir=None,
273                       full_device_images=False,
274                       set_suite_as=None):
275        """Adds a file to one of the dictionaries.
276
277        Args:
278            file_path: string, the path to the file.
279            root_dir: string, the temporary directory that file_path is in.
280                      The default value is file_path if file_path is a
281                      directory. Otherwise, the default value is file_path's
282                      parent directory.
283            set_suite_as: string, the test suite name to use for the given
284                          artifact. Used when the file name does not follow
285                          the standard "android-*ts.zip" file name pattern.
286        """
287        file_name = os.path.basename(file_path)
288        if os.path.isdir(file_path):
289            self.SetFetchedDirectory(file_path, root_dir, full_device_images)
290        elif self._IsImageFile(file_path):
291            self.SetDeviceImage(file_name, file_path)
292        elif re.match("android-[vcgs]ts.zip", file_name):
293            test_suite = (file_name.split("-")[-1]).split(".")[0]
294            self.SetTestSuitePackage(test_suite, file_path)
295        elif file_name == "android-vtslab.zip":
296            self.SetHostControllerPackage("vtslab", file_path)
297        elif file_name.startswith("vti-global-config"):
298            self.SetConfigPackage(
299                "prod" if "prod" in file_name else "test", file_path)
300        elif set_suite_as:
301            self.SetTestSuitePackage(set_suite_as, file_path)
302        elif file_path.endswith(".zip"):
303            self.SetDeviceImageZip(file_path, full_device_images)
304        else:
305            rel_path = (os.path.relpath(file_path, root_dir) if root_dir else
306                        os.path.basename(file_path))
307            self.SetAdditionalFile(rel_path, file_path)
308
309    def PrintDeviceImageInfo(self):
310        """Prints device image info."""
311        logging.info(self.GetDeviceImage())
312
313    def PrintGetTestSuitePackageInfo(self):
314        """Prints test suite package info."""
315        logging.info(self.GetTestSuitePackage())
316
317    def GetFetchedArtifactType(self):
318        """Gets the most recently fetched artifact type.
319
320        Returns:
321            string, type of the artifact.
322        """
323        return self._last_fetched_artifact_type
324