如何修改PyCryptodome加密类避免大文件内存溢出崩溃？

问题背景 #

我用PyCryptodome实现了一个加密类，代码如下：

class Cipher:
    def __init__(self, password: str, key: bytes = None):
        self.key = PBKDF2(password, key if key else password_as_key(password), dkLen=32)
        self.cipher = AES.new(self.key, AES.MODE_CBC, iv=get_random_bytes(16))

    def encrypt(self, data: bytes) -> bytes:
        return self.cipher.encrypt(pad(data, AES.block_size)) + self.cipher.iv

    def decrypt(self, data: bytes) -> bytes:
        encrypted_data = data[:-16]
        iv = data[-16:]
        self.cipher = AES.new(self.key, AES.MODE_CBC, iv=iv)
        return unpad(self.cipher.decrypt(encrypted_data), AES.block_size)

这个类处理常规文件没问题，但处理超大文件时会因内存不足崩溃。环境为Windows 11 Home、8GB内存（可用7.68GB）、12代英特尔酷睿i5-1235U处理器。需要添加防护机制避免崩溃，最好能显示加载界面，内存可用时继续操作。

解决方案 #

1. 核心优化：分块处理文件

原代码一次性将整个文件加载到内存，这是超大文件崩溃的根源。改成按固定块大小分块读写，内存占用会控制在块大小范围内。

修改后的Cipher类（新增文件分块处理方法）：

from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
import os

class Cipher:
    def __init__(self, password: str, key: bytes = None):
        self.key = PBKDF2(password, key if key else password.encode(), dkLen=32)
        # 初始化时不再创建cipher实例，分块处理时按需创建

    def encrypt_file(self, input_path: str, output_path: str, block_size: int = 1024 * 1024):
        # 生成随机IV并写入输出文件开头
        iv = get_random_bytes(16)
        cipher = AES.new(self.key, AES.MODE_CBC, iv=iv)
        
        with open(input_path, 'rb') as infile, open(output_path, 'wb') as outfile:
            outfile.write(iv)  # 先写IV，解密时需要读取
            
            while True:
                chunk = infile.read(block_size)
                if not chunk:
                    break
                # 最后一块需要补全padding
                if len(chunk) % AES.block_size != 0:
                    chunk = pad(chunk, AES.block_size)
                outfile.write(cipher.encrypt(chunk))

    def decrypt_file(self, input_path: str, output_path: str, block_size: int = 1024 * 1024):
        with open(input_path, 'rb') as infile, open(output_path, 'wb') as outfile:
            # 读取开头的IV
            iv = infile.read(16)
            cipher = AES.new(self.key, AES.MODE_CBC, iv=iv)
            
            while True:
                chunk = infile.read(block_size)
                if not chunk:
                    break
                decrypted_chunk = cipher.decrypt(chunk)
                # 最后一块需要去除padding
                if len(chunk) < block_size:
                    decrypted_chunk = unpad(decrypted_chunk, AES.block_size)
                outfile.write(decrypted_chunk)

注：这里假设password_as_key是自定义的密钥派生方法，如果没有可以直接用password.encode()替代

2. 内存监控与防护

用psutil监控可用内存，当内存低于阈值时暂停处理，避免崩溃：

1. 安装依赖：pip install psutil
2. 在分块循环中添加内存检查逻辑：

import psutil
import time

def check_memory_available(threshold_mb: int = 512) -> bool:
    mem = psutil.virtual_memory()
    return mem.available >= threshold_mb * 1024 * 1024

# 在encrypt_file/decrypt_file的循环中修改：
while True:
    # 等待内存可用
    while not check_memory_available():
        time.sleep(1)  # 每秒检查一次
    
    chunk = infile.read(block_size)
    if not chunk:
        break
    # 后续加密/解密逻辑不变

3. 加载界面实现（以Tkinter为例）

如果是GUI应用，添加进度条和加载提示，实时反馈处理状态：

import tkinter as tk
from tkinter import ttk
import threading

class EncryptGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("文件加密工具")
        
        self.progress = ttk.Progressbar(root, orient="horizontal", length=300, mode="determinate")
        self.progress.pack(pady=20)
        
        self.status_label = ttk.Label(root, text="等待开始")
        self.status_label.pack(pady=10)
        
        self.start_btn = ttk.Button(root, text="开始加密", command=self.start_encrypt)
        self.start_btn.pack(pady=10)
        
        self.cipher = Cipher("your_password")
        self.input_path = "large_file.bin"
        self.output_path = "large_file_encrypted.bin"

    def update_progress(self, current, total):
        self.progress["value"] = (current / total) * 100
        self.status_label.config(text=f"已处理：{current//(1024*1024)}MB / {total//(1024*1024)}MB")
        self.root.update_idletasks()

    def encrypt_task(self):
        total_size = os.path.getsize(self.input_path)
        processed = 0
        
        iv = get_random_bytes(16)
        cipher = AES.new(self.cipher.key, AES.MODE_CBC, iv=iv)
        
        with open(self.input_path, 'rb') as infile, open(self.output_path, 'wb') as outfile:
            outfile.write(iv)
            while True:
                while not check_memory_available():
                    time.sleep(1)
                
                chunk = infile.read(1024*1024)
                if not chunk:
                    break
                if len(chunk) % AES.block_size != 0:
                    chunk = pad(chunk, AES.block_size)
                outfile.write(cipher.encrypt(chunk))
                
                processed += len(chunk)
                self.update_progress(processed, total_size)
        
        self.status_label.config(text="加密完成")
        self.start_btn.config(state="normal")

    def start_encrypt(self):
        self.start_btn.config(state="disabled")
        threading.Thread(target=self.encrypt_task).start()

if __name__ == "__main__":
    root = tk.Tk()
    app = EncryptGUI(root)
    root.mainloop()

内容的提问来源于stack exchange，提问作者user21815993