encrypt.py (3713B)
1 import os 2 from cryptography.hazmat.backends import default_backend 3 from cryptography.hazmat.primitives.asymmetric import rsa 4 from cryptography.hazmat.primitives import serialization 5 from cryptography.hazmat.primitives import hashes 6 from cryptography.hazmat.primitives.asymmetric import padding 7 from cryptography.fernet import Fernet 8 from prompts import success 9 10 11 def RSA(target_file, public_key_source): 12 """Encrypts the passed file with the passed RSA public key 13 14 Keyword arguments: 15 target_file -- the filepath to the file to be encrypted 16 public_key_source -- the filepath to the public key 17 """ 18 with open(public_key_source, "rb") as public_key_file: 19 public_key = serialization.load_pem_public_key( 20 public_key_file.read(), 21 backend=default_backend() 22 ) 23 try: 24 with open(target_file) as read_file: 25 file_data = read_file.read() 26 file_data = bytes(file_data, "utf-8") 27 except UnicodeDecodeError: 28 with open(target_file, "rb") as read_file: 29 file_data = read_file.read() 30 data = public_key.encrypt( 31 file_data, 32 padding.OAEP( 33 mgf=padding.MGF1(algorithm=hashes.SHA256()), 34 algorithm=hashes.SHA256(), 35 label=None 36 ) 37 ) 38 data += b"0" 39 with open(target_file, "wb") as write_file: 40 write_file.write(data) 41 42 43 def Symmetric(target_file, symmetric_key_source): 44 """Encrypts the passed file with the passed symmetric key 45 46 Keyword arguments: 47 target_file -- the filepath to the file to be encrypted 48 symmetric_key_source -- the filepath to the symmetric key 49 """ 50 with open(symmetric_key_source, "rb") as symmetric_key_file: 51 symmetric_key_data = symmetric_key_file.read() 52 symmetric_key = Fernet(symmetric_key_data) 53 try: 54 with open(target_file) as read_file: 55 file_data = read_file.read() 56 file_data = bytes(file_data, "utf-8") 57 except UnicodeDecodeError: 58 with open(target_file, "rb") as read_file: 59 file_data = read_file.read() 60 data = symmetric_key.encrypt(file_data) 61 data += b"1" 62 with open(target_file, "wb") as write_file: 63 write_file.write(data) 64 65 def enc_manager(target_files, save_folder): 66 """Encrypt all files passed to the function with the symmetric key, 67 and then replace the symmetric key file's contents with an encrypted 68 version, encrypted with the public key. 69 70 Keyword arguments: 71 target_file_raw -- a string composed of file locations seperated by colons 72 save_folder -- the location of the saved key trio 73 """ 74 pub_src = save_folder + "/public_key.pem" 75 sym_src = save_folder + "/symmetric_key.key" 76 if not os.path.exists(sym_src): 77 for fl in target_files: 78 RSA(fl, pub_src) 79 else: 80 for fl in target_files: 81 if os.path.getsize(fl) > 446: 82 Symmetric(fl, sym_src) 83 else: 84 RSA(fl, pub_src) 85 with open(pub_src, "rb") as pub_file, \ 86 open(sym_src, "rb") as sym_file: 87 public_key = serialization.load_pem_public_key( 88 pub_file.read(), 89 backend=default_backend() 90 ) 91 symmetric_key_data = sym_file.read() 92 encrypted_key = public_key.encrypt( 93 symmetric_key_data, 94 padding.OAEP( 95 mgf = padding.MGF1(algorithm = hashes.SHA256()), 96 algorithm = hashes.SHA256(), 97 label = None 98 ) 99 ) 100 with open(sym_src, "wb") as crypto_key_file: 101 crypto_key_file.write(encrypted_key) 102 success("enc")