import binascii
import hashlib
import struct


class ChunkParser:
    def __init__(self):
        self.chunkdata_partition_tlv_type = 0x01  # Example type
        self.chunkdata_cmd_tlv_type = 0x02  # Example type
        self.chunkdata_value_tlv_type = 0x03  # Example type
        self.chunkhash_info_tlv_type = 0x04  # Example type
        
        # Define formats for struct.pack and struct.unpack
        self.chunkdata_partiton_fmt = '<HH'
        self.chunkdata_cmd_fmt = '<HH'
        self.chunkdata_data_fmt = '<HI'
        self.chunkhash_header_fmt = '<3H'
        self.chunkhash_data_fmt = '<HH'
        self.chunksign_value_fmt = '<HI'


    def parse_chunklist(self, package_file, offset):
        print(f'start read chunklist offset:{offset}')
        package_file.seek(offset)  # Start reading from the beginning
        while True:
            # Read partition TLV
            partition_data = package_file.read(struct.calcsize(self.chunkdata_partiton_fmt))
            print(f'partition_data:{partition_data}')
            if not partition_data:
                break  # End of file
            
            partition_type, partition_length = struct.unpack(self.chunkdata_partiton_fmt, partition_data[:4])
            partition_name = package_file.read(partition_length).decode('utf-8')
            print(f"Partition Type: {partition_type}, Partition Name: {partition_name}")
            if(partition_type != 18):
                print(f'stop read chunklist')
                break
            
            # Read command TLV
            cmd_data = package_file.read(struct.calcsize(self.chunkdata_cmd_fmt))
            cmd_type, cmd_length = struct.unpack(self.chunkdata_cmd_fmt, cmd_data[:4])
            cmd_info = package_file.read(cmd_length).decode('utf-8')
            print(f"Command Type: {cmd_type}, Command Info: {cmd_info}")

            # Read data TLV
            data_data = package_file.read(struct.calcsize(self.chunkdata_data_fmt))
            print(f'self.chunkdata_data_fmt :{data_data}')
            data_type, data_length = struct.unpack(self.chunkdata_data_fmt, data_data[:6])
            data_value = package_file.read(data_length)
            print(f"Data Type: {data_type}, Data Length: {data_length}, Data Value: {data_value}")

            # Process the data based on cmd_info (extract patch or new data)
            if 'pkgdiff' in cmd_info:
                with open('new_patch.dat', 'ab') as patch_file:
                    patch_file.write(data_value)
            elif 'new' in cmd_info:
                with open('new_new.dat', 'ab') as new_file:
                    new_file.write(data_value)
            with open('new_transfer.list', 'ab') as list_file:
                list_file.write((cmd_info + '\n').encode())
                
    def parse_hash_info(self, package_file, offset):
        print(f'start read hash_info offset:{offset}')
        package_file.seek(offset)  # Start reading from the beginning
        while True:
            hash_info_data = package_file.read(struct.calcsize(self.chunkhash_header_fmt))
            if not hash_info_data:
                break  # End of file
            
            hash_type, num_images, image_number = struct.unpack(self.chunkhash_header_fmt, hash_info_data)
            print(f"Hash Type: {hash_type}, Number of Images: {num_images}, Image Number: {image_number}")
            if(hash_type != 22):
                print(f'stop read hash_info offset')
                break
                
    def parse_hash_data(self, package_file, offset):
        print(f'start read hash_data offset:{offset}')
        package_file.seek(offset)
        while True:
            partition_data = package_file.read(struct.calcsize(self.chunkdata_partiton_fmt))
            print(f'partition_data:{partition_data}')
            if not partition_data:
                break  # End of file
            
            partition_type, partition_length = struct.unpack(self.chunkdata_partiton_fmt, partition_data[:4])
            if(partition_type != 23):
                print(f'stop read hash_data offset')
                break
            partition_name = package_file.read(partition_length).decode('utf-8')
            print(f"Partition Type: {partition_type}, Partition Name: {partition_name}")
            
            hash_data = package_file.read(struct.calcsize(self.chunkhash_data_fmt))
            print(f'self.chunkdata_data_fmt :{hash_data}')
            data_type, data_length = struct.unpack(self.chunkhash_data_fmt, hash_data[:4])
            data_value = package_file.read(data_length)
            print(f"Data Type: {data_type}, Data Length: {data_length}, Data Value: {data_value}")

            large_data = package_file.read(struct.calcsize('<2HI')) 
            print(f'self.chunkdata_data_fmt :{large_data}')
            large_data_type, large_data_length, large_data_value = struct.unpack('<2HI', large_data[:8])
            print(f"image large Type: {large_data_type}, image large Length: {large_data_length}, image large Value: {large_data_value}")
            
            with open('new_hash.txt', 'ab') as file:
                file.write(data_value)
                
    def parse_full_sign(self, package_file, offset):
        print(f'start read full image sign offset:{offset}')
        package_file.seek(offset)
        while True:
            sign_data = package_file.read(struct.calcsize(self.CHUNK_SIGN_VALUE_FMT))
            print(f'self.chunkdata_data_fmt :{sign_data}')
            data_type, data_length = struct.unpack(self.CHUNK_SIGN_VALUE_FMT, sign_data[:6]) 
            data_value = package_file.read(data_length)
            print(f"Data Type: {data_type}, Data Length: {data_length}, Data Value: {data_value}")
            if(data_type != 26):
                print(f'stop read full image sign')
                break
            with open('new_sign.txt', 'ab') as file:
                hex_data = binascii.hexlify(data_value)
                file.write(hex_data + b'\n')
                
                
    def parse_full_stream(self, package_file, offset):
        print(f'start read full stream offset:{offset}')
        package_file.seek(offset)
        while True:
            partition_name = package_file.read(struct.calcsize(self.chunkdata_partiton_fmt))
            print(f'self.chunkdata_data_fmt :{partition_name}')
            partition_type, partition_length = struct.unpack(self.chunkdata_partiton_fmt, partition_name[:4])
            partiton_value = package_file.read(partition_length)
            print(f"Data Type: {partition_type}, Data Length: {partition_length}, Data Value: {partiton_value}")
            if(partition_type != 18):
                print(f'stop read full image')
                break
            cmd = package_file.read(struct.calcsize(self.chunkdata_partiton_fmt))
            print(f'self.chunkdata_data_fmt :{cmd}')
            cmd_type, cmd_length = struct.unpack(self.chunkdata_partiton_fmt, cmd[:4])
            cmd_value = package_file.read(cmd_length).decode('utf-8')
            print(f"Data Type: {cmd_type}, Data Length: {cmd_length}, Data Value: {cmd_value}")
            with open('new_full.list', 'ab') as file:
                file.write((cmd_value + '\n').encode())
            
            image_data = package_file.read(struct.calcsize(self.chunkdata_data_fmt))
            data_type, data_length = struct.unpack(self.chunkdata_data_fmt, image_data[:6])
            data_value = package_file.read(data_length)
            with open('new_iamge.img', 'ab') as file:
                file.write(data_value)


