file-encrypt/main.py

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

SALT_SIZE = 16
IV_SIZE = 16
KEY_SIZE = 32
PBKDF2_ITERATIONS = 100_000

# 辅助函数：使用AES加密数据
def encrypt_data(key, data):
    iv = os.urandom(IV_SIZE)  # 生成随机的初始化向量（IV）
    cipher = Cipher(algorithms.AES(key), 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[:IV_SIZE]  # 提取初始化向量（IV）
    encrypted_content = data[IV_SIZE:]
    cipher = Cipher(algorithms.AES(key), 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=KEY_SIZE,
        salt=salt,
        iterations=PBKDF2_ITERATIONS,
        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(SALT_SIZE)
    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} -> {encrypted_dir}")

# 解密并解压单个文件或分卷
def decrypt_file(encrypted_dir, password, output_dir):
    # 读取所有分卷或单个加密文件
    try:
        entries = os.listdir(encrypted_dir)
    except FileNotFoundError:
        print(f"[ERROR] Encrypted directory not found: {encrypted_dir}")
        return

    parts = []
    for f in entries:
        # 分卷文件名形如 <name>.part1 / .part2 ...；单文件为 <name>.enc
        base, _, ext = f.rpartition('.')
        if ext.startswith("part") or ext == "enc":
            parts.append(os.path.join(encrypted_dir, f))

    if not parts:
        print(f"[WARN] No encrypted parts found in: {encrypted_dir}")
        return

    parts = sorted(parts)  # 确保按顺序读取

    encrypted_data = b""
    salt = None
    for i, part in enumerate(parts):
        with open(part, 'rb') as part_file:
            if i == 0:
                # 首个分卷（或单文件 .enc）的前 SALT_SIZE 字节为 salt
                salt = part_file.read(SALT_SIZE)
            encrypted_data += part_file.read()

    if not salt or len(salt) != SALT_SIZE:
        print(f"[ERROR] Invalid salt in first part of: {encrypted_dir}")
        return

    # 解密数据
    try:
        key = derive_key(password, salt)
        decrypted_data = decrypt_data(key, encrypted_data)
        # 解压数据（密码错误时这里会抛异常）
        decompressed_data = decompress_data(decrypted_data)
    except (zlib.error, ValueError):
        print(f"[ERROR] Failed to decrypt {encrypted_dir}. Wrong password or corrupted data.")
        return

    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):
    if not os.path.isdir(directory):
        print(f"[ERROR] Source directory not found: {directory}")
        sys.exit(1)
    os.makedirs(output_dir, exist_ok=True)
    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):
    if not os.path.isdir(directory):
        print(f"[ERROR] Encrypted directory not found: {directory}")
        sys.exit(1)
    os.makedirs(output_dir, exist_ok=True)
    for root, dirs, _ in os.walk(directory):
        for d in dirs:
            encrypted_dir = os.path.join(root, d)
            decrypt_file(encrypted_dir, password, output_dir)

def parse_volume_size(text):
    """解析分卷大小，支持 KB/MB/GB 后缀，如 10MB。"""
    units = {"KB": 1024, "MB": 1024**2, "GB": 1024**3}
    text = text.strip().upper()
    for suffix, factor in units.items():
        if text.endswith(suffix):
            return int(text[:-len(suffix)]) * factor
    return int(text)  # 纯数字视为字节

def main():
    parser = argparse.ArgumentParser(
        description="基于密码的文件加密工具（AES-256-CFB + PBKDF2 + zlib 压缩 + 分卷存储）。"
    )
    sub = parser.add_subparsers(dest="action", required=True)

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("-p", "--password", help="加密/解密密码。省略时将交互式输入（不在命令行暴露密码）。")
    common.add_argument("-v", "--volume-size", default="10MB",
                        help="分卷大小，支持 KB/MB/GB 后缀（默认 10MB）。")

    enc = sub.add_parser("encrypt", parents=[common], help="加密目录")
    enc.add_argument("source", help="待加密的源目录")
    enc.add_argument("output", help="加密输出目录")

    dec = sub.add_parser("decrypt", parents=[common], help="解密目录")
    dec.add_argument("source", help="加密数据所在目录")
    dec.add_argument("output", help="解密输出目录")

    args = parser.parse_args()

    # 密码获取：优先参数，否则交互输入
    password = args.password if args.password else getpass.getpass("Enter password: ")
    if not password:
        print("[ERROR] Password cannot be empty.")
        sys.exit(1)

    try:
        volume_size = parse_volume_size(args.volume_size)
    except ValueError:
        print(f"[ERROR] Invalid volume size: {args.volume_size}")
        sys.exit(1)

    if args.action == "encrypt":
        encrypt_directory(args.source, password, volume_size, args.output)
    elif args.action == "decrypt":
        decrypt_directory(args.source, password, args.output)

if __name__ == "__main__":
    main()