4.11 GETTING AN UPGRADE

EXERCISE 4.11: GETTING AN UPGRADE

Help Alice and Bob out. Rewrite the RSA encryption/decryption program to use the cryptography module instead of gmpy2 operations.

---

# ex4_11.py (way more secure than listing4_4.py)
import os
from cryptography.hazmat.backends import default_backend 
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives import serialization 
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import padding

def rsa_encrypt(m: bytes, public_key: rsa.RSAPublicKey): 
    return public_key.encrypt(
        plaintext=m, 
        padding=padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()), 
            algorithm=hashes.SHA256(), 
            label=None,
        ),
    )

def rsa_decrypt(c: bytes, private_key: rsa.RSAPrivateKey): 
    return private_key.decrypt(
        ciphertext=c, 
        padding=padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()), 
            algorithm=hashes.SHA256(), 
            label=None,
        ),
    )

def main(): 
    public_key_file = None 
    private_key_file = None 
    public_key = None 
    private_key = None 

    while True: 
        print("RSA Crypto") 
        print("------------------")
        print("\tPrivate key file: {}".format(private_key_file))
        print("\tPublic key file: {}".format(public_key_file)) 
        print("\t1. Encrypt Message")
        print("\t2. Decrypt Message.")
        print("\t3. Load public key file ")
        print("\t4. Load private key file.") 
        print("\t5. Create and load new public and private key files.")
        print("\t6. Quit.\n")

        choice = input(">> ") 

        if choice == "1": 
            if not public_key: 
                print("\nNo public key loaded\n") 
            else: 
                message = input("\nPlaintext: ").encode() 
                ciphertext = rsa_encrypt(message, public_key) 
                print("\nCiphertext (hexlified): {}\n".format(ciphertext.hex()))
        elif choice == '2':
            if not private_key: 
                print("\nNo private key loaded\n")
            else: 
                ciphertext_hex = input("\nCiphertext (hexlified): ") 
                ciphertext = bytes.fromhex(ciphertext_hex)
                message = rsa_decrypt(ciphertext, private_key)
                print("\nPlaintext: {}\n".format(message)) 
        elif choice == '3': 
            public_key_temp_file = input("\nEnter public key file: ")
            if not os.path.exists(public_key_temp_file): 
                print(f"File {public_key_temp_file} does not exist.")
                continue 
            with open(public_key_temp_file, 'rb') as f: 
                public_key = serialization.load_pem_public_key(
                    data=f.read(),
                    backend=default_backend()
                )
                public_key_file = public_key_temp_file
                print("\nPublic Key file loaded.\n")

                # unload private key if any
                private_key_file = None 
                private_key = None 
        elif choice == '4':
            private_key_file_temp = input("\nEnter private key file: ")
            if not os.path.exists(private_key_file_temp):
                print(f"File {private_key_file_temp} does not exist.")
                continue 
            with open(private_key_file_temp, 'rb') as f: 
                private_key = serialization.load_pem_private_key(
                    data=f.read(),
                    backend=default_backend(),
                    password=None
                ) 
                private_key_file = private_key_file_temp 
                print("\nPrivate Key file loaded.\n") 

                # load public key for the given private key 
                # (unload previous public key if any) 
                public_key = private_key.public_key() 
                public_key_file = None 
        elif choice == '5': 
            private_key_file_temp = input("\nEnter a file name for new private key: ") 
            public_key_file_temp = input("\nEnter a file name for a new public key: ") 
            
            if os.path.exists(private_key_file_temp) or os.path.exists(public_key_file_temp): 
                print("File already exists") 
                continue

            with open(private_key_file_temp, "wb+") as private_key_file_obj: 
                with open(public_key_file_temp, "wb+") as public_key_file_obj: 
                    private_key = rsa.generate_private_key(
                        public_exponent=65537, 
                        key_size=2048, 
                        backend=default_backend(),
                    )
                    public_key = private_key.public_key() 

                    private_key_file_obj.write(
                        private_key.private_bytes(
                            encoding=serialization.Encoding.PEM, 
                            format=serialization.PrivateFormat.TraditionalOpenSSL, 
                            encryption_algorithm=serialization.NoEncryption(), 
                        )
                    )

                    public_key_file_obj.write(
                        public_key.public_bytes(
                            encoding=serialization.Encoding.PEM, 
                            format=serialization.PublicFormat.SubjectPublicKeyInfo,
                        )
                    )

                    public_key_file = None 
                    private_key_file = private_key_file_temp
        elif choice == '6':
            print("\n\nTerminating. This program will self destruct in 5 seconds. \n")
            break
        else: 
            print("\n\nUnknown Option {}.\n".format(choice))

if __name__ == '__main__':
    main() 

Note that the UI of the previous code is identical to the UI of the code given in Exercise 4.1. Thus the following video still is correct for the above code.

Here is a 2 minute youtube video explaining it all.