Creating a class and calling it from other .py files

Question

My objective with this code was to create a class and call it from other .py files. The code itself is working, but I would like to know what could be improved or be better implemented.

I wrote another program just to call this class and its methods and it worked. I'm using functions that will only work on Windows because this is an exercise, not production code. My code has only one class (StudingClassFirstTime) and 13 methods. It is using the external library PyNaCl.

import getpass # Allow user input a password without echo it on terminal
import os # To be able to use os.system('cls')
from nacl import pwhash, secret, utils, exceptions # Encrypt functions to be used


class StudingClassFirstTime:

    # Constructor mode: Here I will define the KDF (argon2id) parameters that will be used by default
    def __init__(self, kdf=pwhash.argon2id.kdf,
                 ops=pwhash.argon2id.OPSLIMIT_SENSITIVE,
                 mem=pwhash.argon2id.MEMLIMIT_SENSITIVE,
                 size=secret.SecretBox.KEY_SIZE):
        self.kdf = pwhash.argon2id.kdf
        self.ops = pwhash.argon2id.OPSLIMIT_SENSITIVE
        self.mem = pwhash.argon2id.MEMLIMIT_SENSITIVE
        self.size = secret.SecretBox.KEY_SIZE

    # Load an encryption key stored on file
    # This file should have only one encryption key inside it
    # If you have one file.key, you can encrypt as many files as you want, but, all of them will use the very same key
    # If you encrypt the same file twice with the same key, the output will not be equal thanks to NONCE
    def load_key(self, key_file):
        with open(key_file, 'rb') as my_key:
            key_of_file = my_key.read()
        return key_of_file

    # Generate an encryption key and return it's value(key)
    def generate_key(self):
        self.salt = utils.random(pwhash.argon2id.SALTBYTES) # 16 bytes == 128 bits
        password = self.password() # Ask a user to input a password
        key = self.kdf(self.size, password, self.salt, opslimit=self.ops, memlimit=self.mem)
        return key # Key have 32 bytes == 256 bits

    # Generate a raw 256 bits (32 bytes) as a key
    # Do not need any user password, or KDF function
    def generate_key_file(self):
        password = utils.random(32) # 32 bytes == 256 bits
        file_key = secret.SecretBox(password).encode()
        return file_key

    # Generate a encryption key and write it to a file
    def write_key_file(self, key_file):
        key = self.generate_key_file()
        with open(key_file, 'ab') as my_key:
            if os.path.getsize(key_file) > 0: # Check if the file have more than 0 bytes
                my_key.write(b'\n' + key_file)   # If positive, first add a new line, then, write the key
            else:                             # If negative, just write the key, without add new line
                my_key.write(key_file)

    # Encrypt a message, generating an encryption key based on password.
    # The key will not be write to file.
    def encrypt(self, message):
        encrypt_key = self.generate_key()
        encrypt_box = secret.SecretBox(encrypt_key)
        encrypted = encrypt_box.encrypt(message)
        assert len(encrypted) == len(message) + encrypt_box.NONCE_SIZE + encrypt_box.MACBYTES  # NONCE_SIZE == 24 bytes // MACBYTES == 16 bytes
        return self.salt + encrypted

    # Encrypt a message with a key stored on file
    # To use it, first you should use generate_key_file() or already have a key contains your key.
    def encrypt_with_key(self, key_file, message):
        encrypt_key = self.load_key(key_file) # Load the encryption key stored on file
        encrypt_box = secret.SecretBox(encrypt_key)
        encrypted = encrypt_box.encrypt(message)
        assert len(encrypted) == len(message) + encrypt_box.NONCE_SIZE + encrypt_box.MACBYTES
        return encrypted

    # Decrypt a message, using a user supplied password
    def decrypt(self, message):
        try:
            password = getpass.getpass(prompt="\nType password to DECRYPT: ", stream=None).encode() # Ask a user for the password to decrypt the message
            salt = message[:16] # The first 16 bytes of message [0 included : 16 not included] is the salt value
            decrypt_key = self.kdf(secret.SecretBox.KEY_SIZE, password, salt, opslimit=self.ops, memlimit=self.mem) # Derivate the enc. key from password
            decrypt_box = secret.SecretBox(decrypt_key)
            decrypted = decrypt_box.decrypt(message[16:])
            assert len(decrypted) == len(message[16:]) - decrypt_box.NONCE_SIZE - decrypt_box.MACBYTES
            return decrypted

        except exceptions.CryptoError:
            return 'error'

    # Decrypt a message with a key stored on file
    # To use it, first you should use encrypt_with_key() method.
    def decrypt_with_key(self, key_file, message):
        try:
            decrypt_key = self.load_key(key_file)
            decrypt_box = secret.SecretBox(decrypt_key)
            decrypted = decrypt_box.decrypt(message)
            assert len(decrypted) == len(message) - decrypt_box.NONCE_SIZE - decrypt_box.MACBYTES
            return decrypted

        except exceptions.CryptoError:
            return 'error'

    # Ask a user to type a password and verify it and return the password
    def password(self):
        while True:
            password1 = getpass.getpass(prompt="\nType your password: ", stream=None).encode() # getpass do not echo the password on terminal
            password2 = getpass.getpass(prompt="Type your password again: ", stream=None).encode()
            if password1 != password2:
                input("\nPassword does not match! Press any key to try again.")
                os.system('cls')
                continue
            else:
                return password1

    # Similar as write_key_file(), but, will not write the key on file.
    # This will be done by the function that call it.
    # The objective here is allow the caller (function) to store multiple keys in one single file, instead of, one file.key per encrypted file
    def key_db(self):
        key = self.generate_key_file()
        return key

    # Similar to encrypt_with_key(), but the key will be supplied.
    def encrypt_db(self, encrypt_key_db, message):
        encrypt_box = secret.SecretBox(encrypt_key_db)
        encrypted = encrypt_box.encrypt(message)
        assert len(encrypted) == len(message) + encrypt_box.NONCE_SIZE + encrypt_box.MACBYTES
        return salt + encrypted

    # The inverse of encrypt_db()
    def decrypt_db(self, key_db, message):
        try:
            decrypt_box = secret.SecretBox(bytes.fromhex(key_db))
            decrypted = decrypt_box.decrypt(message)
            assert len(decrypted) == len(message) - decrypt_box.NONCE_SIZE - decrypt_box.MACBYTES
            return decrypted

        except exceptions.CryptoError:
            return 'error'

