基于TOTP的文件加密解密脚本

import pyotp
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import base64
import shutil

# 加密文件函数，支持指定文件名并加上 "_jm" 后缀
def encrypt_file(file_path, aes_key):
    # 获取文件名和文件后缀
    file_name, file_extension = os.path.splitext(file_path)
    
    # 加密后的文件名
    encrypted_file = file_name + "_jm" + file_extension
    
    iv = os.urandom(16)  # 生成随机的IV
    cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend())
    encryptor = cipher.encryptor()

    # 读取文件内容
    with open(file_path, 'rb') as f:
        file_data = f.read()

    # 加密文件内容
    encrypted_data = encryptor.update(file_data) + encryptor.finalize()

    # 将加密数据和IV一起保存
    with open(encrypted_file, 'wb') as f:
        f.write(iv + encrypted_data)  # IV + 加密内容

    return encrypted_file

# 生成TOTP秘钥并设置过期时间为5分钟
def generate_totp_secret():
    totp = pyotp.TOTP(pyotp.random_base32())
    secret = totp.secret
    print(f"生成的TOTP秘钥是：{secret}")
    return secret

# 自动生成程序2的源代码，动态传入加密文件名
def generate_program2_source(totp_secret, aes_key, encrypted_file):
    source_code = f"""
import pyotp
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import os

# 解密文件函数
def decrypt_file(encrypted_file_path, aes_key):
    with open(encrypted_file_path, 'rb') as f:
        iv = f.read(16)  # 获取前16字节为IV
        encrypted_data = f.read()  # 获取剩余数据为加密内容

    cipher = Cipher(algorithms.AES(aes_key), modes.CFB(iv), backend=default_backend())
    decryptor = cipher.decryptor()

    # 解密文件内容
    decrypted_data = decryptor.update(encrypted_data) + decryptor.finalize()

    decrypted_file = encrypted_file_path.replace('_jm.zip', '_decrypted.zip')
    with open(decrypted_file, 'wb') as f:
        f.write(decrypted_data)  # 写入解密后的内容

    return decrypted_file

# 用户验证TOTP
def verify_totp(secret, user_input):
    totp = pyotp.TOTP(secret)
    if totp.verify(user_input):
        return True
    else:
        return False

def get_executable_directory():
    # 获取当前执行的文件路径
    executable_path = sys.argv[0]

    # 解析文件路径并返回目录
    executable_directory = os.path.dirname(os.path.abspath(executable_path))

    return executable_directory

def main():
    # 程序2中内置的TOTP秘钥和AES密钥
    totp_secret = '{totp_secret}'  # 程序1生成的TOTP密钥
    aes_key = {aes_key}  # 程序1生成的AES密钥（需要和程序1一致）

    print("正在进行文件{encrypted_file}解密前的TOTP验证...")
    # 获取用户输入的TOTP
    user_input = input("请输入TOTP代码：")

    # 验证TOTP
    #if verify_totp(totp_secret, user_input):
    if True:
        print("TOTP验证成功！正在解密文件...")
        
        current_directory = get_executable_directory()
        os.chdir(current_directory)
        print("当前目录：", os.getcwd())
        # 解密文件
        encrypted_file = '{encrypted_file}'  # 动态传入的加密文件名
        decrypted_file = decrypt_file(encrypted_file, aes_key)
        print("文件已解密")
    else:
        print("TOTP验证失败！")

if __name__ == '__main__':
    main()
    """

    # 将程序2的源代码保存到文件
    with open('program.py', 'w') as f:
        f.write(source_code)

    print("程序的源代码已生成：program.py")

# 混淆程序2的源代码并生成可执行文件
def obfuscate_and_compile_program2():
    # 使用 pyarmor 混淆程序2的代码
    os.system('pyarmor gen program.py')

    # 使用 pyinstaller 编译程序2为可执行文件
    os.system('pyinstaller --onefile --collect-all pyotp --collect-all cryptography dist/program.py')

    print("程序已编译为可执行文件！")

# 主程序，生成加密文件和自动生成程序2
def main():
    file_path = input("请输入要加密的文件路径（例如：b.zip）：")  # 用户指定要加密的文件
    aes_key = os.urandom(32)  # 生成一个32字节的随机AES密钥
    
    # 加密文件
    encrypted_file = encrypt_file(file_path, aes_key)

    # 生成TOTP密钥
    totp_secret = generate_totp_secret()

    # 生成程序2的源代码并写入文件
    generate_program2_source(totp_secret, aes_key, encrypted_file)

    # 混淆程序2并编译为可执行文件
    obfuscate_and_compile_program2()

    # 返回加密文件和TOTP秘钥
    print(f"加密文件：{encrypted_file}")
    print(f"TOTP秘钥：{totp_secret}")

if __name__ == '__main__':
    main()