encryptFiles
This commit is contained in:
197
encryptFiles/main.py
Normal file
197
encryptFiles/main.py
Normal file
@ -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()
|
||||
Reference in New Issue
Block a user