def get_file_sha256(update_package):
    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 calculate_file_hash(file_path):
    hash_sha256 = hashlib.sha256()  # You can change this to sha1(), md5(), etc. if needed
    with open(file_path, 'rb') as f:
        # Read the file in chunks to avoid using too much memory
        for chunk in iter(lambda: f.read(4096), b""):
            hash_sha256.update(chunk)
    return hash_sha256.hexdigest() 
  
          
def read_update_bin(file_path):
    with open(file_path, 'rb') as f:
        return f.read()


def parse_update_bin(data, offsets):
    parsed_data = {}
    for name, (offset, length) in offsets.items():
        parsed_data[name] = data[offset:offset + length]
    return parsed_data


def reconstruct_patch(parsed_data):
    # 假设patch是一个字节串，您可以根据需要进行修改
    patch = b''
    for name in parsed_data:
        patch += parsed_data[name]
    return patch


def verify_patch(original_patch, reconstructed_patch):
    return original_patch == reconstructed_patch


def main():
    chunk_parser = ChunkParser()
    with open('update.bin', 'rb') as package_file:
        # 全量流式旧步骤说明（参数仅供历史参考),参数需要在制作包的打印信息中获取:
        # 步骤1 - parse_full_stream 偏移量 632, 使用chunk_parser.parse_full_stream(package_file, 632)
        # 步骤2 - parse_hash_info 偏移量 1975016790，使用chunk_parser.parse_hash_info(package_file, 1975016790)
        # 步骤3 - parse_hash_data 偏移量 1975016796,使用chunk_parser.parse_hash_data(package_file, 1975016796)
        # 步骤4 - parse_full_sign 偏移量 1975017344,使用chunk_parser.parse_full_sign(package_file, 1975017344)
        chunk_parser.parse_chunklist(package_file, 632)
        chunk_parser.parse_hash_info(package_file, 21590734)
        chunk_parser.parse_hash_data(package_file, 21590740)
        chunk_parser.parse_full_sign(package_file, 21590922)
    

if __name__ == "__main__":
    main()