'''
The first 56 bytes of encrypted text are:
[0:16] = SALT (16) # Will be present only if the user supply a password
[16:40] = NONCE (24) # Always present
[40:56] = MACBYTES (16) # Always present
[56:] = Encrypted Text
'''

Answer 1

First lets discuss the class and method naming:

StudingClassFirstTime First of all, the word Studing is not in my vocabulary. But please just name classes after what they are supposed to do, and don't include the word class; we know it is a class.

• load_key: this loads a 32 byte key from file, which is fine, however it doesn't include file in the identifier as in later methods;
• generate_key: this derives a key from a password; this is not the same thing as generating a key, which in cryptography generally means taking bytes from a random number generator;
• generate_key_file: here a file is suggested, however no file is being produced;
• write_key_file: here a key file (or, as talked about before, the contents) is generated and encrypted first, which is not in the method description;
• encrypt: OK, fine, here you perform the encryption which always uses a key, however the key parameter is missing - you just assume that the reader expects that the key is generated;
• encrypt_with_key: one would suggest that this is a pleonasm - a key is required to encrypt anything;
• password: password is not a verb - so it should be named ask_user_for_password for instance.
• key_db: similar problem as above, no verb - the extensive commenting required shows that this is not a well thought out name;
• encrypt_db: this begs the question: what DB, did we define a DB somewhere?
• decrypt_db: OK, but the comment here isn't necessary, everybody should know that decrypt is the inverse of encrypt.

For write_key_file it is not symmetric with load_key which is generally what would be expected.

For encrypt_db, the DB is not present in the method call, instead there is a parameter message, which we have to guess is the DB. Moreover, encrypt_key_db to me suggests that there is a database of keys, not that this represents the key to encrypt the DB.

The only difference between the the encrypt_db and the encrypt_with_key seems to be that the key is in a file or not. That is not reflected by the API.

---

