#!/usr/bin/env python
import json
import unpaddedbase64
from Crypto import Random
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.Util import Counter

class EncryptionError(Exception):
    pass

def decrypt(ciphertext: bytes, key: str, hash: str, iv: str):
    """Decrypt an encrypted attachment.
    Args:
        ciphertext (bytes): The data to decrypt.
        key (str): AES_CTR JWK key object.
        hash (str): Base64 encoded SHA-256 hash of the ciphertext.
        iv (str): Base64 encoded 16 byte AES-CTR IV.
    Returns:
        The plaintext bytes.
    Raises:
        EncryptionError if the integrity check fails.
    """
    expected_hash = unpaddedbase64.decode_base64(hash)

    h = SHA256.new()
    h.update(ciphertext)

    if h.digest() != expected_hash:
        raise EncryptionError("Mismatched SHA-256 digest.")

    try:
        byte_key: bytes = unpaddedbase64.decode_base64(key)
    except (BinAsciiError, TypeError):
        raise EncryptionError("Error decoding key.")

    try:
        # Drop last 8 bytes, which are 0
        byte_iv: bytes = unpaddedbase64.decode_base64(iv)[:8]
    except (BinAsciiError, TypeError):
        raise EncryptionError("Error decoding initial values.")

    ctr = Counter.new(64, prefix=byte_iv, initial_value=0)

    try:
        cipher = AES.new(byte_key, AES.MODE_CTR, counter=ctr)
    except ValueError as e:
        raise EncryptionError(e)

    return cipher.decrypt(ciphertext)


# if __name__ == "__main__":
    # with open('images/output', 'wb') as output:
        # with open('images/LUJAssHxtTWsnYPbSlTcMdvl.octet-stream', 'rb') as cipher:
            # with open('images/LUJAssHxtTWsnYPbSlTcMdvl.metadata', 'r') as rawmeta:
                # meta = json.load(rawmeta)
                # key = meta['file']['key']
                # decrypted = decrypt_attachment(cipher.read(), key['k'], meta['file']['hashes']['sha256'], meta['file']['iv'])
                # output.write(decrypted)