Mail Crypt CLI encrypted email wrapper for Gmail

Question

I've spent the past few days writing a python based client that allows you to locally encrypt your message and send it through Gmail. The encryption scheme is based off of hybrid RSA/AES encryption, where the message body is encrypted with AES and the one time session key is RSA PKCS1 OEAP. It also provides valdation of the sender by hashing and signing the hash with your private key. For this to run it needs pycryptodome but it's imported as Crypto, as well as less secure connections need to be enabled on the Gmail account. Any feedback would be greatly appreciated.

"""This is a wrapper for gmail to be able to send emails encrypted locally with
   256 bit AES and 4096 bit RSA hybrid encryption."""

import email
import imaplib
import pickle
import smtplib
import socket
from getpass import getpass
from Crypto.Hash import SHA512
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Cipher import AES
from Crypto.PublicKey import RSA
from Crypto.Random import get_random_bytes
from Crypto.Signature import pss


# Needs pycryptodome


def logo():
    """Opening screen prompt"""

    print('    __  ___      _ __   ______                 __')
    print('   /  |/  /___ _(_) /  / ____/______  ______  / /_')
    print('  / /|_/ / __ `/ / /  / /   / ___/ / / / __ \/ __/')
    print(' / /  / / /_/ / / /  / /___/ /  / /_/ / /_/ / /_  ')
    print('/_/  /_/\__,_/_/_/   \____/_/   \__, / .___/\__/')
    print('                               /____/_/          ')


def usr_help():
    """Displays help message."""
    print('Getting Started:\n'
          '\n'
          '     -First you must generate your keys.\n'
          '     -After that you must share your public key with who ever\n'
          '         you intend to communicate with, they must also share their\n'
          '         public key with you.\n'
          '\n'
          '(1) Send email - This allows you to send emails to people you have\n '
          '                 shared and recieved public keys from.\n'
          '\n'
          '(2) Check email - This loads your inbox and allows you to select a message\n'
          '                     to read\n'
          '\n'
          '(3) Delete emails - This loads your inbox and allows you to select emails\n'
          '                     to delete\n'
          '\n'
          '(4) Import public key - This loads any public key sent to your inbox.\n'
          '\n'
          '(5) Share public key - This sends your public key to someones email to \n'
          '                         to be imported and used\n'
          '\n'
          '(6) Generate key pair - This generates your public and private key pair.\n'
          '                         This must be done when your first use mail crypt.\n'
          '                         If you generate a new key pair it will overwrite \n'
          '                         your existing key pair, so you will not be able to\n'
          '                         read old emails and you will have to re-share your\n'
          '                         public key\n'
          '\n'
          '(7) Help - Displays this message.\n'
          '\n'
          '(00) Exit - This closes connectins, logs out, and ends mail cyrpt.\n'
          '\n')
    input("Press any key to return to the menu.")


def splash():
    """front page gets Gmail login info
        username and password is passed in as string and returned to the main loop"""

    print('\n' * 8)
    print("Welcome to MailCrypt the secure Gmail wrapper.")
    print('Connecting...')
    username = input("Enter email address to use.\n")
    passwd = getpass()

    return username, passwd


def load_inbox(imap_server_instance):
    """loads emails from inbox. Opens a connection to the Gmail IMAP server, should be able to swap
        that with any email provider that has an IMAP sever that can interface with the imaplib
        module. By default it loads only the inbox folder.

        Calls the decrypt_msg function to decrypt the cyher text email"""
    try:
        imap_server_instance.select('inbox')
        _, data = imap_server_instance.uid('search', None, 'ALL')
        data = data[0].split()

        for item in data:
            # Loops through all emails in inbox and displays there UID and who sent it.
            _, email_data = imap_server_instance.uid('fetch', item, '(RFC822)')
            raw = email_data[0][1]
            msg = email.message_from_bytes(raw)
            print('UID:', item.decode(), end=" From: ")
            print(msg['From'])

        msg_to_fetch = input('What message do you want to read? (UID)\n')
        _, email_data = imap_server_instance.uid('fetch', msg_to_fetch, '(RFC822)')
        raw = email_data[0][1]
        msg = email.message_from_bytes(raw)
        # The call to decrypt_msg function
        clear_msg, authentic, sig_stat = decrypt_msg(msg.get_payload(), msg['From'])

        if authentic:
            print('Message tag valid')
        else:
            print('Message tag INVALID!')
        if sig_stat:
            print('Message signature is valid.')
        else:
            print('Message signature INVALID!')

        print(clear_msg)
        input('Press any key to return to the menu.')

    except imaplib.IMAP4.error:
        print("Error please try again. (005)")


