"""
Thin wrapper around pyaes that fixes a mode and encryption style, along with a standard format
for encrypted text to be stored in that allows for easily determining the difference between an encrypted and
non-encrypted file.
"""
from typing import Tuple

import base64
import os
import re
import json
import hashlib

import pyaes

HEADER_TEXT = "OKPY ENCRYPTED FILE FOLLOWS\n" + "-" * 100 + "\n"

# matches keys
KEY_PATTERN = r"[a-z2-7]{52}9999"


def generate_key() -> str:
    """
    Generates a random key
    """
    return to_safe_string(os.urandom(32))


def is_valid_key(key: str) -> bool:
    """
    Returns if this is a valid key
    """
    return re.match("^" + KEY_PATTERN + "$", key) is not None


def get_keys(document: str) -> list:
    """
    Gets all valid keys in the given document
    """
    return re.findall(KEY_PATTERN, document)


def encode_and_pad(data: str, to_length: int) -> bytes:
    """
    Encodes as utf-8 and pads the given data sequence to the given length with
        null characters.

    Padding is to preserve the IND-CPA of messages of different lengths, by
        making sure to pad to a common length.
    """
    encoded = data.encode("utf-8")
    if to_length is None:
        return encoded
    if len(encoded) > to_length:
        raise ValueError(
            "Cannot pad data of length {} to size {}".format(len(encoded), to_length)
        )
    return encoded + b"\0" * (to_length - len(encoded))


def un_pad_and_decode(padded: bytes) -> str:
    """
    Un-pads the given data sequence by stripping trailing null characters and
        recodes it at utf-8.
    """
    return padded.rstrip(b"\0").decode("utf-8")


def hash_key(key : str) -> str:
    """
    Hashes the given key using SHA-256, so we can quickly determine if a
        ciphertext matches a given key.
    """
    return hashlib.sha256(key.encode("utf-8")).hexdigest()


def dump_ct(iv : bytes, ciphertext : bytes, hashed_key : str) -> str:
    """
    Dump the given initialization vector, ciphertext, and hashed key into a
        single string. Add a header.
    """
    return HEADER_TEXT + json.dumps(
        dict(
            iv=to_safe_string(iv),
            ciphertext=to_safe_string(ciphertext),
            hashed_key=hashed_key,
        )
    )


def load_ct(string : str) -> Tuple[bytes, bytes, str]:
    """
    Takes a string created by dump_ct and undoes the transformation, resulting
        in a tuple (iv, ciphertext, hashed_key)
    """
    if not string.startswith(HEADER_TEXT):
        raise ValueError("Invalid ciphertext: does not start with the header")

    string = string[len(HEADER_TEXT) :]

    result = json.loads(string.strip())
    return (
        from_safe_string(result["iv"]),
        from_safe_string(result["ciphertext"]),
        result["hashed_key"],
    )


def encrypt(data: str, key: str, pad_length: int = None) -> str:
    """
    Encrypt the given data using the given key. Tag the result so that it is clear that this is an encrypted file.
    """
    data_as_bytes = encode_and_pad(data, pad_length)

    iv = os.urandom(16)

    encrypter = pyaes.Encrypter(aes_mode_of_operation(key, iv))
    ciphertext = encrypter.feed(data_as_bytes) + encrypter.feed()

    return dump_ct(iv, ciphertext, hash_key(key))


def is_encrypted(encoded_ciphertext: str) -> bool:
    return encoded_ciphertext.startswith(HEADER_TEXT)


def decrypt(encoded_ciphertext: str, key: str) -> str:
    """
    Decrypt the given ciphertext with the given key. The ciphertext must correspond to the format as generated by
        encrypt(data, key)
    """
    iv, ciphertext, hashed_key = load_ct(encoded_ciphertext)

    if hashed_key != hash_key(key):
        raise InvalidKeyException

    decrypter = pyaes.Decrypter(aes_mode_of_operation(key, iv))
    padded_plaintext = decrypter.feed(ciphertext) + decrypter.feed()

    return un_pad_and_decode(padded_plaintext)


def to_safe_string(unsafe_bytes: bytes) -> str:
    # use 9 instead of = for padding so that the string looks more homogenous
    return base64.b32encode(unsafe_bytes).decode("ascii").replace("=", "9").lower()


def from_safe_string(safe_string: str) -> bytes:
    return base64.b32decode(safe_string.upper().replace("9", "=").encode("ascii"))


def aes_mode_of_operation(key : str, iv : bytes) -> pyaes.AESModeOfOperationCBC:
    return pyaes.AESModeOfOperationCBC(from_safe_string(key), iv)


class InvalidKeyException(Exception):
    pass
