cryptography与zlib系列：数据分块解密与解压

cryptography与zlib系列：数据分块解密与解压

数据分块压缩与加密可以查看我的另外一篇博客，这里与之相对应的是解密与解压的操作。解压之后的数据转为BytesIO对象用于后续读取，适用于大部分数据格式，而不需要从文件中读取。

完整示例

from zlib import decompress
from struct import unpack
from concurrent.futures import ThreadPoolExecutor
from io import BytesIO


def decrypt_decompress_chunk(
    chunk_data, cipher_suite
):
    # Decrypt the chunk
    decrypted_chunk = cipher_suite.decrypt(chunk_data)
    # Decompress the chunk
    decompressed_chunk = decompress(decrypted_chunk)

    return decompressed_chunk

def get_model_decrypt_decompress_data(
    file_path: str,
    cipher_suite=cipher_suite,
    byteio_flag=True
):
    # Open the encrypted file
    with open(file_path, "rb") as fr:
        encrypted_chunks = fr.read()
    
    # 解析加密块和长度信息
    chunks = []
    offset = 0
    while offset < len(encrypted_chunks):
        chunk_size = unpack('<I', encrypted_chunks[offset:offset+4])[0]
        chunk = encrypted_chunks[offset+4:offset+4+chunk_size]
        chunks.append(chunk)
        offset += 4 + chunk_size
    
    # 创建线程池并解密解压缩
    with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
        decompressed_chunks = list(
            executor.map(
                lambda x: decrypt_decompress_chunk(x, cipher_suite), 
                chunks
            )
        )
    
    # 合并解压缩数据
    decompressed_model_data = b''.join(decompressed_chunks)
    if byteio_flag:
        
        decompressed_model_data = BytesIO(decompressed_model_data)
    # 重置文件指针到开始位置，多线程使用需要使用锁保护
    decompressed_model_data.seek(0)
    return decompressed_model_data





if __name__ == "__main__":
    from time import time
    from cryptography.fernet import Fernet

    model_encrypted_path = "model_test.enc"
    model_decrypted_path = "model_test.onnx"
    key =  b''
    tic = time()
    cipher_suite = Fernet(key)
    decompressed_model_data = get_model_decrypt_decompress_data(
        model_encrypted_path,
        cipher_suite,
    )
    toc = time()
    print(f"Decompress and decrypt time: {toc - tic:.2f} seconds")


    tic = time()
    with open(model_decrypted_path, "wb") as fw:
        fw.write(decompressed_model_data)
    toc = time()
    print(f"Write to file time: {toc - tic:.2f} seconds")