file-encrypt/main.py

import os
import zlib
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.hashes import SHA256
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend


# 辅助函数：使用AES加密数据
def encrypt_data(key, data):
    iv = os.urandom(16)  # 生成随机的初始化向量（IV）
    cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=default_backend())
    encryptor = cipher.encryptor()
    encrypted_data = iv + encryptor.update(data) + encryptor.finalize()
    return encrypted_data

# 辅助函数：使用AES解密数据
def decrypt_data(key, data):
    iv = data[:16]  # 提取初始化向量（IV）
    encrypted_content = data[16:]
    cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=default_backend())
    decryptor = cipher.decryptor()
    decrypted_data = decryptor.update(encrypted_content) + decryptor.finalize()
    return decrypted_data

# 辅助函数：通过密码派生加密密钥
def derive_key(password, salt):
    kdf = PBKDF2HMAC(
        algorithm=SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend()
    )
    return kdf.derive(password.encode())

# 压缩文件数据
def compress_data(data):
    return zlib.compress(data)

# 解压文件数据
def decompress_data(data):
    return zlib.decompress(data)

# 加密并压缩单个文件
def encrypt_file(file_path, password, volume_size, output_dir):
    # 读取文件内容
    with open(file_path, 'rb') as f:
        data = f.read()

    # 压缩文件内容
    compressed_data = compress_data(data)

    # 加密压缩后的内容
    salt = os.urandom(16)
    key = derive_key(password, salt)
    encrypted_data = encrypt_data(key, compressed_data)

    # 如果数据大小超过分卷限制，分卷存储
    file_name = os.path.basename(file_path)
    encrypted_dir = os.path.join(output_dir, file_name)
    os.makedirs(encrypted_dir, exist_ok=True)

    if len(encrypted_data) > volume_size:
        start = 0
        part_number = 1
        while start < len(encrypted_data):
            end = start + volume_size
            part_data = encrypted_data[start:end]
            part_file_path = os.path.join(encrypted_dir, f"{file_name}.part{part_number}")
            with open(part_file_path, 'wb') as part_file:
                if part_number == 1:
                    part_file.write(salt)
                part_file.write(part_data)
            start = end
            part_number += 1
    else:
        # 如果数据不需要分卷，直接保存
        output_path = os.path.join(encrypted_dir, f"{file_name}.enc")
        with open(output_path, 'wb') as enc_file:
            enc_file.write(salt + encrypted_data)

    print(f"Encrypted and saved: {file_path} to {encrypted_dir}")

# 解密并解压单个文件或分卷
def decrypt_file(encrypted_dir, password, output_dir):
    # 读取所有分卷或单个加密文件
    parts = [os.path.join(encrypted_dir, f) for f in os.listdir(encrypted_dir) if f.rsplit('.')[-1].startswith("part") or f.endswith(".enc")]
    parts = sorted(parts)  # 确保按顺序读取
    print(f"Parts found: {parts}")
    encrypted_data = b""
    salt = None
    for i, part in enumerate(parts):
        with open(part, 'rb') as part_file:
            if i == 0:
                salt = part_file.read(16)
            encrypted_data += part_file.read()

    if salt is None:
        raise ValueError("Salt not found in the first part.")

    # 解密数据
    key = derive_key(password, salt)
    decrypted_data = decrypt_data(key, encrypted_data)

    # 解压数据
    decompressed_data = decompress_data(decrypted_data)
    output_file = os.path.join(output_dir, os.path.basename(encrypted_dir))
    with open(output_file, 'wb') as f:
        f.write(decompressed_data)

    print(f"Decrypted and extracted to: {output_file}")

# 加密文件夹中的所有文件
def encrypt_directory(directory, password, volume_size, output_dir):
    for root, _, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            encrypt_file(file_path, password, volume_size, output_dir)

# 解密文件夹中的所有加密文件
def decrypt_directory(directory, password, output_dir):
    for root, dirs, _ in os.walk(directory):
        for dir in dirs:
            encrypted_dir = os.path.join(root, dir)
            decrypt_file(encrypted_dir, password, output_dir)

if __name__ == "__main__":
    action = "encrypt"  # 或 "decrypt"
    # action = "decrypt"  # 或 "decrypt"
    password = "123456"
    volume_size = 10 * 1024 * 1024  # 每个分卷大小（10 MB）
    source_dir = "./data"  # 原始数据文件夹
    encrypted_dir = "./encrypted"  # 加密文件夹
    decrypted_dir = "./decrypted"  # 解密文件夹

    if action == "encrypt":
        encrypt_directory(source_dir, password, volume_size, encrypted_dir)
    elif action == "decrypt":
        decrypt_directory(encrypted_dir, password, decrypted_dir)
    else:
        print("Invalid action. Use 'encrypt' or 'decrypt'.")