1#
2# Copyright (C) 2018 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 host_controller.command_processor import base_command_processor
26
27# Name of android-info.txt file which contains prerequisite data for the img.zip
28_ANDROID_INFO_TXT_FILENAME = "android-info.txt"
29
30
31class CommandRepack(base_command_processor.BaseCommandProcessor):
32    """Command processor for repack command."""
33
34    command = "repack"
35    command_detail = ("Repackage the whole device image files, including GSI "
36                      "images if exist.")
37
38    # @Override
39    def SetUp(self):
40        """Initializes the parser for repack command."""
41        self.arg_parser.add_argument(
42            "--dest",
43            default="gs://vts-release/img_package",
44            help="Google Cloud Storage base URL to which the file is uploaded."
45        )
46        self.arg_parser.add_argument(
47            "--additional_files",
48            nargs="*",
49            default=[],
50            help="Additional files that need to be added to the zip file.")
51
52    # @Override
53    def Run(self, arg_line):
54        """Runs an repack command."""
55        args = self.arg_parser.ParseLine(arg_line)
56        try:
57            device_zipfile_path = self.console.device_image_info[
58                common.FULL_ZIPFILE]
59        except KeyError as e:
60            logging.exception(e)
61            logging.error(
62                "please execute this command after fetching at least one "
63                "device img set.")
64            return False
65
66        with zipfile.ZipFile(device_zipfile_path, 'r') as zip_ref:
67            zip_ref.extract(_ANDROID_INFO_TXT_FILENAME,
68                            common.FULL_ZIPFILE_DIR)
69            self.console.tools_info[_ANDROID_INFO_TXT_FILENAME] = os.path.join(
70                common.FULL_ZIPFILE_DIR, _ANDROID_INFO_TXT_FILENAME)
71
72        tempdir_base = os.path.join(os.getcwd(), "tmp")
73        tmpdir_rezip = tempfile.mkdtemp(dir=tempdir_base)
74
75        dest_url_base, new_zipfile_name = os.path.split(
76            self.GetDestURL(args.dest))
77
78        new_zipfile_path = os.path.join(tmpdir_rezip, new_zipfile_name)
79        with zipfile.ZipFile(
80                new_zipfile_path, "w", allowZip64=True) as zip_ref:
81            for img_path in self.console.device_image_info:
82                if img_path not in (common.FULL_ZIPFILE,
83                                    common.FULL_ZIPFILE_DIR,
84                                    common.GSI_ZIPFILE,
85                                    common.GSI_ZIPFILE_DIR):
86                    logging.info("Adding %s into the zip archive.", img_path)
87                    zip_ref.write(
88                        self.console.device_image_info[img_path],
89                        img_path,
90                        compress_type=zipfile.ZIP_DEFLATED)
91            if args.additional_files:
92                additional_file_list = self.ReplaceVars(args.additional_files)
93                for file_path in additional_file_list:
94                    file_name = os.path.basename(file_path)
95                    logging.info(
96                        "Adding additional file %s into the zip archive.",
97                        file_name)
98                    zip_ref.write(
99                        file_path,
100                        os.path.join(common._ADDITIONAL_FILES_DIR, file_name),
101                        compress_type=zipfile.ZIP_DEFLATED)
102            zip_ref.write(
103                self.console.tools_info[_ANDROID_INFO_TXT_FILENAME],
104                _ANDROID_INFO_TXT_FILENAME,
105                compress_type=zipfile.ZIP_DEFLATED)
106
107        logging.info("Repackaged image set: %s", new_zipfile_path)
108        logging.info("Uploading %s to %s.", new_zipfile_name, dest_url_base)
109
110        self.console.onecmd("upload --src=%s --dest=%s/%s" %
111                            (new_zipfile_path, dest_url_base,
112                             new_zipfile_name))
113
114        shutil.rmtree(tmpdir_rezip, ignore_errors=True)
115
116    def GetDestURL(self, dest_base_url):
117        """Generates the destination URL to GCS bucket based on dest_base_url.
118
119        Args:
120            dest_base_url: string, URL to a GCS bucket.
121
122        Returns:
123            A string, device/gsi img sets branch/target info and the final
124            .zip file name appended to the dest_base_url.
125        """
126        device_branch = re.sub(
127            "git_", "", self.console.detailed_fetch_info["device"]["branch"])
128        device_target = self.console.detailed_fetch_info["device"]["target"]
129        device_build_id = self.console.detailed_fetch_info["device"][
130            "build_id"]
131        new_zipfile_name = ("%s_%s_%s.zip" % (device_branch, device_target,
132                                              device_build_id))
133        dest_url_base = os.path.join(dest_base_url, device_branch,
134                                     device_target)
135
136        if common.GSI_ZIPFILE in self.console.device_image_info:
137            gsi_branch = re.sub(
138                "git_", "", self.console.detailed_fetch_info["gsi"]["branch"])
139            gsi_target = self.console.detailed_fetch_info["gsi"]["target"]
140            gsi_build_id = self.console.detailed_fetch_info["gsi"]["build_id"]
141            new_zipfile_name = new_zipfile_name[:-4] + ("_%s_%s_%s.zip" % (
142                gsi_branch, gsi_target, gsi_build_id))
143            dest_url_base = os.path.join(dest_url_base, gsi_branch, gsi_target)
144
145        ret = os.path.join(dest_url_base, new_zipfile_name)
146        self.console.repack_dest_path = ret
147        return ret
148