Helper functions for the Cryptography library

saferdata.py

import base64
import os
import time

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

"""
Requires the cryptography library.  At the moment cryptography is at version 38.0.3
To install: pip install cryptography

Derive key iteration timing on Lenovo P15S laptop with 11th Gen Intel i7-1165G7
100000 --> 46.1 ms
200000 --> 100.9 ms
300000 --> 146.3 ms
400000 --> 180.4 ms
500000 --> 245.1 ms
600000 --> 264.0 ms
700000 --> 298.2 ms
800000 --> 341.3 ms
900000 --> 377.2 ms

500000 iterations at ~1/4s is a good tradeoff.
"""


def derive_key(password: str, salt=os.urandom(16), iterations=500000):
    """Takes a password, and optionally a salt and/or iteration count
       and produces a cryptographically strong key.

    Args:
        password (str)  : a password string
        salt (bytes)    : salt to keep keys unique
        iterations (int): number of iterations to improve time-hardness

    Returns:
        key : a 32 byte encryption key for use in encrypting/decrypting
        salt: a 16 byte random number that keeps identical passwords from
              producing identical keys and needs to be saved
    """
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=iterations,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password.encode("utf-8")))

    return key, salt


def is_encrypted(token: str) -> bool:
    """Checks if the specified string is an encrypted token
    by checking for the enc$ prefix.

    Args:
        token (str): a string that is potentially an encrypted token

    Returns:
        bool: True if the string begins with enc$, otherwise False
    """
    if token.startswith("enc$"):
        return True

    return False


def encrypt_string(plaintext: str, key: bytes) -> str:
    """Encrypts the plaintext string using the key

    Args:
        plaintext (str): plaintext string
        key (bytes)    : key used for encryption

    Returns:
        str: the encrypted token prefixed with enc$
    """
    f = Fernet(key)
    enc = f.encrypt(data=plaintext.encode("utf-8"))

    return f"enc${enc.decode()}"


def decrypt_string(token: str, key: bytes) -> str:
    """Attempts to decrypt the token string using the key, if
    it begins with the enc$ prefix.

    Args:
        token (str): encrypted string
        key (bytes): key used for decryption

    Returns:
        str: plaintext string
    """
    if is_encrypted(token):
        f = Fernet(key)
        dec = f.decrypt(token[3:].encode("utf-8"))

        return f"{dec.decode()}"
    else:
        return ""


if __name__ == "__main__":
    my_password = "this.isn't.my.password"
    key, salt = derive_key(my_password)

    secret = "This is a secret message.  Don't tell anyone."
    print(f"Original text: {secret}")
    print(f"Is encrypted?: {is_encrypted(secret)}\n")

    enc_token = encrypt_string(secret, key)
    print(f"Encrypted text: {enc_token}")
    print(f"Is encrypted? : {is_encrypted(enc_token)}\n")

    dec_token = decrypt_string(enc_token, key)
    print(f"Decrypted text: {dec_token}")
    print(f"Is encrypted? : {is_encrypted(dec_token)}\n\n")

    print("Iteration timing")
    for iterations in range(100000, 1000000, 100000):
        start = time.time_ns()
        derive_key(my_password, iterations=iterations)
        print(f"{iterations} --> {(time.time_ns()-start)/1000000:.1f} ms")