1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3
4# Copyright (c) 2023 Huawei Device Co., Ltd.
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16"""
17Description : Generate the update.bin hash check data
18"""
19import os
20import struct
21import hashlib
22import enum
23from log_exception import UPDATE_LOGGER
24
25# hash data sample
26# hash info module:
27# hash info:1 32 3 4194304
28# hash value module:
29# /version_list                   (32bytes)
30# 1 176
31# 1 176 bf10259a1fc1b2f780a49ce6XXXXXXXX
32# hash sign module:
33# hash sign:45ef8ec12e56e3b82c9a05XXXXXX
34
35HashType = enum.Enum('HashType', ('SHA256', 'SHA384', 'SHA512'))
36HashAlgo = {HashType.SHA256 : hashlib.sha256,
37            HashType.SHA384 : hashlib.sha384,
38            HashType.SHA512 : hashlib.sha512}
39HASH_TYPE_SIZE = 2
40HASH_LENGTH_SIZE = 4
41HASH_TLV_SIZE = HASH_TYPE_SIZE + HASH_LENGTH_SIZE
42UPGRADE_HASHINFO_SIZE = 10
43HASH_DATA_HEADER_SIZE = 38
44HASH_DATA_ADDR_SIZE = 8
45COMPONENT_NAME_SIZE = 32
46# hash block size
47HASH_BLOCK_SIZE = 4 * 1024 * 1024
48
49"""
50Format
51H: unsigned short
52I: unsigned int
53B: unsigned char
54s: char[]
55"""
56HASH_TLV_FMT = "<HI"
57HASH_INFO_FMT = "<3HI"
58HASH_DATA_HEADER_FMT = "<32sHI"
59HASH_DATA_ADDR_FMT = "<2I"
60
61
62class CreateHash(object):
63    """
64    Create the component hash data
65    """
66
67    def __init__(self, hash_type, count):
68        self.hashinfo_tlv_type = 0x06
69        self.hashdata_tlv_type = 0x07
70        self.sign_tlv_type = 0x08
71        self.hash_type = hash_type
72        self.hash_digest_size = HashAlgo[hash_type]().digest_size
73        self.component_num = count
74        self.block_size = HASH_BLOCK_SIZE
75        self.hashinfo_value = bytes()
76        self.hashdata = bytes()
77        self.signdata = bytes()
78        self.hashdata_list = []
79
80    def write_hashinfo(self):
81        try:
82            hashinfo_tlv = struct.pack(HASH_TLV_FMT, self.hashinfo_tlv_type, UPGRADE_HASHINFO_SIZE)
83            hashinfo_header = struct.pack(HASH_INFO_FMT, self.hash_type.value, self.hash_digest_size,
84                self.component_num, self.block_size)
85        except struct.error:
86            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
87            return False
88
89        # write hashinfo
90        self.hashinfo_value = hashinfo_tlv + hashinfo_header
91        return True
92
93    def write_hashdata(self):
94        try:
95            hashdata_len = len(self.hashdata)
96            hashdata_tlv = struct.pack(HASH_TLV_FMT, self.hashdata_tlv_type, hashdata_len)
97        except struct.error:
98            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
99            return False
100
101        UPDATE_LOGGER.print_log("Write hashdata hash len %d" % hashdata_len)
102        # write hashdata
103        self.hashdata = hashdata_tlv + self.hashdata
104        UPDATE_LOGGER.print_log("Write hashdata hash tlv complete")
105        return True
106
107    def write_signdata(self, signdata):
108        try:
109            signdata_len = len(signdata)
110            signdata_tlv = struct.pack(HASH_TLV_FMT, self.sign_tlv_type, signdata_len)
111        except struct.error:
112            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
113            return False
114
115        # write signdata
116        self.signdata = signdata_tlv + signdata
117        UPDATE_LOGGER.print_log("Write hashdata sign tlv complete")
118        return True
119
120    def calculate_hash_data(self, data):
121        hash_algo = HashAlgo[self.hash_type]()
122        hash_algo.update(data)
123        return hash_algo.digest()
124
125    def write_component_hash_data(self, component):
126        UPDATE_LOGGER.print_log("calc component hash")
127        try:
128            with open(component.file_path, "rb") as component_file:
129                component_len = os.path.getsize(component.file_path)
130                block_num = component_len // HASH_BLOCK_SIZE
131                component_name = component.component_addr.decode().ljust(COMPONENT_NAME_SIZE, "\0")
132                UPDATE_LOGGER.print_log(
133                    "calc component hash component name:%s %d" % (component_name, len(component_name)))
134                total_block = block_num + 1 if component_len % HASH_BLOCK_SIZE > 0 else block_num
135                self.hashdata += struct.pack(HASH_DATA_HEADER_FMT, component_name.encode(),
136                    total_block, component_len)
137                UPDATE_LOGGER.print_log("calc component hash  block_num:%d" % total_block)
138                write_len = 0
139                for i in range(0, block_num):
140                    component_file.seek(write_len)
141                    component_data = component_file.read(HASH_BLOCK_SIZE)
142                    write_len += HASH_BLOCK_SIZE
143                    self.hashdata += struct.pack(HASH_DATA_ADDR_FMT, (i * HASH_BLOCK_SIZE if i != 0 else 0),
144                        write_len - 1) + self.calculate_hash_data(component_data)
145                if component_len - write_len > 0 :
146                    component_file.seek(write_len)
147                    component_data = component_file.read(component_len - write_len)
148                    self.hashdata += struct.pack(HASH_DATA_ADDR_FMT, (write_len if write_len != 0 else 0),
149                        component_len - 1) + self.calculate_hash_data(component_data)
150        except (struct.error, IOError):
151            return False
152        UPDATE_LOGGER.print_log("calc component hash complete  ComponentSize:%d" % component_len)
153        return True
154
155    def parse_hashinfo(self, data):
156        # parse hashinfo
157        hash_type_value = 0
158        try:
159            hash_type_value, self.hash_digest_size, self.component_num, self.block_size = \
160                struct.unpack(HASH_INFO_FMT, data[:UPGRADE_HASHINFO_SIZE])
161            self.hash_type = HashType(hash_type_value)
162        except struct.error:
163            return False
164
165        UPDATE_LOGGER.print_log("parese hashinfo complete, %d %d %d %d" % (hash_type_value,
166            self.hash_digest_size, self.component_num, self.block_size))
167        return True
168
169    def parse_hashdata(self, data):
170        offset = 0
171        try:
172            for i in range(0, self.component_num):
173                img_name, hash_num, img_size = struct.unpack(HASH_DATA_HEADER_FMT,
174                    data[offset: HASH_DATA_HEADER_SIZE + offset])
175                UPDATE_LOGGER.print_log("parese hashinfo complete, %s %d %d" % (img_name,
176                    hash_num, img_size))
177                offset += HASH_DATA_HEADER_SIZE
178                self.hashdata_list.append((img_name.decode(), hash_num, img_size))
179                for j in range(0, hash_num):
180                    hash_data_star, hash_data_end = struct.unpack(HASH_DATA_ADDR_FMT,
181                        data[offset: HASH_DATA_ADDR_SIZE + offset])
182                    hash_data = data[HASH_DATA_ADDR_SIZE + offset:HASH_DATA_ADDR_SIZE + self.hash_digest_size + offset]
183                    offset += (HASH_DATA_ADDR_SIZE + self.hash_digest_size)
184                    self.hashdata_list.append((hash_data_star, hash_data_end, hash_data))
185        except struct.error:
186            return False
187
188        UPDATE_LOGGER.print_log("parese hashdata complete")
189        return True
190
191    def parse_signdata(self, data):
192        # parse signdata
193        self.signdata = data
194        UPDATE_LOGGER.print_log("parese hashdata sign complete")
195        return True
196
197    def parse_print_hashdata(self, save_path):
198        hash_check_fd = os.open(os.path.join(save_path + "hash_check_file_parse"), os.O_RDWR | os.O_CREAT, 0o755)
199        with os.fdopen(hash_check_fd, "wb+") as hash_check_file_p:
200            hash_check_file_p.write(("hash info:").encode())
201            hash_check_file_p.write(("%s %s %s %s\n" % (
202                HashType(self.hash_type.value).name, str(self.hash_digest_size),
203                str(self.component_num), str(self.block_size))).encode())
204
205            offset = 0
206            for i in range(0, self.component_num):
207                hash_check_file_p.write(("%s\n" % (self.hashdata_list[offset][0])).encode())
208                hash_check_file_p.write(("%s %s\n" % (
209                    str(self.hashdata_list[offset][1]), str(self.hashdata_list[offset][2]))).encode())
210                for j in range(0, self.hashdata_list[offset][1]):
211                    index = offset + 1
212                    hashdata_hexstr = "".join("%02x" % b for b in self.hashdata_list[j + index][2])
213                    hash_check_file_p.write(("%s" % (
214                        str(self.hashdata_list[j + index][0]), str(self.hashdata_list[j + index][1]),
215                        hashdata_hexstr)).encode())
216
217                offset += (1 + self.hashdata_list[offset][1])
218
219            signdata_hexstr = "".join("%02x" % b for b in self.signdata)
220            hash_check_file_p.write(("hash sign:").encode())
221            hash_check_file_p.write(signdata_hexstr.encode())
222