Now lets talk about the way the program/class is structured. It seems to be doing a lot of things at the same time. For one, it mixes file IO, user input, clearing of the terminal as well as the "business logic" of performing encryption / decryption and generation of the keys from passwords.

Generally we try and separate these three as much as possible. What you certainly now want is to make your business logic UI dependent. This is exactly what the class does when calling getpass.getpass from the business logic. This could be called "reversal of control". Generally there is a controller class that performs the business logic, which is in turn controlled by the UI class. At most the controller class allows the UI class to insert callback functionality (e.g. for updating a progress bar). The UI class or main method would instead configure the password through a method or constructor call. That makes the code reusable for e.g. GUI applications instead of console applications.

---

Finding Argon to perform the password hashing is of course fine, and a good choice.

---

For me as a user it is unclear what this class is for. Yes, it encrypts things, but what does that tell me? Why would it be used to encrypt, encrypt with a key or encrypt a DB? Generally you should focus on just one 'use case', e.g. encrypting a DB.

Generally, programmers tend to present one interface to the user when it comes to performing writing one or multiple items. Having a separate function for this that does things slightly differently just confuses matters. One way to avoid this is to have a class that is specific for a (password protected) key, which you can then use to encrypt separate items that make up the DB using separate method calls.

Unfortunately the use case / requirements for this class are sufficiently unclear for me to propose a better API for it.

---

With regards to cryptography:

• always indicate the protocol in the description of the class; if you haven't got a separate protocol description then you'll have to provide it here;
• a key and a password are separate entities: password = getpass.getpass is fine, but password = utils.random(32) is not (32 random bytes are not human readable, and hence not a password).

Answer 2

Short comments as docstrings

Docstrings are triple-quoted strings that come immediately after you define a class or function, and they are special in the sense that Python can tell those strings are to be seen as comments that document the code.

For most of your methods you have a short comment immediately before the function, which is great. My suggestion is that you move it to "inside" the function.

For example, instead of

    # Generate an encryption key and return it's value(key)
    def generate_key(self):
        self.salt = utils.random(pwhash.argon2id.SALTBYTES) # 16 bytes == 128 bits
        password = self.password() # Ask a user to input a password
        key = self.kdf(self.size, password, self.salt, opslimit=self.ops, memlimit=self.mem)
        return key # Key have 32 bytes == 256 bits

do

    def generate_key(self):
        """Generate an encryption key and return it's value(key)."""
        self.salt = utils.random(pwhash.argon2id.SALTBYTES) # 16 bytes == 128 bits
        password = self.password() # Ask a user to input a password
        key = self.kdf(self.size, password, self.salt, opslimit=self.ops, memlimit=self.mem)
        return key # Key have 32 bytes == 256 bits

Answer 3

Write as little code as possible inside a try

I see you are writing some of your code in a EAFP (easier to ask for forgiveness than permission) style, as opposed to a LBYL style (look before you leap). This means that you are using try: ... except: ... blocks, which is perfectly fine and often encouraged in Python (article on EAFP vs LBYL in Python).

However, when you write EAFP code, it is advisable to write the least amount of code possible inside the try statement. The idea is that we don't want our except to catch an error other than the one we are trying to prevent. So, while I am not familiar with the libraries you are using, you can probably take something like

    # The inverse of encrypt_db()
    def decrypt_db(self, key_db, message):
        try:
            decrypt_box = secret.SecretBox(bytes.fromhex(key_db))
            decrypted = decrypt_box.decrypt(message)
            assert len(decrypted) == len(message) - decrypt_box.NONCE_SIZE - decrypt_box.MACBYTES
            return decrypted

        except exceptions.CryptoError:
            return 'error'

and rewrite it as

    # The inverse of encrypt_db()
    def decrypt_db(self, key_db, message):
        try:
            decrypt_box = secret.SecretBox(bytes.fromhex(key_db))
            decrypted = decrypt_box.decrypt(message)
        except exceptions.CryptoError:
            return 'error'

        assert len(decrypted) == len(message) - decrypt_box.NONCE_SIZE - decrypt_box.MACBYTES
        return decrypted

Now, depending on what line might actually raise a CryptoError, you would further change that part of the code.