def delete_email(imap_server_instance):
    """Moves email to deleted folder and deletes."""
    uids = []
    try:
        imap_server_instance.select('inbox')
        _, data = imap_server_instance.uid('search', None, 'ALL')
        data = data[0].split()

        for item in data:
            # Loops through all emails in inbox and displays there UID and who sent it.
            _, email_data = imap_server_instance.uid('fetch', item, '(RFC822)')
            raw = email_data[0][1]
            msg = email.message_from_bytes(raw)
            uids.append(item.decode())
            print('UID:', item.decode(), end=" From: ")
            print(msg['From'])

    except imaplib.IMAP4.error:
        print("Error please try again. (005)")

    while True:
        email_to_delete = input("What email would you like to delete?(UID)\n Enter 00 to exit.\n")
        if email_to_delete in uids:
            imap_server_instance.uid('store', email_to_delete, '+X-GM-LABELS', '(\\Trash)')
        elif email_to_delete == '00':
            imap_server_instance.expunge()
            break
        else:
            print('Invalid choice!')


def share_public_key(username, smtp_server_instance):
    """Shares your curently stored public key"""

    recipient = input('Enter email to send public key to.\n')

    with open('my_public_key.pem', 'r') as public_key_save:
        public_key_string = public_key_save.read()

    try:
        smtp_server_instance.sendmail(username, recipient, public_key_string)
        print('email sent')

    except smtplib.SMTPRecipientsRefused:
        print("Could not reach recipients. (003)")

    except smtplib.SMTPDataError:
        print("Unknown error. (004)")


def send_email():
    """Generates the email but doesent send

    Calls the encrypt_msg function"""

    print('Enter destination email address.')
    recipient = input()
    print('Type your message here, press enter to send.')
    msg_body = input()

    enc_msg_body = encrypt_msg(msg_body.encode(), recipient)
    return enc_msg_body, recipient


def import_public_key(imap_server_instance):
    """Imports public key from inbox by logging in and scaning all emails for
        the key word 'PUBLIC', newer keys found will over write the older keys in the
        public key dictionary and file

        Calls the public_key_store function"""

    print('Updating public keys.')
    try:
        imap_server_instance.select('inbox')
        _, data = imap_server_instance.uid('search', None, 'ALL')
        data = data[0].decode().split()

        for item in data:
            _, email_data = imap_server_instance.uid('fetch', item, '(RFC822)')
            raw = email_data[0][1].decode()
            msg = email.message_from_string(raw)
            if 'PUBLIC' in msg.get_payload():
                public_key_store(msg['From'], msg.get_payload())

    except imaplib.IMAP4.error:
        print('Unexpected error. (005)')


def send(username, recipient, msg_body, smtp_server_instance):
    """Opens connection to gmail smtp server and sends the email"""

    try:
        smtp_server_instance.sendmail(username, recipient, msg_body)
        print('email sent')

    except smtplib.SMTPRecipientsRefused:
        print("Could not reach recipients. (003)")

    except (smtplib.SMTPDataError, smtplib.SMTPSenderRefused):
        print("Unknown error. (004)")


