Mail Crypt Library for encrypted email

Question

(Last Post) Mail Crypt CLI encrypted email wrapper for Gmail

I've spent the past few days re-writing the logic of my program as a library to improve the structure and readability of my code. The idea behind MailCrypt is for it to be a python based client/library that allows you to locally encrypt messages and send them through Gmail or any SMTP/IMAP server. 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 validation 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.

"""Library for send and receiveing encrypted emails."""

import pickle
import email
import imaplib
import smtplib
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


class MailCrypt:
    """Core compents for encryption/decryption and key generation."""
    aes_session_key_length = 32

    def __init__(self, private_key, public_key):
        """Takes in public key, private key."""
        self.private_key = private_key
        self.public_key = public_key

    def encrypt_msg(self, message, recpient_public_key):
        """Generates a session key to use with AES to encrypt the message, then
            encrypts the session key with the recipients public key. Everything
            is returned in hex format to be better sent over email."""
        session_key = get_random_bytes(self.aes_session_key_length)
        aes_cipher = AES.new(session_key, 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(recpient_public_key).encrypt(session_key).hex()
                + ' ' + pss.new(self.private_key).sign(SHA512.new(message)).hex()
        )

    def decrypt_msg(self, message, sender_public_key):
        """Splits the message into its sections
            Indicies:
            0 aes_cipher_text
            1 tag
            2 nonce
            3 enc_session_key
            4 signature

            Decrypts the session key, then decrypts the message body.
            If aes_cipher.verify throws an error that means an invalid
                tag was provided
            If pss.new throws an error that means the message signature
                is invalid"""
        seperated_msg = [
            value.encode().fromhex(value)
            for value in message.split(' ')
        ]
        aes_cipher = AES.new(
            PKCS1_OAEP.new(self.private_key).decrypt(seperated_msg[3]),
            AES.MODE_EAX,
            nonce=seperated_msg[2],
        )
        clear_text = aes_cipher.decrypt(seperated_msg[0])

        try:
            aes_cipher.verify(seperated_msg[1])
            tag_status = True
        except ValueError:
            tag_status = False

        try:
            pss.new(sender_public_key).verify(SHA512.new(clear_text), seperated_msg[4])
            sig_status = True
        except ValueError:
            sig_status = False

        return tag_status, sig_status, clear_text

    def update_keys(self, private_key, public_key):
        """Updates the keys in the mailcrypt object if new keys are generated"""
        self.public_key = public_key
        self.private_key = private_key


class Keys:
    """Handles key pair creation and storage."""
    rsa_key_length = 4096
    my_private_key = None
    my_public_key =None

    def __init__(self):
        self.key_dict = {}

    def load_my_keys(self, passwd):
        """Loads saved keys in from file."""
        with open('private_key.pem', 'r') as fprivate_key_save, \
                open('my_public_key.pem', 'r') as fpublic_key_save:
            self.my_public_key = RSA.import_key(fpublic_key_save.read())
            self.my_private_key = RSA.import_key(fprivate_key_save.read(), passphrase=passwd)

    def load_keys(self):
        """Read public keys in from file. Must be in same folder
            as script is run from."""
        with open('public_key_bank.pkl', 'rb') as fpublic_key_file:
            self.key_dict = pickle.load(fpublic_key_file)

    def export_keys(self):
        """Saves key_dict to file."""
        with open('public_key_bank.pkl', 'wb') as fpublic_key_file:
            pickle.dump(self.key_dict, fpublic_key_file)

    def add_key(self, address, key):
        """Adds key to to key_dict."""
        self.key_dict[address] = key

    def retreive_key(self, address):
        """Retrievs public key based on email."""
        return RSA.import_key(self.key_dict[address])

    def generate_keys(self, passwd):
        """Generates public and private key pairs and exports them as .pem files."""
        private_key = RSA.generate(self.rsa_key_length)
        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=passwd))

    def generate_keys_test(self):
        """"Used for testing, returns key pair."""
        private_key = RSA.generate(self.rsa_key_length)
        public_key = private_key.publickey()
        return private_key, public_key

    @classmethod
    def change_rsa_key_length(cls, length):
        """Changes the key length for the key pair generation
            must be powers of 256"""
        cls.rsa_key_length = length

    @classmethod
    def change_aes_key_length(cls, length):
        """Changes the AES session key length must be 8, 16, or 32"""
        cls.aes_session_key_length = length


