key.py (3270B)
1 import os 2 from cryptography.hazmat.backends import default_backend 3 from cryptography.hazmat.primitives.asymmetric import rsa, padding 4 from cryptography.hazmat.primitives import serialization, hashes 5 from cryptography.fernet import Fernet 6 from prompts import success 7 8 9 def rsa_key(pub, priv, passkey): 10 """Generate a private and public key to the provided filepaths. 11 12 Keyword arguments: 13 pub -- path to save the public key 14 priv -- path to save the private key 15 """ 16 private_key = rsa.generate_private_key( 17 public_exponent=65537, 18 key_size=4096, 19 backend=default_backend() 20 ) 21 public_key = private_key.public_key() 22 private_key_text = private_key.private_bytes( 23 encoding=serialization.Encoding.PEM, 24 format=serialization.PrivateFormat.PKCS8, 25 encryption_algorithm=( 26 serialization.BestAvailableEncryption( 27 bytes(passkey, "utf-8") 28 ) if passkey != "" else serialization.NoEncryption() 29 ) 30 ) 31 public_key_text = public_key.public_bytes( 32 encoding=serialization.Encoding.PEM, 33 format=serialization.PublicFormat.SubjectPublicKeyInfo 34 ) 35 with open(priv, "wb") as private_file, \ 36 open(pub, "wb") as public_file: 37 private_file.write(private_key_text) 38 public_file.write(public_key_text) 39 40 41 def mixed_key(pub, priv, sym, passkey): 42 """Generate a private, public, and symmetric key to the 43 provided filepaths. 44 45 Keyword arguments: 46 pub -- path to save the public key 47 priv -- path to save the private key 48 sym -- path to save the symmetric key 49 """ 50 rsa_key(pub, priv, passkey) 51 symmetric_key = Fernet.generate_key() 52 with open(sym, "wb") as sym_file: 53 sym_file.write(symmetric_key) 54 55 56 def key_manager(target_files, save_folder, passkey): 57 """Call either the `rsa_key` function or the `mixed_key` 58 function, according to the needs of the target files. 59 60 Keyword arguments: 61 target_files -- a list of the files to be encrypted 62 save_folder -- the directory to save the keys 63 passkey -- the password to enhance the RSA encryption (OPTIONAL) 64 """ 65 priv_src = save_folder + "/private_key.pem" 66 pub_src = save_folder + "/public_key.pem" 67 sym_src = save_folder + "/symmetric_key.key" 68 rsa = True 69 for fl in target_files: 70 if os.path.getsize(fl) > 446: 71 rsa = False 72 if rsa: 73 rsa_key(pub_src, priv_src, passkey) 74 else: 75 mixed_key( 76 pub_src, 77 priv_src, 78 sym_src, 79 passkey 80 ) 81 82 def just_key_manager(mode, save_folder, passkey): 83 """Call either the `rsa_key` function or the `mixed_key` 84 function, according to passed mode. 85 86 Keyword arguments: 87 mode -- either 0 (RSA) or 1 (Mixed) 88 save_folder -- the directory to save the keys 89 passkey -- the password to enhance the RSA encryption (OPTIONAL) 90 """ 91 rsa = True if mode == 0 else False 92 pub_src = save_folder + "/public_key.pem" 93 priv_src = save_folder + "/private_key.pem" 94 sym_src = save_folder + "/symmetric_key.key" 95 if rsa: 96 rsa_key(pub_src, priv_src, passkey) 97 else: 98 mixed_key(pub_src, priv_src, sym_src, passkey) 99 success("key")