wangchuanli
9f00a2d36d
- 使用 argparse 提供 encrypt/decrypt 子命令，支持交互式密码与分卷大小解析 - 增加输入目录检查、密码错误提示和分卷识别容错 - 添加 .gitignore、LICENSE 和 README.md - 将 CFB 模式迁移至 decrepit 模块并提取加密参数常量
215 行
7.6 KiB
Python
215 行
7.6 KiB
Python
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()
|