1
\$\begingroup\$

This is a revised copy of the Mail Crypt library based on the recommendations from the last post, as well as reviewing recommendations from the first post that I misunderstood.
[Last post] Mail Crypt Library for encrypted email
[First Post] Mail Crypt CLI encrypted email wrapper for Gmail

Change Log:

  1. Removed try/except block from tag/sig status so end user gets the error.
  2. Un-bundled context manager so public and private keys are handled separately.
  3. encrypted message return uses join and hex with for loop rather that .hex after every item.
  4. decrypt message now unpack each item into there own var rather than as a list.
  5. import keys now unpacks into 2 var rather than list.
  6. get_emails unpack just what is need into data and rest is pass to underscore.
  7. created fetch function in email to replace 3 identical fetch calls in Email. (get_emails, read_encrypted, read)
  8. Keys.export_keys renamed to Keys.save_keys to have symmetric nomenclature.
  9. Split Keys class into CreateKeys, PublicKeysStore, and PersonalKeys.
  10. rsa_key_length is no longer a class constant but now an arg, change_rsa_key_length class method removed.
  11. split load my keys into get_public and get_private.
  12. removed the private and public key args from MailCrypt so keys must now be passed into encrypt and decrypt function.
  13. static method decorator added to decrypt message.
  14. update keys method removed from MailCrypt since MailCrypt no longer handles keys.
  15. removed constant of aes_session_key_length from MailCrypt, its now an arg passed into encrypt_msg with a defalut of 32.
  16. encrypt message is now a static method.
  17. Email class split into EmailSMTP and EmailIMAP (Grouped based on which server is used).
  18. SMTP port is no longer a constant, can be passed in or default used.
  19. removed change smtp port method since port is now an arg.
  20. fetch now replaces read in EmailIMAP.
  21. PublicKey object now must be passed into EmailSMTP to handle public key look up.
