"""Save as decrypt.py. Python 3.10+ and cryptography; PUBLIC demo data and password.
Install: python -m pip install cryptography
Run: python decrypt.py payload.json
"""
import base64
import json
import re
import sys
from pathlib import Path

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


def decode(value):
    if not isinstance(value, str) or not re.fullmatch(r"[A-Za-z0-9_-]*", value) or len(value) % 4 == 1:
        raise ValueError("Invalid unpadded Base64URL.")
    return base64.urlsafe_b64decode(value + "=" * ((4 - len(value) % 4) % 4))


def decrypt(payload, password):
    kdf = payload.get("kdf", {})
    if (payload.get("version") != 2 or payload.get("algorithm") != "AES-256-GCM"
            or payload.get("keyLength") != 256 or payload.get("encoding") != "base64url"
            or payload.get("tagLength") != 128 or kdf.get("name") != "PBKDF2"
            or kdf.get("hash") != "SHA-256"):
        raise ValueError("Expected a HashyTools AES v2 envelope.")
    iterations = kdf.get("iterations")
    if type(iterations) is not int or not 10000 <= iterations <= 2000000:
        raise ValueError("PBKDF2 iterations must be between 10000 and 2000000.")
    if not isinstance(password, str) or not password:
        raise ValueError("A non-empty password is required.")
    salt, iv, tag = decode(kdf.get("salt")), decode(payload.get("iv")), decode(payload.get("tag"))
    ciphertext, aad = decode(payload.get("ciphertext")), decode(payload.get("aad", ""))
    if len(salt) != 16 or len(iv) != 12 or len(tag) != 16:
        raise ValueError("Expected a 16-byte salt, 12-byte IV and 16-byte tag.")
    key = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt,
                    iterations=iterations).derive(password.encode("utf-8"))
    # AESGCM expects ciphertext followed by the 16-byte authentication tag.
    return AESGCM(key).decrypt(iv, ciphertext + tag, aad).decode("utf-8")


if __name__ == "__main__":
    try:
        if len(sys.argv) != 2:
            raise ValueError("Usage: python decrypt.py payload.json")
        data = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
        result = decrypt(data.get("envelope", data), "correct horse battery staple")
        if "envelope" in data and result != data.get("plaintext"):
            raise ValueError("Test-vector mismatch.")
        print(result)
    except Exception as error:
        print(f"Decryption failed: {type(error).__name__}: {error}", file=sys.stderr)
        sys.exit(1)