Encrypting Integer Primary Keys in Python

>>> PrimaryKeyEncryptor.generate_secret()
'134810bef265eabbfcc33c8956b9c870'

>>> PrimaryKeyEncryptor('134810bef265eabbfcc33c8956b9c870').encrypt(1)
'ec68e9d12ae27ce7ce9bad7488ecf665'

>>> PrimaryKeyEncryptor('134810bef265eabbfcc33c8956b9c870').encrypt(2)
'd2f225fd542442c6e12ab19943749391'

>>> PrimaryKeyEncryptor('134810bef265eabbfcc33c8956b9c870').decrypt('ec68e9d12ae27ce7ce9bad7488ecf665')
1

>>> PrimaryKeyEncryptor('134810bef265eabbfcc33c8956b9c870').decrypt('d2f225fd542442c6e12ab19943749391')
2

>>> PrimaryKeyEncryptor('134810bef265eabbfcc33c8956b9c870').decrypt('d2f225fd542442c6e12ab19943749393')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/test.py", line 105, in decrypt
    raise ValueError('The encrypted primary key is invalid')
ValueError: The encrypted primary key is invalid

"""This module contains a python class that was designed for encrypting sequential integer primary keys

    Typical usage example:

    # Generate a secret (done once)
    secret = PrimaryKeyEncryptor.generate_secret()

    # Instantiate encryptor
    encryptor = PrimaryKeyEncryptor(secret)

    # Encrypt
    encrypted = encryptor.encrypt(primary_key)

    # Decrypt
    decrypted_primary_key = encryptor.decrypt(encrypted)
"""
from os import urandom
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.algorithms import AES as Algorithm
from cryptography.hazmat.primitives.ciphers.modes import ECB as Mode


class PrimaryKeyEncryptor:
    """This class uses AES-128 in ECB mode to generate a secure mapping between a 64 bit integer and a 32 character hex
    string.

    Note:
        In many other contexts, the use of AES-128 in ECB mode is insecure. Please do not use this example for anything
        other than the explit purpose of encrypting sequential integer primary keys generated by a storage engine.
    """

    def __init__(self, secret: str):
        """Initialize a PrimaryKeyEncryptor class.

        Args:
            secret: A 32 character (16 byte) hex string which is used as a key for encryption and decryption.

        Raises:
            ValueError: If the secret is not a 32 byte hex string.
        """

        secret_bytes = bytes.fromhex(secret)

        if len(secret_bytes) != 16:
            raise ValueError('The secret for the PrimaryKeyEncryptor must be 16 bytes in hexadecimal format')

        algorithm = Algorithm(secret_bytes)
        mode = Mode()

        self.cipher = Cipher(algorithm, mode, backend=default_backend())

    @staticmethod
    def generate_secret() -> str:
        """Generate a secret that can be used to instantiate the PrimaryKeyEncryptor class.

        Returns:
            A 32 charachter (16 byte) hex string which is suitable to use as a key for encryption and decryption.
        """

        return urandom(16).hex()

    def encrypt(self, primary_key: int) -> str:
        """Encrypt an integer primary key

        Args:
            primary_key: An integer to encrypt, this cannot be more than 64 bits in size.

        Returns:
            The encrypted integer primary key as a 32 charachter (16 byte) hex string.

        Raises:
            ValueError: If the primary key is not an integer of 64 bits or less.
        """
        primary_key_bytes = primary_key.to_bytes(8, byteorder='big')

        encryptor = self.cipher.encryptor()

        cipher_bytes = encryptor.update(primary_key_bytes * 2) + encryptor.finalize()

        return cipher_bytes.hex()

    def decrypt(self, encrypted_primary_key: str) -> int:
        """Encrypt an integer primary key

        Args:
            encrypted_primary_key: The encrypted integer primary key as a 32 character (16 byte) hex string.

        Returns:
            The decrypted integer primary key.

        Raises:
            ValueError: If the encrypted primary key is not valid.
        """
        cipher_bytes = bytes.fromhex(encrypted_primary_key)

        if len(cipher_bytes) != 16:
            raise ValueError('The encrypted primary key must be 16 bytes in hexadecimal format')

        decryptor = self.cipher.decryptor()

        plain_bytes = decryptor.update(cipher_bytes) + decryptor.finalize()

        if plain_bytes[:8] != plain_bytes[8:]:
            raise ValueError('The encrypted primary key is invalid')

        return int.from_bytes(plain_bytes[:8], byteorder='big')