import argparse
from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA


def encrypt(text, pub):
    public_key = RSA.importKey(pub)
    public_key = PKCS1_OAEP.new(public_key)
    encrypted_msg = public_key.encrypt(text)
    return encrypted_msg


def decrypt(text, priv):
    private_key = RSA.importKey(priv)
    private_key = PKCS1_OAEP.new(private_key)
    decrypted_msg = private_key.decrypt(text)
    return decrypted_msg


def main():
    # get arguments from command line
    parser = argparse.ArgumentParser()
    parser.add_argument('--msg', type=str)
    args = parser.parse_args()

    # generate public and private keys to use for encryption and decryption
    key = RSA.generate(3072)
    priv = key.export_key('PEM')
    pub = key.publickey().exportKey('PEM')

    # print original message
    original = args.msg
    print("Original message: " + original)

    # print encrypted message
    ciphertext = encrypt(original.encode(), pub)
    print("Encrypted message: ")
    print(ciphertext)

    # print decrypted message
    plaintext = decrypt(ciphertext, priv)
    print("Decrypted message: " + plaintext.decode())


if __name__ == '__main__':
    main()