def generate_keys():
    """This uses the pycrypt RSA module to create private and public keys, the
    RSA.generate key length can be changed for more or less secuity. (must be expont of 256)

        Calls public_key_store func"""
    while True:
        if input('Generating a new set of keys will permainatly delete your old ones.\n '
                 'Enter y to continue\n').lower() != 'y':
            print('Exiting.')
            break
        print('Generating new key pair.')
        private_key = RSA.generate(4096)
        public_key = private_key.publickey()

        with open('my_public_key.pem', 'wb') as public_key_save, \
                open('private_key.pem', 'wb') as private_key_save:
            public_key_save.write(public_key.export_key('PEM'))
            private_key_save.write(private_key.export_key('PEM',
                                                          passphrase=getpass('Enter password to secure private key.\n')))

        return None


def encrypt_msg(msg_body, recipient):
    """uses public key retrieve function to pull public key out of saved keys"""
    publickey = public_key_retrieve(recipient)
    session_key = get_random_bytes(32)

    if publickey:
        aes_cipher = AES.new(session_key, AES.MODE_EAX)
        nonce = aes_cipher.nonce
        aes_cipher_text, tag = aes_cipher.encrypt_and_digest(msg_body)

        rsa_cipher = PKCS1_OAEP.new(publickey)
        enc_session_key = rsa_cipher.encrypt(session_key)

        try:
            with open('private_key.pem', 'r') as private_key_save:
                privatekey = RSA.import_key(private_key_save.read(),
                                            passphrase=getpass('\nEnter password to unlock private key.'))

                msg_hash = SHA512.new(msg_body)
                signature = pss.new(privatekey).sign(msg_hash)

        except ValueError:
            return 'Invalid private key or password.\n'

        full_msg = aes_cipher_text.hex() + ' ' + tag.hex() + ' ' +  nonce.hex() + ' ' +  enc_session_key.hex()  + ' ' +  signature.hex()
        return full_msg

    print('No public key for that email is stored.')
    return None


def decrypt_msg(full_msg, sender_publickey):
    """decrypts message imported from your inbox
       index of full_msg components in split list
       0 aes_cipher_text
       1 tag
       2 nonce
       3 enc_session_key
       4 signature
       """

    seperated_msg = full_msg.split(' ')

    # Encodes, and converts items in list from hex to bytes
    for item in enumerate(seperated_msg):
        seperated_msg[item[0]] = seperated_msg[item[0]].encode().fromhex(seperated_msg[item[0]])

    try:
        with open('private_key.pem', 'r') as private_key_save:
            privatekey = RSA.import_key(private_key_save.read(),
                                        passphrase=getpass('Enter password to unlock private key.'))
            decrypt = PKCS1_OAEP.new(privatekey)
            psk = decrypt.decrypt(seperated_msg[3])

    except ValueError:
        return 'Invalid private key or password.\n', False, False

    aes_cipher = AES.new(psk, AES.MODE_EAX, nonce=seperated_msg[2])
    clear_text = aes_cipher.decrypt(seperated_msg[0])

    try:
        aes_cipher.verify(seperated_msg[1])
        authentic = True

    except ValueError:
        authentic = False

    msg_hash = SHA512.new(clear_text)
    valid_sig = pss.new(public_key_retrieve(sender_publickey))

    try:
        valid_sig.verify(msg_hash, seperated_msg[4])
        sig_stat = True

    except ValueError:
        sig_stat = False

    return clear_text, authentic, sig_stat


def public_key_store(email_address, new_public_key):
    """Handles storeing and updating stored public keys"""
    try:
        # Tries to open existing file
        with open('public_key_bank.pkl', 'rb') as public_key_file:
            public_key_bank = pickle.load(public_key_file)

    except FileNotFoundError:
        public_key_bank = {}

    public_key_bank[email_address] = new_public_key

    with open('public_key_bank.pkl', 'wb') as public_key_file:
        pickle.dump(public_key_bank, public_key_file)


def public_key_retrieve(email_address):
    """retreives public key from saved"""
    try:
        with open('public_key_bank.pkl', 'rb') as public_key_file:
            public_key_bank = pickle.load(public_key_file)
            if email_address in public_key_bank:
                requested_key = public_key_bank[email_address]
                requested_key = RSA.import_key(requested_key)

                return requested_key

            return False

    except FileNotFoundError:
        print('Public key file has not been created. (006)')


