encryptFiles

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()