JeffCrypt.py

import binascii
import getpass
import io
import lzma
import time
from pathlib import Path
from secrets import token_bytes
from uuid import uuid4

import qrcode
from cryptography.exceptions import InvalidSignature, InvalidTag
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt


class JeffCryptManager:
    # Path to directory to store crypt materials
    CRYPT_MATERIAL_DIRECTORY: Path = Path(__file__).resolve().parent / "crypt_materials"
    # Dump all the generated qr code pngs here
    QR_CODE_DIRECTORY: Path = Path(__file__).resolve().parent / "qr_codes"

    # Size of the KDF salt
    __KDF_SALT_LEN: int = 16
    # Size of the KDF output
    __KDF_OUTPUT_LEN: int = 32
    # Size of the nonce
    __NONCE_LEN: int = 12

    def __init__(self) -> None:
        # Forward declare the types
        self.pepper: bytes
        self.signing_key: Ed25519PrivateKey | None = None
        self.verifier_key: Ed25519PublicKey
        self.fresh_install: bool = False

        # Initialize the required directories if they don't exit
        for directory in [self.CRYPT_MATERIAL_DIRECTORY, self.QR_CODE_DIRECTORY]:
            self.__initialize_directory(directory)

        # Initialize the KDF pepper
        self.__initialize_kdf_pepper()

        # Initialize the signing/verifier keys
        self.__initialize_asymmetric_keys()

    @classmethod
    def __initialize_directory(cls, directory: Path) -> None:
        # Check if the directory exists
        if not directory.exists():
            # Check it's not an file file
            if not directory.is_file():
                # Make the directory
                # If this throws an exception, let it bubble up and crash the program
                directory.mkdir()
            else:
                msg = f"Could not create {directory} since it already exists as a file"
                raise RuntimeError(msg)

    def __initialize_asymmetric_keys(self) -> None:
        # Key len for both types is 32 bytes long
        KEY_LEN: int = 32

        # Figure of the path to the key file we need
        signing_key_path: Path = self.CRYPT_MATERIAL_DIRECTORY / "signing_key.priv"
        verifier_key_path: Path = self.CRYPT_MATERIAL_DIRECTORY / "signing_key.pub"

        # Check if key file already exists on disk
        if signing_key_path.exists() and signing_key_path.is_file():
            # Read from the file on disk
            with signing_key_path.open("r", encoding="utf-8") as reader:
                # This can blow up, let it blow up and bubble up
                hexed_key: str = reader.read()
                key_bytes: bytes = binascii.a2b_hex(hexed_key)

                # Quick and dirty check on expected input length
                if len(key_bytes) != KEY_LEN:
                    msg: str = "Signing key length read from disk is incorrect length"
                    raise ValueError(msg)

                self.signing_key = Ed25519PrivateKey.from_private_bytes(key_bytes)
                self.verifier_key = self.signing_key.public_key()

        # If we didn't have to write the pepper, we might only have the public key
        # and we're only a verfier
        elif not self.fresh_install and verifier_key_path.exists() and verifier_key_path.is_file():
            with verifier_key_path.open("r", encoding="utf-8") as reader:
                # This can blow up, let it blow up and bubble up
                hexed_verifier_key: str = reader.read()
                verifier_key_bytes: bytes = binascii.a2b_hex(hexed_verifier_key)

                # Quick and dirty check on expected input length
                if len(verifier_key_bytes) != KEY_LEN:
                    msg = "Signing key length read from disk is incorrect length"
                    raise ValueError(msg)

                self.verifier_key = Ed25519PublicKey.from_public_bytes(verifier_key_bytes)

        # Fresh install, generate and write keys to disk
        elif self.fresh_install:
            # Attempt to write the signing key it to disk
            self.signing_key = Ed25519PrivateKey.generate()
            self.verifier_key = self.signing_key.public_key()

            # Convert signer key to hex
            key_bytes = self.signing_key.private_bytes_raw()  # type: ignore[attr-defined]
            hexed_signer_key: str = binascii.b2a_hex(key_bytes).decode()
            # Write the hexed signer key to disk
            with signing_key_path.open("w", encoding="utf-8") as writer:
                writer.write(hexed_signer_key)

            # convert verifier key to hex
            key_bytes = self.verifier_key.public_bytes_raw()  # type: ignore[attr-defined]
            hexed_verifier_key = binascii.b2a_hex(key_bytes).decode()
            # Attempt to write the verifier key to disk
            with verifier_key_path.open("w", encoding="utf-8") as writer:
                writer.write(hexed_verifier_key)

        # Some invalid state
        else:
            msg = "Crypt materials forlder in an invalid state"
            raise ValueError(msg)

    def __initialize_kdf_pepper(self) -> None:
        # We're going a bit insane, but lets have a pepper that's 128 bytes long
        PEPPER_LEN: int = 128

        # Figure of the path to the key file we need
        pepper_path: Path = self.CRYPT_MATERIAL_DIRECTORY / "pepper"
        hexed_pepper: str

        # Check if pepper file already exists on disk
        if pepper_path.exists():
            # Read from the file on disk
            with pepper_path.open("r", encoding="utf-8") as reader:
                # This can blow up, let it blow up and bubble up
                hexed_pepper = reader.read()
                pepper_bytes: bytes = binascii.a2b_hex(hexed_pepper)

                if len(pepper_bytes) != PEPPER_LEN:
                    msg = "Stored pepper is not the correct length"
                    raise ValueError(msg)

                self.pepper = pepper_bytes
        else:
            # Generate a pepper and write it to disk
            self.pepper = token_bytes(PEPPER_LEN)

            hexed_pepper = binascii.b2a_hex(self.pepper).decode()
            # Attempt to write it to disk
            with pepper_path.open("w", encoding="utf-8") as writer:
                writer.write(hexed_pepper)

            # If the pepper file doesn't exist, assume we're a fresh install/usage
            self.fresh_install = True

    def __generate_kdf_value(self, input_value: bytes, salt: bytes | None) -> tuple[bytes, bytes]:
        # If no salt is provided, generate one
        if salt is None:
            salt: bytes = token_bytes(self.__KDF_SALT_LEN)

        # Append the pepper to generate the input salt
        input_salt: bytes = salt + self.pepper

        # Generate the KDF from the salt + pepper
        kdf: Scrypt = Scrypt(input_salt, self.__KDF_OUTPUT_LEN, 2**18, 8, 1)

        # Derive the key
        key: bytes = kdf.derive(input_value)

        return salt, key

    def can_generate_messages(self) -> bool:
        return self.signing_key is not None

    def generate_authenticated_message(self, input_string: str, password: bytes) -> bytes:
        # Check we have a signing flag, otherwise we shouldn't be able to generate the messages
        if not self.can_generate_messages():
            msg = "Signing key has not been configured, cannot generated authenticated messages"
            raise AttributeError(msg)

        # KDF the key using the input password and pepper
        kdf_salt: bytes
        symmetric_key: bytes
        kdf_salt, symmetric_key = self.__generate_kdf_value(password, None)

        input_as_bytes: bytes = input_string.encode()
        input_to_encrypt: bytes = input_as_bytes

        # Attempt to compress the data with maximum compression
        compressed_data: bytes = lzma.compress(input_as_bytes, preset=9)

        if len(compressed_data) < len(input_as_bytes):
            # Tag the data with a 0x1 to show it's compressed
            input_to_encrypt = compressed_data + bytes([0x1])
        else:
            # Tag the data with a 0x0 to show it's uncompressed
            input_to_encrypt += bytes([0x0])

        # Signature is the input as bytes
        signature: bytes = self.signing_key.sign(input_as_bytes)  # type: ignore[union-attr]

        # Generate a random nonce
        # TODO: Update to use a generate_nonce function which keeps track of used nonce
        #       such that we don't reuse them, but this will do for now
        nonce: bytes = token_bytes(self.__NONCE_LEN)

        # Encrypt the message
        cipher_text: bytes = ChaCha20Poly1305(symmetric_key).encrypt(nonce, input_to_encrypt, signature)

        # Return the nonce + kdf_salt + signature + cipher_text
        return nonce + kdf_salt + signature + cipher_text

    def generate_authenticated_message_and_qr_code(self, input_str: str, password: bytes) -> tuple[bytes, str, str]:
        # Check we have a signing flag, otherwise we shouldn't be able to generate the messages
        if not self.can_generate_messages():
            msg: str = "Signing key has not been configured, cannot generated authenticated messages"
            raise AttributeError(msg)

        # Generate the message
        generated_message: bytes = self.generate_authenticated_message(input_str, password)

        # Generate the QR code
        qr_code_image = qrcode.QRCode()
        qr_code_image.add_data(generated_message)
        image = qr_code_image.make_image()

        # Attempt to write the image to disk
        # Generate a random file_name
        filename: str = str(uuid4()) + ".png"
        image.save(str(self.QR_CODE_DIRECTORY / filename))

        qr_code_buffer = io.StringIO()
        qr_code_image.print_ascii(qr_code_buffer)

        return generated_message, filename, qr_code_buffer.getvalue()

    def authenticate_message(self, input_bytes: bytes, password: bytes) -> bytes:
        SIG_LEN: int = 64
        TAG_LEN: int = 16

        # The input should be at least 12 bytes for the nonce, 32 bytes for the KDF salt,
        # 64 for the signature and then 16 bytes for the tag
        MIN_LEN = self.__NONCE_LEN + self.__KDF_SALT_LEN + SIG_LEN + TAG_LEN
        if len(input_bytes) < MIN_LEN:
            msg = "message is less than the expected amount"
            raise ValueError(msg)

        # Split the message into its components
        offset: int = 0
        nonce: bytes = input_bytes[offset : offset + self.__NONCE_LEN]
        offset += self.__NONCE_LEN

        kdf_salt: bytes = input_bytes[offset : offset + self.__KDF_SALT_LEN]
        offset += self.__KDF_SALT_LEN

        signature: bytes = input_bytes[offset : offset + SIG_LEN]
        offset += SIG_LEN

        cipher_text: bytes = input_bytes[offset:]

        # Attempt to regenerate the symmetric key
        symmetric_key_bytes: bytes
        _, symmetric_key_bytes = self.__generate_kdf_value(password, kdf_salt)

        try:
            plain_text: bytes = ChaCha20Poly1305(symmetric_key_bytes).decrypt(nonce, cipher_text, signature)
        except (TypeError, ValueError, InvalidTag) as exc:
            # If something went wrong in the decrypt, we have an invalid message
            msg: str = "Failed to decrypt data"
            raise ValueError(msg) from exc

        # Check if the data has been compressed
        data_to_verify: bytes = plain_text

        # Check for the compression tag
        if data_to_verify[-1]:  # noqa: SIM108
            # Decompress it
            data_to_verify = lzma.decompress(plain_text)
        else:
            # Strip off the compression tag
            data_to_verify = data_to_verify[:-1]

        # Now to verify the signature
        try:
            # Check the signature
            self.verifier_key.verify(signature, data_to_verify)

        except InvalidSignature as exc:
            # Signature verification failed
            msg: str = "Failed to verify the signature of the data"
            raise ValueError(msg) from exc
        else:
            # If no exception was thrown, we have a successful verification
            return data_to_verify