print('\n' * 100)
logo()
try:
    with smtplib.SMTP_SSL('smtp.gmail.com', 465) as smtp_server, \
            imaplib.IMAP4_SSL('imap.gmail.com') as imap_server:
        while True:
            USERNAME, PASSWD = splash()
            try:
                smtp_server.ehlo()
                smtp_server.login(USERNAME, PASSWD)
                imap_server.login(USERNAME, PASSWD)

                break

            except smtplib.SMTPAuthenticationError:
                print('Incorrect username or password! (001)')
            except imaplib.IMAP4.error:
                print('Incorrect username or password! (002)')
            except (OSError, smtplib.SMTPServerDisconnected, imaplib.IMAP4_SSL.error):
                print('Conection error. (008)')
                break

        while True:
            print('\n' * 100)
            logo()
            print('\n' * 5)
            OPTION = input('What would you like to do\n'
                           '(1) Send email\n'
                           '(2) Check email\n'
                           '(3) Delete email\n'
                           '(4) Import public key\n'
                           '(5) Share public key\n'
                           '(6) Generate key pair.\n'
                           '(7) Help\n'
                           '(00) Exit\n')
            print('\n' * 100)

            if OPTION == '1':
                # Send email
                CIPHER_TEXT, DESTINATION_EMAIL = send_email()
                send(USERNAME, DESTINATION_EMAIL, CIPHER_TEXT, smtp_server)

            elif OPTION == '2':
                # Check email
                load_inbox(imap_server)

            elif OPTION == '3':
                # Delete email
                delete_email(imap_server)

            elif OPTION == '4':
                # Import public key
                import_public_key(imap_server)

            elif OPTION == '5':
                # Share public key
                share_public_key(USERNAME, smtp_server)

            elif OPTION == '6':
                # Generate key pair
                generate_keys()

            elif OPTION == '7':
                usr_help()

            elif OPTION == '00':
                # End program loop and logs out of smtp/imap servers
                try:
                    smtp_server.quit()
                    imap_server.close()
                    imap_server.logout()
                except imaplib.IMAP4.error:
                    pass
                PASSWD = '0000000000000000000000000'
                break

            else:
                print('Invalid option!')

except socket.gaierror:
    print('Please check your network connection and try again. (007)')
except KeyboardInterrupt:
    pass

Answer 1

You should separate the UI from the actual logic of the program.
Whilst a nice UI is cool I believe that it's impeding your ability to structure your code properly. And so I'd recommend making Mail Crypt a Python library first, and a CLI second.

Since this is a library we want:

• Any errors raised to pass through to the user.
• No print statements.
• All values to be provided to the method. (not inside the method)

Additionally since load_inbox, delete_email, and import_public_key require the same code we can make an Email class to interact with them.

This can get the following code:

class MailCrypt:
    def __init__(self, smtp, imap):
        self.smtp = smtp
        self.imap = imap

    def send(self, account, recipient, message):
        message = encrypt_msg(message.encode(), recipient)
        self.smtp.sendmail(account, recipient, message)

    def emails(self):
        self.imap.select('inbox')
        uids = self.imap.uid('search', None, 'ALL')
        for uid in data[1][0].split():
            email = self.imap.uid('fetch', uid, '(RFC822)')
            yield Email(
                self.imap,
                email.message_from_bytes(email[1][0][1]),
                uid.decode(),
            )

    def delete_all(self):
        self.imap.expunge()

    def share_public_key(self, account, recipient):
        with open('my_public_key.pem', 'r') as f:
            public_key = f.read()
        smtp_server_instance.sendmail(account, recipient, public_key)

    def generate_keys(self, passphrase):
        private_key = RSA.generate(4096)
        public_key = private_key.publickey()

        with open('my_public_key.pem', 'wb') as fpub, \
                open('private_key.pem', 'wb') as fpri:
            fpub.write(public_key.export_key('PEM'))
            fpri.write(private_key.export_key('PEM', passphrase=passphrase))