"""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."""
    @staticmethod
    def encrypt_msg(message, recpient_public_key, private_key, aes_session_key_length=32):
        """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(aes_session_key_length)
        aes_cipher = AES.new(session_key, AES.MODE_EAX)
        aes_cipher_text, tag = aes_cipher.encrypt_and_digest(message)
        pub = PKCS1_OAEP.new(recpient_public_key).encrypt(session_key)
        priv = pss.new(private_key).sign(SHA512.new(message))

        return ' '.join(
            part.hex() for part in (
                aes_cipher_text,
                tag,
                aes_cipher.nonce,
                pub,
                priv,
            )
        )

    @staticmethod
    def decrypt_msg(message, sender_public_key, private_key):
        """Splits the message into its sections
            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"""
        aes_cipher_text, tag, nonce, enc_session_key, signature = (
            value.encode().fromhex(value)
            for value in message.split(' '))

        aes_cipher = AES.new(
            PKCS1_OAEP.new(private_key).decrypt(enc_session_key),
            AES.MODE_EAX,
            nonce=nonce,
        )
        clear_text = aes_cipher.decrypt(aes_cipher_text)

        aes_cipher.verify(tag)
        pss.new(sender_public_key).verify(SHA512.new(clear_text), signature)

        return clear_text


class PublicKey:
    """Handles public key storage and retrieval."""
    def __init__(self):
        self.key_dict = {}

    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 save_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):
        """Retrieves public key based on email."""
        return RSA.import_key(self.key_dict[address])

    def remove_key(self, address):
        """Removes key from dict"""
        self.key_dict.pop(address)


class PersonalKeys:
    """Handles users public and private keys."""
    def __init__(self):
        self.my_private_key = None
        self.my_public_key = None

    def get_private(self, passwd):
        """Loads private key in from file."""
        with open('private_key.pem', 'r') as fprivate_key_save:
            self.my_private_key = RSA.import_key(fprivate_key_save.read(), passphrase=passwd)

    def get_public(self):
        """Loads public key in from file."""
        with open('my_public_key.pem', 'r') as fpublic_key_save:
            self.my_public_key = RSA.import_key(fpublic_key_save.read())


class CreateKeys:
    """Handles key pair creation and storage."""
    def __init__(self, rsa_key_length=4096):
        self.rsa_key_length = rsa_key_length

    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:
            fpub.write(public_key.export_key('PEM'))

        with open('private_key.pem', 'wb') as fpri:
            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


class EmailSMTP:
    """Handles the SMTP functionality."""
    def __init__(self, username, passwd, server_address, pub_key_obj, per_key_obj, port=465):
        self.smtp = smtplib.SMTP_SSL(server_address, port)
        self.smtp.ehlo()
        self.username = username
        self.smtp.login(username, passwd)
        self.publickeys = pub_key_obj
        self.private = per_key_obj

    def close_connection(self):
        """Closes connection to server."""
        self.smtp.close()

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

    def send_encrypted(self, recipient, message):
        """Sends encrypted message."""
        message = MailCrypt.encrypt_msg(message.encode(),
                                        self.publickeys.retreive_key(recipient),
                                        self.private.my_private_key)

        self.smtp.sendmail(self.username, recipient, message)

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


class EmailIMAP:
    """Handles the IMAP sever funcionality"""
    def __init__(self, username, passwd, server_address):
        self.imap = imaplib.IMAP4_SSL(server_address)
        self.imap.login(username, passwd)
        self.imap.select('inbox')

    def close_connection(self):
        """Logs out and closes connection to the server"""
        self.imap.logout()
        self.imap.close()

    def get_emails(self):
        """Yeilds uid and senders address for every message
            in the inbox folder."""
        _, data, *_ = self.imap.uid('search', None, 'ALL')
        for uid in data[0].decode().split(' '):
            sender, _ = self.fetch(uid)
            yield uid, sender

    def fetch(self, uid):
        """returns sender address and message payload"""
        _, email_data = self.imap.uid('fetch', uid, '(RFC822)')
        msg = email.message_from_bytes(email_data[0][1])
        return msg['From'], msg.get_payload()

    def read_encrypted(self, uid, sender_public, private):
        """Fetches email from given uid and returns clear text."""
        _, payload = self.fetch(uid)
        return MailCrypt.decrypt_msg(payload, sender_public, private)

    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 uid, sender in self.get_emails():
            sender, msg_body = self.fetch(uid)
            if 'PUBLIC' in msg_body:
                yield sender, msg_body

\$\endgroup\$
  • 1
    \$\begingroup\$ When are you going to add the CLI again? \$\endgroup\$ – Peilonrayz Apr 16 at 15:48
  • \$\begingroup\$ I'm planning on building a CLI for it once this supporting library meets the standard and more or less meets the approval of others in this community. I was also considering releasing this library to be available through pip if there is enough interest in it. \$\endgroup\$ – Joe Smith Apr 16 at 21:04
3
\$\begingroup\$

IMO good design is:

  1. It solves the purpose it sets out to solve.
  2. It only solves the problem you set out to solve.
  3. It doesn't add extra problems for the user.
  4. It doesn't solve extra problems for the user, making it hard to maintain.

MailCrypt

  1. ❌ It is not clear what it solves.
  2. -
  3. ❌ The code is more annoying to use than if the code were just functions.
  4. ❌ The code is more annoying to maintain than if it were just functions.

PublicKey

  1. πŸ’š It stores public keys.
  2. πŸ’š
  3. ❌

    • Save and load keys can be abstracted away.
    • If the program doesn't exit as expected all the new keys are gone.
    • Why can't I specify my own file location?
  4. πŸ’š

PersonalKeys and CreateKeys

  1. 🧑 Together they handle personal keys.
  2. πŸ’š They handle personal keys.
  3. ❌

    • Why do I need two classes for public keys?
    • Why does get_private and get_public not give me what I asked for?
    • Why can I not specify my own file location?
    • Why does generating my keys not give them to me?
    • Why doesn't my_public_key update when I generate new keys?


    Overall, why do I have to jump through loops to interact with my keys?

    Please convert the following code to use your classes to see which is easier to use.

    personal = PersonalKeys(...)
    public = personal.public
    _, private = personal.generate()
    p, _ = personal.generate(save=True)
    assert personal.public == p
    assert personal.public != public
    assert personal.private != private
    personal.save(public, private)
    assert personal.public == public
    assert personal.private == private
    
  4. 🧑 Why does generate_keys not call generate_keys_test?

EmailSMTP and EmailIMAP

  1. ❌ Together they handle emails (collection) and email (items).
  2. ❌ They are handling both the collection and the item.
  3. ❌ The interface is clunky. Contrast with usage of my previous answer.

    mc = MailCrypt()
    for email in mc.emails():
        try:
            data = email.read()
        except <insert decrypt error>:
            email.import_key()
            data = email.read()
    
        if 'delete' in data:
            email.delete()
    
  4. ❌ It's solving two things at once.


The CLI gives your code purpose and shows you the usability problems of your code, helping to find both 1 and 3.

| improve this answer | |
\$\endgroup\$
  • 1
    \$\begingroup\$ The use of Unicode icons is fun. \$\endgroup\$ – Reinderien Apr 17 at 0:11
  • \$\begingroup\$ As always I greatly appreciate your feedback. As for the CLI I have written a bare bones one for each submission for testing and to check it meets functionality needed. Could you clarify the following: Are you suggesting just re-naming the MailCrypt class or removal of it entirely and have the two methods just be function and not be contained in a class? When you suggest the abstraction of save/load keys I'm not sure what you mean. (I know what abstraction means) As for not being able to specify file location, that is on the way. For SMTP/IMAP I grouped them based on what they interact with \$\endgroup\$ – Joe Smith Apr 17 at 0:11
  • \$\begingroup\$ SMTP handling all outgoing mail and IMAP handling all incoming mails. I did not quite understand why in your previous response half of the email functionality was in MailCrypt and the other half was in Email with SMTP and IMAP intermixed. \$\endgroup\$ – Joe Smith Apr 17 at 0:11
  • 1
    \$\begingroup\$ @JoeSmith "just be function and not be contained in a class". Please reread the code I provided for PublicStore in my previous answer - same functionality no save / load. I'm not sure if you want me to reply to the other things, since there seems to be no questions. \$\endgroup\$ – Peilonrayz Apr 17 at 1:48
  • 1
    \$\begingroup\$ @JoeSmith To note whilst I have said that encrypt_msg would be better as just a function. I don't think that's the best solution. My answer isn't a "do this" list this time, instead really think about the points I've raised and see if you'd put it somewhere else. \$\endgroup\$ – Peilonrayz Apr 17 at 1:55
3
\$\begingroup\$

Typo

recpient_public_key -> recipient_public_key

Type hints

Adding some would help for this code to be self-documenting, and to a certain extent, verified by some static analysis tools. For an example, your constructor:

def encrypt_msg(message, recpient_public_key, private_key, aes_session_key_length=32):

could be (I'm guessing a little)

def encrypt_msg(message: str, recpient_public_key: bytes, private_key: bytes, aes_session_key_length: int = 32):

Your other function signatures could similarly benefit, as well as member variables:

    self.key_dict = {}

I see the key is an email address string, and I assume the value is a stringy key:

self.key_dict: Dict[str, str] = {}

Context managers

EmailSMTP, since it needs to close self.smtp, should implement __enter__/__exit__ to do so, and any code you have that calls this should use it in a with. EmailIMAP should do the same.

| improve this answer | |
\$\endgroup\$
  • 1
    \$\begingroup\$ Yep, PEP8 is a little odd that way. It says When combining an argument annotation with a default value, however, do use spaces around the = sign \$\endgroup\$ – Reinderien Apr 16 at 14:18
  • 1
    \$\begingroup\$ I think it's a reasonable exception to the rule :) Otherwise I can see myself misreading some default arguments as part of the type. \$\endgroup\$ – Peilonrayz Apr 16 at 14:21

Your Answer

By clicking β€œPost Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.