create_chunk.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2024 Hunan OpenValley Digital Industry Development Co., Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Description : Generate the update.bin chunk info data
"""
import os
import struct
import hashlib
import enum
import tempfile
import zipfile
from log_exception import UPDATE_LOGGER
from utils import OptionsManager
from utils import ZIP_EVENT

CHUNK_LIST_COUNT_SIZE = 4
CHUNK_LIST_SIZE = 8
CHUNK_RESERVE = 244
UPGRADE_CHUNKINFO_SIZE = 256

CHUNK_TLV_FMT = "<HI"
CHUNK_INFO_FMT = "<HI"
CHUNK_DATA_HEADER_FMT = "<HI"
CHUNK_SIGN_HEADER_FMT = '<HI'
CHUNK_DATA_PARTITION_FMT = '<HH'
CHUNK_DATA_CMD_FMT = '<HH'
CHUNK_DATA_DATA_FMT = '<HI'
CHUNK_HASH_HEADER_FMT = '<3H'
CHUNK_HASH_PARTITION_FMT = '<HH'
CHUNK_HASH_VALUE_FMT = '<HH'
CHUNK_LARGE_PARTITION_FMT = '<HH'
CHUNK_SIGN_PARTITON_FMT = "<HH"
CHUNK_SIGN_VALUE_FMT = '<HI'


OPTIONS_MANAGER = OptionsManager()


def get_file_sha256(update_package):
    """
    Get the SHA256 value of the package file
    """
    sha256obj = hashlib.sha256()
    maxbuf = 8192
    with open(update_package, 'rb') as package_file:
        while True:
            buf = package_file.read(maxbuf)
            if not buf:
                break
            sha256obj.update(buf)
    hash_value_hex = sha256obj.hexdigest()
    hash_value = sha256obj.digest()
    return str(hash_value_hex).upper()


def get_chunk_sha256(chunk):
    """
    Get the SHA256 value of the chunk
    """
    sha256obj = hashlib.sha256()
    sha256obj.update(chunk)
    hash_value_hex = sha256obj.hexdigest()
    hash_value = sha256obj.digest()
    return str(hash_value_hex).upper()


class CreateChunk(object):
    """
    Create the image chunk data
    """
    def __init__(self, num, count):
        self.chunkinfo_tlv_type = 0x10
        self.chunkdata_tlv_type = 0x11
        self.chunkdata_partition_tlv_type = 0x12
        self.chunkdata_cmd_tlv_type = 0x13
        self.chunkdata_value_tlv_type = 0x14
        self.chunkhash_tlv_type = 0x15
        self.chunkhash_info_tlv_type = 0x16
        self.chunkhash_partition_tlv_type = 0x17
        self.chunkhash_value_tlve_type = 0x18
        self.chunkimage_large_tlv_type = 0x19
        self.chunksign_value_tlv_type = 0x1a
        
        self.chunk_list_count = count
        self.chunk_list_num = num
        self.chunk_list_length = 0
        self.chunkinfo_value = bytes()
        self.pkg_chunk = bytes()
        self.hash_chunk = bytes()
        self.sign_chunk = bytes()
        self.chunkdata_value = bytes()
        self.chunk_list = []
        self.hash_list = []
        self.sign_list = []
        self.all_image_hash_data = []
        self.write_chunk_hashdata = bytes()
        self.signdata = bytes()

        
    def write_chunkinfo(self, package_file, startoffset):
        UPDATE_LOGGER.print_log("write_chunkinfo Startoffset %s!"\
            % startoffset)
        try:
            chunkinfo_tlv = struct.pack(CHUNK_TLV_FMT, self.chunkinfo_tlv_type, UPGRADE_CHUNKINFO_SIZE)
            chunkinfo_header = struct.pack(CHUNK_INFO_FMT, self.chunk_list_count, self.chunk_list_num)
        except struct.error:
            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError
        
        # write chunkinfo
        self.chunkinfo_value = chunkinfo_tlv + chunkinfo_header
        package_file.seek(startoffset)
        package_file.write(self.chunkinfo_value)
        return True

    def write_chunklist(self, image, package_file, startoffset):
        UPDATE_LOGGER.print_log("write pkg chunklist StartOffset:%s"\
            % startoffset)
        try:
            chunk_count = 0
            patch_index = 0
            new_index = 0 
            for chunk in OPTIONS_MANAGER.image_transfer_dict_contents[image].splitlines()[4:]:
                chunk_count += 1
                partiton_info = image
                cmd_info = chunk
                
                # Step 1: Pack partition name        
                partition_tlv = struct.pack(CHUNK_DATA_PARTITION_FMT, self.chunkdata_partition_tlv_type, len(partiton_info)) + partiton_info.encode('utf-8')
                print(f"Packed partition TLV{chunk_count}: {partition_tlv}")
                # Write to binary file
                package_file.seek(startoffset)
                package_file.write(partition_tlv)
                startoffset += len(partition_tlv)
            
                # Step 2: Pack command info 
                cmd_len = len(cmd_info)
                cmd_tlv = struct.pack(CHUNK_DATA_CMD_FMT, self.chunkdata_cmd_tlv_type, cmd_len) + cmd_info.encode('utf-8')
                print(f"Packed command TLV{chunk_count}: {cmd_tlv}")
                print(f'length:{cmd_len}')
                # Write to binary file
                package_file.seek(startoffset)
                package_file.write(cmd_tlv)
                startoffset += len(cmd_tlv)
                
                # Step 3: Pack patch dependency data
                data_tlv, patch_index, new_index = self.pack_dependency_data(image, cmd_info, chunk, patch_index, new_index)
                # Write to bin file 
                data_len = len(data_tlv)
                package_file.seek(startoffset)
                package_file.write(data_tlv)
                startoffset += data_len
            
        except (struct.error) as e:
            UPDATE_LOGGER.print_log(f"Unexpected error: {e}", log_type=UPDATE_LOGGER.ERROR_LOG)
            UPDATE_LOGGER.print_log("write chunk error!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError
        except Exception as e:
            UPDATE_LOGGER.print_log(f"An unexpected error occurred: {e}", log_type=UPDATE_LOGGER.ERROR_LOG)
            UPDATE_LOGGER.print_log("write chunklist complete", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError
        
        return startoffset
    
    def write_hash_info(self, image_number, package_file, startoffset):
        UPDATE_LOGGER.print_log("write image hash info StartOffset:%s"\
            % startoffset)
        try:
            # pack hash info head
            hash_info_tlv = struct.pack(CHUNK_HASH_HEADER_FMT, 
                                        self.chunkhash_info_tlv_type, 
                                        2, image_number)
            # Write to bin file
            package_file.seek(startoffset)
            package_file.write(hash_info_tlv)
            startoffset += len(hash_info_tlv)
            
            print(f"Packed image name TLV: {hash_info_tlv}")
            
        except struct.error:
            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError
        
        return startoffset  
        
    
    def write_image_hashdata(self, image_file, package_file, startoffset):
        UPDATE_LOGGER.print_log("write image hash StartOffset:%s"\
            % startoffset)
        try:
            # Step 1: Pack hash partition name
            image_name_tlv = struct.pack(CHUNK_HASH_PARTITION_FMT,
                                         self.chunkhash_partition_tlv_type, 
                                         len(image_file)) + image_file.encode('utf-8')
            partition_len = len(image_name_tlv)
            # Write to bin file
            package_file.seek(startoffset)
            package_file.write(image_name_tlv)
            startoffset += partition_len
            
            print(f"Packed image name TLV: {image_name_tlv}")
            
            # Step 2: Pack target partition hash value
            each_image_file = os.path.join(
                OPTIONS_MANAGER.target_package_dir,
                '%s.img' % image_file)
            image_hash_data = get_file_sha256(each_image_file)
            image_hash_tlv = struct.pack(CHUNK_HASH_VALUE_FMT, 
                                         self.chunkhash_value_tlve_type,
                                         len(image_hash_data)) + image_hash_data.lower().encode('utf-8')
            image_hash_data_len = len(image_hash_tlv)
            print(f'hash data:{image_hash_data.lower}')
            # Write to bin file
            package_file.seek(startoffset)
            package_file.write(image_hash_tlv)
            startoffset += image_hash_data_len
            
            print(f"Packed hash data TLV: {image_hash_tlv}")
            
        except struct.error:
            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError
        return startoffset
      
    def write_image_large(self, image_file, package_file, startoffset):
        UPDATE_LOGGER.print_log("write image large StartOffset:%s"\
            % startoffset)
        try:
            if OPTIONS_MANAGER.full_img_list:
                image_length = len(OPTIONS_MANAGER.full_image_new_data[image_file])
            else:
                image_length = len(OPTIONS_MANAGER.diff_image_new_data[image_file])
            # Step 1: Pack hash partition name
            image_large_tlv = struct.pack(CHUNK_LARGE_PARTITION_FMT,
                                         self.chunkimage_large_tlv_type,
                                         8) + image_length.to_bytes(8, byteorder='little')
            partition_large_len = len(image_large_tlv)
            # Write to bin file
            package_file.seek(startoffset)
            package_file.write(image_large_tlv)
            startoffset += partition_large_len

            print(f"Packed image large TLV: {image_large_tlv}")
            
        except struct.error:
            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError    
        return startoffset
    
    def write_all_image_signdata(self, signdata):
        try:
            signdata_len = len(signdata)
            signdata_tlv = struct.pack(CHUNK_SIGN_VALUE_FMT, 
                                       self.chunksign_value_tlv_type, 
                                       signdata_len)
        except struct.error:
            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError

        # write signdata
        self.signdata = signdata_tlv + signdata
        UPDATE_LOGGER.print_log("Write hashdata sign tlv complete") 
        return True
    
    def write_chunklist_full_image(self, image_name, package_file, chunks, block_sets, startoffset):
        try:
            UPDATE_LOGGER.print_log("write chunk StartOffset:%s"\
            % startoffset)
            
            chunk_index = 0
            for chunk in chunks:
                # Step 1: Pack partition name
                partition_tlv = struct.pack(CHUNK_DATA_PARTITION_FMT, 
                                            self.chunkdata_partition_tlv_type, 
                                            len(image_name)) + image_name.encode()
                # Write to binary file
                package_file.seek(startoffset)
                package_file.write(partition_tlv)
                startoffset += len(partition_tlv)
                
                # Step 2: Pack command info
                cmd_str = ("%s %s %d,%s,%s" % ("new", get_chunk_sha256(chunk), 2,
                                        min(block_sets[chunk_index]), 
                                        max(block_sets[chunk_index]) + 1))
                cmd_tlv = struct.pack(CHUNK_DATA_CMD_FMT, 
                                      self.chunkdata_cmd_tlv_type, 
                                      len(cmd_str)) + cmd_str.encode()
                # Write to binary file
                package_file.seek(startoffset)
                package_file.write(cmd_tlv)
                startoffset += len(cmd_tlv)
                
                # Step 3: Pack the sliced image data
                data_tlv = struct.pack(CHUNK_DATA_DATA_FMT, 
                                       self.chunkdata_value_tlv_type, 
                                       len(chunk)) + chunk
                # Write to binary file
                package_file.seek(startoffset)
                package_file.write(data_tlv)
                startoffset += len(data_tlv)
                chunk_index += 1
        except struct.error:
            UPDATE_LOGGER.print_log("Pack fail!", log_type=UPDATE_LOGGER.ERROR_LOG)
            raise RuntimeError
        
        return startoffset
    
    def pack_dependency_data(self, image, cmd_info, chunk, patch_index, new_index):
        if "pkgdiff" in cmd_info:
            if OPTIONS_MANAGER.image_patch_dic[image]:
                data_value = OPTIONS_MANAGER.image_patch_dic[image][patch_index]
                if not data_value:
                    UPDATE_LOGGER.print_log("data_value is empty, using chunk instead.", 
                                            log_type=UPDATE_LOGGER.INFO_LOG)
                    data_value = chunk
                    data_tlv = struct.pack(CHUNK_DATA_DATA_FMT, 
                                            self.chunkdata_value_tlv_type, 
                                            len(data_value)) + data_value.encode('utf-8')
                    patch_index += 1
                else:
                    data_tlv = struct.pack(CHUNK_DATA_DATA_FMT, 
                                            self.chunkdata_value_tlv_type, 
                                            len(data_value)) + data_value
                    patch_index += 1
            else:
                data_value = chunk
                UPDATE_LOGGER.print_log("patch.data is empty!", log_type=UPDATE_LOGGER.ERROR_LOG)
                raise RuntimeError
            # Determine if the line in transfer.list contains new, if it does, take new.data with it.
        elif "new" in cmd_info:
            if OPTIONS_MANAGER.image_new_dic[image]:
                data_value = OPTIONS_MANAGER.image_new_dic[image][new_index]
                if not data_value:
                    UPDATE_LOGGER.print_log("data_value is empty, using chunk instead.", 
                                            log_type=UPDATE_LOGGER.INFO_LOG)
                    data_value = chunk
                    data_tlv = struct.pack(CHUNK_DATA_DATA_FMT, 
                                            self.chunkdata_value_tlv_type, 
                                            len(data_value)) + data_value.encode('utf-8')
                    new_index += 1
                else:
                    data_tlv = struct.pack(CHUNK_DATA_DATA_FMT, 
                                            self.chunkdata_value_tlv_type, 
                                            len(data_value)) + data_value
                    new_index += 1
            else: 
                data_value = chunk 
                UPDATE_LOGGER.print_log("new.data is empty!", log_type=UPDATE_LOGGER.ERROR_LOG)
                raise RuntimeError
        else:
            # If none of the above instructions are met, put 0 in data_value
            data_value = b''
            data_tlv = struct.pack(CHUNK_DATA_DATA_FMT, 
                                    self.chunkdata_value_tlv_type, 0) + data_value
        return data_tlv, patch_index, new_index