class Email:
    """Handles all of the imap and smtp connections and functionality"""
    smtp_port = 465

    def __init__(self, smtp_server_address, imap_server_address, username, passwd, mailcrypt_instance, keys_instance):
        """Opens a connections to specified imap and smtp servers, logs in with the
            given username and passwords and navagaits to the inbox folder."""
        self.smtp = smtplib.SMTP_SSL(smtp_server_address, self.smtp_port)
        self.imap = imaplib.IMAP4_SSL(imap_server_address)
        self.username = username
        self.passwd = passwd
        self.smtp.ehlo()
        self.smtp.login(username, passwd)
        self.imap.login(username, passwd)
        self.mailcrypt = mailcrypt_instance
        self.imap.select('inbox')
        self.keys = keys_instance

    @classmethod
    def change_smtp_port(cls, port):
        """Allows you to change the port for the smtp connection if
            useing a non-standard configuraation."""
        cls.smtp_port = port

    def send(self, recipient, message):
        """Sends plain text email."""
        self.smtp.sendmail(self.username, recipient, message)

    def read(self, uid):
        """Fetches plain text email based on uid and returns the
            message body."""
        _, email_data = self.imap.uid('fetch', uid, '(RFC822)')
        msg = email.message_from_bytes(email_data[0][1])
        return msg.get_payload()

    def send_encrypted(self, recipient, message):
        """Sends encrypted message."""
        message = self.mailcrypt.encrypt_msg(message.encode(), self.keys.retreive_key(recipient))
        self.smtp.sendmail(self.username, recipient, message)

    def read_encrypted(self, uid):
        """Fetches email from given uid and returns clear text."""
        _, email_data = self.imap.uid('fetch', uid, '(RFC822)')
        msg = email.message_from_bytes(email_data[0][1])
        return self.mailcrypt.decrypt_msg(msg.get_payload(), self.keys.retreive_key(msg['From']))

    def mark_delete(self, uid):
        """Moves the specified email to trash folder. If useing email
            provider other than gmail 'Trash' needs to be changed to
            whatever folder that service uses."""
        self.imap.uid('store', uid, '+X-GM-LABELS', '(\\Trash)')

    def delete_all(self):
        """Empties the trash folder."""
        self.imap.expunge()

    def import_key(self):
        """Checks message payloads for public keys, if found it
            yeilds the senders email address and the public key."""
        for message in self.get_emails():
            msg_body = self.read(message[0])
            if 'PUBLIC' in msg_body:
                yield message[1], msg_body

    def get_emails(self):
        """Yeilds tuple with uid and senders address for every message
            in the inbox folder."""
        _, data = self.imap.uid('search', None, 'ALL')
        for uid in data[0].split():
            _, email_data = self.imap.uid('fetch', uid, '(RFC822)')
            msg = email.message_from_bytes(email_data[0][1])
            yield uid.decode(), msg['From']

    def share_public_key(self, recipient):
        """Sends public key."""
        self.send(recipient, self.mailcrypt.public_key.export_key('PEM'))

Answer 1

Typo

seperated_msg -> separated_msg

Retrievs -> Retrieves

Hex-able join

This:

    return (
            aes_cipher_text.hex()
            + ' ' + tag.hex()
            + ' ' + aes_cipher.nonce.hex()
            + ' ' + PKCS1_OAEP.new(recpient_public_key).encrypt(session_key).hex()
            + ' ' + pss.new(self.private_key).sign(SHA512.new(message)).hex()
    )

can be converted into a join, since everything there has hex called on it:

pub = PKCS1_OAEP.new(recpient_public_key).encrypt(session_key)
priv = pss.new(self.private_key).sign(SHA512.new(message))
return ' '.join(
    part.hex() for part in (
        aes_cipher_text,
        tag,
        aes_cipher.nonce,
        pub,
        priv,
    )
)

I also think that pub and priv should be separated due to their length.

Unpacking

On the other end, in decrypt_msg, you should not be using array indices into seperated_msg. Instead,

aes_cipher_text, tag, none, enc_session_key, signature = (
    value.encode().fromhex(value)
    for value in message.split(' ')
)

Somewhat similarly, this:

    for message in self.get_emails():
        msg_body = self.read(message[0])
        if 'PUBLIC' in msg_body:
            yield message[1], msg_body

should unpack whatever message is into two different iteration variables.

Boolean returns

tag_status and sig_status are not very Pythonic ways of capturing error information. If it's important to separate failure types in this way, rather than returning two different booleans, throw (potentially) two different custom exception types that throw from e (rethrow) ValueError.

Statics

These three have problems:

class Keys:
    rsa_key_length = 4096
    my_private_key = None
    my_public_key =None

The first should be capitalized since it's a class constant. The second and third probably should not be statics, and should be instance variables instead.

Symmetric nomenclature

Rather than load_keys / export_keys, consider either load_keys/save_keys or export_keys/import_keys.

Bundled context manager

This:

    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=passwd))

should be separated into two different with statements. Those two files are independent and should not share a context scope.