cx-keys/assets/secret_crypto.py

import Crypto.Cipher.AES
import hashlib
import os

class secret_crypto:
    """
    This class is responsible for creation and manage of the secrets crypto.
    All of the secrets would have own IV. When new secret_crypto is creating,
    IV is random. It could be change by function set_iv. Set IV is necessary
    to decrypt secret.

    Password HAVE NOT being stored in database. It must be given from API, 
    to protect database. Password from the user is hashed by SHA256 to use it 
    as key. IV is random, and other for all of the keys.
    
    If database leaks, secrets are save, because keys is provided from user
    every request for the secret.
    """

    def __init__(
        self, 
        password: str, 
        __iv: bytes | None = None,
        __key: bytes | None = None
    ) -> None:
        """
        This function create new secret_crypto. It could be used encrypt new
        secret. To decrypt secret, app must set IV.

        password: str - Password to encrypt secrets
        __iv: bytes | None - IV, used to clone (default: None)
        __key: bytes | None - Key, used to clone (default: None)
        """

        if __key is None:
            encoded = password.encode("UTF-8")
            hashed = hashlib.sha256(encoded)
            __key = hashed.digest()

        if __iv is None:
            __iv = os.urandom(16)

        self.__iv = __iv
        self.__key = __key
        self.__password = password

    def encrypt(self, secret: str) -> bytes:
        """
        This function encrypt secret, using previously loaded password as key
        and IV. IV would be random, and own for all secrets.

        secret: str - Secret to encode
        return: bytes - Encrypted secret
        """

        return self.__cipher.encrypt(secret.encode("UTF-8"))

    def decrypt(self, crypted: bytes) -> str:
        """
        This function decrypt previously crypted secret using given password
        as key. IV must be restore. When random IV is used, then cipher would
        not decrypt secret.
        
        crypted: bytes - Crypted secret to decrypt
        return: str - Decrypted secret
        """

        return self.__cipher.decrypt(crypted).decode("UTF-8")

    def crypted(self, secret: str) -> [bytes, bytes]:
        """
        This function is complex version of encrypt. It do exacly the same, 
        but it return also IV, which could make code look better.

        secret: str - Secret to encrypt
        return: [bytes, bytes] - [Encrypted secret, IV]
        """

        return self.encrypt(secret), self.iv

    @property
    def password(self) -> str:
        """
        return: str - Password to encrypt
        """

        return self.__password

    @property
    def key(self) -> bytes:
        """
        return: bytes - Password as 256 bit binary key
        """

        return self.__key

    @property
    def iv(self) -> bytes:
        """
        return: bytes - IV of the crypto, get it to save in database
        """

        return self.__iv

    def set_iv(self, iv: bytes) -> object:
        """
        This function set IV. IV must be save in database, to decrypt secret
        later. This function is used to restore IV, previously read from 
        database.

        iv: bytes - IV read from database
        return: secret_crypto - Clone of crypto with set given iv
        """

        return secret_crypto(self.password, iv, self.key)

    @property
    def mode(self) -> int:
        """
        return: int - Mode of the AES used to crypt
        """

        return Crypto.Cipher.AES.MODE_CFB

    @property
    def __cipher(self) -> object:
        """
        return: object - New clean cipher, which could be used to work
        """

        return Crypto.Cipher.AES.new(
            self.key, 
            self.mode, 
            iv = self.iv
        )