def main() -> None:
    password: bytes

    manager: JeffCryptManager = JeffCryptManager()
    if manager.can_generate_messages():
        # We're on a generator machine

        # Stage 1 - Get required details
        input_message: str = input("What is your message to encode?\n")
        password = getpass.getpass("Enter password: ").encode()

        # Stage 2 - Generate authenticated message and QR code
        # Types for unpacking
        authenticated_message: bytes
        file_name: str
        qr_code: str
        # Unpack results immediately
        (
            authenticated_message,
            file_name,
            qr_code,
        ) = manager.generate_authenticated_message_and_qr_code(input_message, password)

        # Stage 3 - Print out details
        print(f"Input message: {input_message}")
        print(f"Authenticated message: {authenticated_message.hex()}")
        print(f"QR code file location: {manager.QR_CODE_DIRECTORY / file_name}")
        print(f"QR code:\n{qr_code}")

        # Stage 4 - Attempt to authenticate generated message
        print("Attempt to decrypt message and validate")
        # For some reason we need to sleep to clean the terminal properly?
        # Flushing stdout doesn't seem to work
        time.sleep(0.5)
        # Prompt for the password again
        password = getpass.getpass("Enter password: ").encode()
        try:
            message: bytes = manager.authenticate_message(authenticated_message, password)
            print(f"Authenticated message: {message.decode()}")
        except ValueError:
            print("Incorrect password entered, failed to authenticate message")
    else:
        # We're running on a verifier

        # Stage 1 - Get required inputs
        try:
            message_to_validate: bytes = binascii.a2b_hex(
                input("What is your message (hex-encoded) to decrypt and verify?\n")
            )
        except binascii.Error:
            print("Invalid input entered")
            return
        password = getpass.getpass("Enter password: ").encode()

        # Stage 2 = Validate and print the data
        print("Attempt to decrypt message and validate")
        try:
            validated_message: bytes = manager.authenticate_message(message_to_validate, password)
            print(f"Authenticated message: {validated_message.decode()}")
        except ValueError as exc:
            print(f"Failed to validate data due to: {exc}")


if __name__ == "__main__":
    main()