1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# Copyright (c) 2021 Huawei Device Co., Ltd.
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.
15import ast
16import os
17import shutil
18import zipfile
19
20from test.fake_data import RSA_PRIVATE_KEY_DATA
21
22UPDATER_BINARY_DATA = b"updater binary data"
23BOARD_LIST_DATA = b"""HI3516
24HI3518
25HI3559"""
26VERSION_MBN_DATA = b"""Hi3516DV300-eng 10 QP1A.190711.020
27Hi3516DV300-eng 10 QP1A.190711.021"""
28
29UPDATER_SPECIFIED_CONFIG_XML_DATA = """<?xml version="1.0"?>
30<package>
31    <head name="Component header information">
32        <info fileVersion="01" prdID="123456"
33        softVersion="Hi3516DV300-eng 10 QP1A.190711.0VERSION_MARK"
34        date="2021-01-23" time="12:30">head info</info>
35    </head>
36    <group name = "Component information">
37        SYSTEM_MARK
38        COMPONENT_MARK
39    </group>
40</package>"""
41UPDATER_PARTITIONS_XML_DATA = b"""<?xml version="1.0" encoding="GB2312" ?>
42<Partition_Info>
43<Part PartitionName="boot" FlashType="emmc" FileSystem="none"
44    Start="0" Length="1M"/>
45<Part PartitionName="kernel" FlashType="emmc" FileSystem="none"
46    Start="1M" Length="15M"/>
47<Part PartitionName="updater" FlashType="emmc" FileSystem="ext3/4"
48    Start="16M" Length="20M"/>
49<Part PartitionName="misc" FlashType="emmc" FileSystem="none"
50    Start="36M" Length="1M"/>
51<Part PartitionName="system" FlashType="emmc" FileSystem="ext3/4"
52    Start="37M" Length="3300M"/>
53<Part PartitionName="vendor" FlashType="emmc" FileSystem="ext3/4"
54    Start="3337M" Length="263M"/>
55<Part PartitionName="userdata" FlashType="emmc" FileSystem="ext3/4"
56    Start="3600M" Length="1464M"/>
57</Partition_Info>"""
58
59SYSTEM_COMPONENT_STR = """<component compAddr="system" compId="12"
60        resType="05" compType="0" compVer="0o00">./system.img</component>"""
61COMPONENT_STR_INC = """<component compAddr="vendor" compId="13"
62        resType="05" compType="1" compVer="0o00">./vendor.img</component>"""
63COMPONENT_STR_FULL = """<component compAddr="vendor" compId="13"
64        resType="05" compType="0" compVer="0o00">./vendor.img</component>"""
65COMPONENT_STR_USERDATA = """<component compAddr="userdata" compId="14"
66        resType="05" compType="0" compVer="0o00">./userdata.img</component>"""
67SOURCE_VERSION_STR = "20"
68TARGET_VERSION_STR = "21"
69
70
71def create_input_package(
72        package_name, package_type="target", is_su=False,
73        is_updater_partitions=False, has_userdata=False,
74        is_inc_full_none=False, is_miss_image=False,
75        is_miss_version_list=False, is_miss_updater_binary=False):
76    """
77    Create input package.
78    :param package_name: package name
79    :param package_type: package type, source or target
80    :param is_su: Is it an updater upgrade package
81    :param is_updater_partitions: Is it an updater partitions upgrade package
82    :param has_userdata: Configuring UserData in XML
83    :param is_inc_full_none: Both full and incremental list lengths are zero
84    :param is_miss_image: Is image missing
85    :param is_miss_version_list: Is VERSION.list missing
86    :param is_miss_updater_binary: Is updater_binary missing
87    :return:
88    """
89    # Create a folder for the input package.
90    package_name_path = "./%s" % package_name
91    if not os.path.exists(package_name_path):
92        os.mkdir(package_name_path)
93
94    if not is_miss_image:
95        create_file(os.path.join(package_name_path, "system.img"), get_target_vendor_data())
96
97    # Judge the type of input package and generate the corresponding vendor.img
98    if package_type == "target":
99        vendor_content = get_target_vendor_data()
100    elif package_type == "source":
101        vendor_content = get_source_vendor_data()
102    else:
103        print("Unknown package type!")
104        raise RuntimeError
105    if not is_miss_image:
106        create_file(os.path.join(package_name_path, "vendor.img"), vendor_content)
107    if not is_miss_updater_binary:
108        create_file(os.path.join(package_name_path, "updater_binary"), UPDATER_BINARY_DATA)
109    # updater upgrade package
110    if is_su:
111        create_file(os.path.join(package_name_path, "uImage"), get_target_vendor_data())
112        create_file(os.path.join(package_name_path, "updater.img"), get_target_vendor_data())
113        create_file(os.path.join(package_name_path, "updater_b.img"), get_target_vendor_data())
114        create_file(os.path.join(package_name_path, "updater_uImage"), get_target_vendor_data())
115    # Create updater_config dir.
116    updater_config_path = "./%s/updater_config" % package_name
117    if not os.path.exists(updater_config_path):
118        os.mkdir(updater_config_path)
119    create_file(os.path.join(updater_config_path, "BOARD.list"), BOARD_LIST_DATA)
120    if not is_miss_version_list:
121        create_file(os.path.join(updater_config_path, "VERSION.mbn"), VERSION_MBN_DATA)
122    # Judge the type of input package and
123    xml_content = \
124        create_updater_specified_config_file(has_userdata, is_updater_partitions, package_type)
125
126    if is_inc_full_none:
127        xml_content = xml_content.replace("SYSTEM_MARK", "")
128        xml_content = xml_content.replace("COMPONENT_MARK", "")
129    else:
130        xml_content = xml_content.replace("SYSTEM_MARK", SYSTEM_COMPONENT_STR)
131        xml_content = xml_content.replace("COMPONENT_MARK", COMPONENT_STR_FULL)
132
133    create_file(os.path.join(updater_config_path, "updater_specified_config.xml"), xml_content)
134    # Create partition_file.xml.
135    if is_updater_partitions:
136        create_file("./partition_file.xml", UPDATER_PARTITIONS_XML_DATA)
137    # Create rsa_private_key2048.pem.
138    create_file("./rsa_private_key2048.pem", RSA_PRIVATE_KEY_DATA)
139    # Create zip package.
140    with zipfile.ZipFile('./%s.zip' % package_name, 'w', zipfile.ZIP_DEFLATED) as package_zip:
141        package_zip.write(package_name_path)
142        for home, dirs, files in os.walk(package_name_path):
143            for each_file_name in files:
144                package_zip.write(os.path.join(home, each_file_name))
145            for each_dir_name in dirs:
146                package_zip.write(os.path.join(home, each_dir_name))
147
148
149def create_updater_specified_config_file(
150        has_userdata, is_updater_partitions, package_type):
151    """
152    generate the corresponding updater_specified_config.xml
153    :param has_userdata: has userdata
154    :param is_updater_partitions: is updater partitions
155    :param package_type: package type
156    :return:
157    """
158    if package_type == "target":
159        xml_content = UPDATER_SPECIFIED_CONFIG_XML_DATA.replace(
160            "VERSION_MARK", TARGET_VERSION_STR)
161        xml_content = xml_content.replace(
162            "COMPONENT_MARK", COMPONENT_STR_INC)
163    elif package_type == "source":
164        xml_content = UPDATER_SPECIFIED_CONFIG_XML_DATA.replace(
165            "VERSION_MARK", SOURCE_VERSION_STR)
166        if is_updater_partitions:
167            xml_content = xml_content.replace(
168                "COMPONENT_MARK", COMPONENT_STR_FULL)
169        elif has_userdata:
170            xml_content = xml_content.replace(
171                "COMPONENT_MARK", COMPONENT_STR_USERDATA)
172    else:
173        print("Unknown package type!")
174        raise RuntimeError
175    return xml_content
176
177
178def create_file(file_path, file_data):
179    """
180    Create file
181    :param file_path: file path
182    :param file_data: file data
183    :return:
184    """
185    with open(file_path, "wb") as w_f:
186        w_f.write(file_data)
187
188
189def clear_package(package_name):
190    """
191    Clean up the constructed input package and files
192    :param package_name: constructed input package name
193    :return:
194    """
195    if os.path.exists("./%s" % package_name):
196        shutil.rmtree("./%s" % package_name)
197    if os.path.exists("./%s.zip" % package_name):
198        os.remove("./%s.zip" % package_name)
199    if os.path.exists("./partition_file.xml"):
200        os.remove("./partition_file.xml")
201    if os.path.exists("./rsa_private_key2048.pem"):
202        os.remove("./rsa_private_key2048.pem")
203
204
205def get_source_vendor_data():
206    """
207    Get source vendor image file data
208    :return:
209    """
210    with open(r"./source_vendor_data", "r") as r_f:
211        content = r_f.read()
212        return ast.literal_eval(content)
213
214
215def get_target_vendor_data():
216    """
217    Get target vendor image file data
218    :return:
219    """
220    with open(r"./target_vendor_data", "r") as r_f:
221        content = r_f.read()
222        return ast.literal_eval(content)
223