diff --git a/encryptFiles/main.py b/encryptFiles/main.py new file mode 100644 index 0000000..4a4a4cf --- /dev/null +++ b/encryptFiles/main.py @@ -0,0 +1,197 @@ +import os +import hashlib +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.hazmat.primitives import hashes +from cryptography.hazmat.primitives.asymmetric import padding +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey,RSAPublicKey +import concurrent.futures +from rich.console import Console +from rich.progress import Progress +from time import time + +console=Console() + +# 获取32位MD5哈希值 +def get_md5_hash(password): + md5_hash = hashlib.md5(password.encode()) + return md5_hash.hexdigest() + +# 生成RSA密钥对 +def generate_rsa_key_pair(): + private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, + backend=default_backend() + ) + public_key = private_key.public_key() + return private_key, public_key + +# 用AES加密公钥，并保存到文件 +def encrypt_and_save_private_key(private_key:RSAPrivateKey, aes_key:str,folder_path:str): + pwd=get_md5_hash(aes_key) + private_key_bytes =private_key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.BestAvailableEncryption(pwd.encode()) + ) + with open(os.path.join(folder_path, 'private_key.der'), 'wb') as file: + file.write(private_key_bytes) + +# 加密文件 +def encrypt_file(file_path, public_key:RSAPublicKey): + out=[] + with open(file_path, 'rb') as file: + while True: + data = file.read(100) + if not data: + break + ciphertext = public_key.encrypt( + data, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None + ) + ) + out.append(ciphertext) + with open(file_path + '.enc', 'wb') as file: + file.write(b'\xc0\x12\x87\x3e\x30\x55\x79\xc2\x25\x43'.join(out)) + os.remove(file_path) + +# 解密文件 +def decrypt_file(file_path, private_key_byte): + private_key=serialization.load_der_private_key(private_key_byte,None) + # console.print(f"[yellow]Decrypting {file_path[:10]}...[/yellow]") + + out=[] + with open(file_path, 'rb') as file: + datas=file.read().split(b'\xc0\x12\x87\x3e\x30\x55\x79\xc2\x25\x43') + for data in datas: + out.append(private_key.decrypt( + data, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None + ) + )) + os.remove(file_path) + with open(file_path[:-4], 'wb') as file: + file.write(b"".join(out)) + # console.print(f"[green]Decrypted {file_path[:10]}[/green]") + return file_path[:-4],out[0] + +# 检查文件夹是否已加密 +def is_folder_encrypted(folder_path): + for filepath in os.listdir(folder_path): + if filepath.endswith('.enc'):return True + return False + +def is_folder_init(folder_path): + return os.path.exists(os.path.join(folder_path, 'private_key.der')) + +# 加密文件夹中的所有文件 +def encrypt_folder(folder_path, public_key): + cnt=0 + for root, dirs, files in os.walk(folder_path): + cnt+=len(files) + with Progress() as progress: + task1 = progress.add_task("[cyan]Encrypting...", total=cnt-2) + for root, dirs, files in os.walk(folder_path): + for file in files: + if file in ["public_key.pem","private_key.der"]: + continue + file_path = os.path.join(root, file) + if os.path.isfile(file_path) and not file_path.endswith('.enc'): + encrypt_file(file_path, public_key) + progress.update(task1, advance=1) + +# 解密文件夹中的所有文件 +def decrypt_folder(folder_path, private_key): + # for filename in os.listdir(folder_path): + # if filename in ["public_key.pem","private_key.der"]: + # continue + # file_path = os.path.join(folder_path, filename) + # if os.path.isfile(file_path) and file_path.endswith('.enc'): + # decrypted_data = decrypt_file(file_path, private_key) + # with open(file_path[:-4], 'wb') as file: + # file.write(decrypted_data) + # progress.advance(task1,1) + with concurrent.futures.ProcessPoolExecutor(max_workers=16) as executor: + futures = [] + for root, dirs, files in os.walk(folder_path): + for file in files: + if file in ["public_key.pem","private_key.der"]: + continue + file_path = os.path.join(root, file) + if os.path.isfile(file_path) and file_path.endswith('.enc'): + future = executor.submit(decrypt_file, file_path, private_key) + futures.append(future) + + with Progress() as progress: + task1 = progress.add_task("[cyan]Decrypting...", total=len(futures)) + while futures: + for f1 in futures: + if f1.done(): + progress.advance(task1, 1) + futures.pop(futures.index(f1)) + break + + +# 主函数 +def main(): + import sys + if len(sys.argv)!=2: + exit(1) + folder_path = sys.argv[1] # 指定文件夹路径 + if not is_folder_init(folder_path): + console.print("[green]Initializing[/green]") + # 生成RSA密钥对 + private_key, public_key = generate_rsa_key_pair() + # 获取用户输入的密码并转换为32位MD5哈希值 + pwd = console.input("Enter password: ",password=True) + encrypt_and_save_private_key(private_key, pwd,folder_path) + with open (os.path.join(folder_path,'public_key.pem'), 'wb') as f: + pub = public_key.public_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PublicFormat.SubjectPublicKeyInfo + ) + f.write(pub) + # 加密文件夹中的所有文件 + st=time() + encrypt_folder(folder_path, public_key) + console.print(f"[green]Encrypted in %.2f s[/green]"%(time()-st)) + else: + if not is_folder_encrypted(folder_path): + console.print("[yellow]Encrypt[/yellow]") + with open(os.path.join(folder_path,'public_key.pem'), 'rb') as f: + padded_data=f.read() + public_key = serialization.load_pem_public_key(padded_data) + st=time() + encrypt_folder(folder_path, public_key) + print(f"[green]Encrypted in %.2f s[/green]"%(time()-st)) + else: + # 加密过，需要解密公钥并解密文件夹中的文件 + console.print("[green]Decrypt[/green]") + with open(os.path.join(folder_path,'private_key.der'), 'rb') as file: + encrypted_private_key = file.read() + pwd = console.input("Enter password: ",password=True) + + public_key = serialization.load_der_private_key( + encrypted_private_key, + get_md5_hash(pwd).encode() + ) + pub = public_key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + ) + st=time() + decrypt_folder(folder_path, pub) + console.print(f"[green]Decrypted in %.2f s[/green]"%(time()-st)) + + +if __name__ == "__main__": + main()