import argparse
from Cryptodome.Cipher import AES
from Cryptodome.Random import get_random_bytes


def encrypt(text, key):
    # make encryptor
    encryptor = AES.new(key, AES.MODE_CFB)

    # takes the data and returns ciphertext
    cipher_text = encryptor.encrypt(text)

    # now to write the ciphertext to the file
    encrypted_file = open("encrypted_file.txt", "wb")
    encrypted_file.write(cipher_text)
    encrypted_file.close()

    # return the cipher_text to show users
    return cipher_text


def decrypt(text, key):
    # make decrypter
    decrypter = AES.new(key, AES.MODE_CFB)
    plain_text = decrypter.decrypt(text)
    return plain_text


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

    # make random 32B long key
    key = get_random_bytes(32)

    # read the text file in and make it a string with no line breaks
    text = args.file.read().replace('\n', '')

    # print encrypted text
    print("Original file contents:")
    print(text)

    # encrypt the file
    encrypted_text = encrypt(str.encode(text), key)

    # print encrypted text
    print("\nEncrypted file created. Contents:")
    print(encrypted_text)

    # decrypt and print decrypted text
    print("\nDecrypted file. Contents:")
    print(decrypt(encrypted_text, key))


if __name__ == '__main__':
    main()