cx-keys/assets/secret_coder.py

import os
import hashlib
import Crypto.Cipher.AES

class bad_password(Exception):
    """
    This exception is user to mark when bad password had been used.
    """

    pass

class secret_coder:
    """
    This class is required for crypting and decrypting secret by passwords.
    It store crypted secret as strings in coded form, which make it easies to 
    store in database, and send over REST api.

    Format:
    <hashed_secret>:<cipher>:<crypted_secret>

     - hashed_secret: SHA256 of the secret, to validate success of the decrypt.
     - cipher: Cipher from AES.
     - crypted_secret: Crypted secret to encrypt with password.
    """

    def __init__(self, password: str) -> None:
        """
        This functio initialize new coder. It require password, which would
        be used to encrypting and decrypting secrets.

        Parameters:
            password (str): Password to use as crypto key
        """

        password = password.encode("UTF-8")
        hashed = hashlib.sha256(password)  
        self.__password = hashed.digest()

    @property
    def password(self) -> bytes:
        """ It return password as crypto key. """

        return self.__password

    def encrypt(self, plain: str) -> str:
        """
        This function encrypt new secret, and return it in specified format.

        Parameters:
            plain (str): Secret to encode

        Returns:
            (str): Encrypted secret
        """

        cipher = self.__iv
        coder = self.__cipher(cipher)
        crypted = coder.encrypt(plain.encode("UTF-8"))
        hashed = self.__hash(plain)

        return self.__pack(hashed, cipher, crypted)

    def decrypt(self, crypted: str) -> str:
        """
        This function decrypt secret in given format, and return decrypted 
        secret. When password which initialize coder is bad, then
        bad_password Exception had been raiser.

        Parameters:
            crypted (str): Crypted secret in coder format

        Returns:
            (str): Decrypted secret
        """

        hashed, cipher, crypted = self.__unpack(crypted) 
        coder = self.__cipher(cipher)
        result = coder.decrypt(crypted)

        try:
            result = result.decode("UTF-8")
        except:
            raise bad_password("Crypto key is not correct for this secret.")

        if hashed != self.__hash(result):
            raise bad_password("Crypto key is not correct for this secret.")
        
        return result
        
    @property
    def mode(self) -> int:
        """ This return AES mode which had been used by coder. """
        
        return Crypto.Cipher.AES.MODE_CFB

    @property
    def __iv(self) -> bytes:
        """ This return new random cipher for AES. """

        return os.urandom(self.__cipher_length)

    @property
    def __cipher_length(self) -> int:
        """ This return lenght of random cipher. """

        return 16

    def __cipher(self, iv: bytes) -> object:
        """ This return new AES coder to work with coders. """

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

    def __hash(self, secret: str) -> bytes:
        """
        This function hash given string, and return it as 
        bytes.

        Parameters:
            secret (str): Content to hash

        Returns:
            (bytes): Hashed content
        """

        secret = secret.encode("UTF-8")
        hashed = hashlib.sha256(secret)
        
        return hashed.digest()

    @property
    def separator(self) -> str:
        """ This return separator used to packing secret. """

        return ":"

    def __pack(self, hashed: bytes, cipher: bytes, crypted: bytes) -> str:
        """
        This function pack full content required for secret, to coder format.

        Parameters:
            hashed (bytes): Hashed secret
            cipher (bytes): Ciper used in crypter
            crypted (bytes): Crypted secret

        Returns:
            (str): Coded crypted secret
        """

        hashed = hashed.hex()
        cipher = cipher.hex()
        crypted = crypted.hex()

        return hashed + self.separator + cipher + self.separator + crypted

    def __unpack(self, crypted: str) -> [bytes, bytes, bytes]:
        """
        This function unpack given coded secret into bytes parts. When 
        secret is not valid raise TypeError.

        Parameters:
            crypted (str): Packed secret

        Returns:
            (bytes): Hashed secret before encryption to validate password
            (bytes): Cipher required by AES to work
            (bytes): Crypted secret to encrypt
        """

        splited = crypted.split(self.separator)

        if len(splited) != 3:
            raise TypeError("Secret is in invalid format.")

        hashed = bytes.fromhex(splited[0])
        cipher = bytes.fromhex(splited[1])
        crypted = bytes.fromhex(splited[2])

        if len(cipher) != self.__cipher_length:
            raise TypeError("Secret is in invalid format.")
        
        return hashed, cipher, crypted