class Email:
    def __init__(self, imap, email, uid):
        self.email = email
        self.uid = uid

    def read(self):
        return decrypt_msg(self.email.get_payload(), self.email['From'])

    def delete(self):
        self.imap.uid('store', self.uid, '+X-GM-LABELS', '(\\Trash)')

    def import_key(self):
        payload = self.email.get_payload()
        if 'PUBLIC' in payload:
            public_key_store(msg['From'], payload)

This is, I hope, clearly much easier to read and work with than the code you have.

This doesn't cover all of the functionality that you had before. For example, I've only written the code to import the key from one email. I have also not written the select an email code that you had duplicated over multiple functions.

However, to implement these is really quite simple:

def import_all(mc):
    for email in mc.emails():
        email.import_key()


def select_email(mc):
    emails = {e.uid: e for e in mc.emails()}
    for email in emails.values():
        print('UID: {} FROM: {}'.format(email.uid, email.email['From']))
    uid = input('What email would you like? ')
    return emails[uid]

The public key store should:

• Be a class.
• Store the keys in memory.
• Raise an error if the key you're asking for is not in the store.

class PublicStore:
    def __init__(self, location):
        self._location = location
        self._keys = self._load()

    def _load(self):
        try:
            with open(self._location, 'rb') as f:
                return pickle.load(f)
        except FileNotFoundError:
            return {}

    def store(email, key):
        self._keys[email] = key
        with open(self._location, 'wb') as f:
            pickle.dump(self._keys, f)

    def get(email):
        return RSA.import_key(self._keys[email])

Your encryption and decryption can be improved by:

• Moving the personal key into a class.

• Getting the personal and sender keys outside of the encrypt or decrypt functions.

• Removing a lot of unneeded variables.

• You can use a comprehension to build the seperated_msg.

• You should let errors propagate. This is not the correct place to handle them.

  If you're not a fan of this then you can alternately create another class and handle the verification in one or two of its methods.

class PersonalKey:
    def __init__(self, public_location, private_location):
        self._public_location = public_location
        self._private_location = private_location

    def get_public(self, passphrase):
        with open(self._public_location, 'r') as f:
            return RSA.import_key(f.read(), passphrase=passphrase)

    def get_private(self, passphrase):
        with open(self._private_location, 'r') as f:
            return RSA.import_key(f.read(), passphrase=passphrase)


def encrypt_msg(message, sender_key, personal_key):
    aes_cipher = AES.new(get_random_bytes(32), AES.MODE_EAX)
    aes_cipher_text, tag = aes_cipher.encrypt_and_digest(message)
    return (
        aes_cipher_text.hex()
        + ' ' + tag.hex()
        + ' ' + aes_cipher.nonce.hex()
        + ' ' + PKCS1_OAEP.new(sender_key).encrypt(session_key).hex()
        + ' ' + pss.new(personal_key).sign(SHA512.new(message)).hex()
    )


def decrypt_msg(message, sender_key, personal_key):
    seperated_msg = [
        value.encode().fromhex(value)
        for value in message.split(' ')
    ]
    aes_cipher = AES.new(
        PKCS1_OAEP.new(personal_key).decrypt(seperated_msg[3]),
        AES.MODE_EAX,
        nonce=seperated_msg[2],
    )
    clear_text = aes_cipher.decrypt(seperated_msg[0])
    aes_cipher.verify(seperated_msg[1])
    pss.new(sender_key).verify(SHA512.new(clear_text), seperated_msg[4])
    return clear_text

Unfortunately, the answer has consumed enough of my time. The above code may not work, and the changes I made in each section conflict with each other.

Please can you take the above answer and complete the final steps to get the code to work together. When you're doing this with the above functions you are not allowed to use print, getpass or input in the above. However, when effectively rebuilding the CLI you can use these functions to your heart's content.

Once you've made your changes please come back and post another question :)