#!/usr/bin/env python3
"""Offline recovery helper for the Base64 value in a TagVault v2 NDEF Text record."""

from __future__ import annotations

import argparse
import base64
import getpass
import hashlib
import json
import sys
import unicodedata
from pathlib import Path

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305

MAGIC = b"TVLT\x02"
KDF_ID_PBKDF2_SHA256 = 1
ITERATIONS = 600_000
SALT_BYTES = 16
HEADER_BYTES = 5 + 1 + 4 + SALT_BYTES
NONCE_BYTES = 12
TAG_BYTES = 16
MAX_RECORD_BYTES = 4096


class RecoveryError(Exception):
    """Input is malformed, unsupported, or cannot be authenticated."""


def decode_text_record_value(encoded: bytes) -> bytes:
    """Decode Base64 copied from the NDEF Text value, allowing line wrapping."""
    try:
        compact = b"".join(encoded.split())
        envelope = base64.b64decode(compact, validate=True)
    except (ValueError, base64.binascii.Error) as exc:
        raise RecoveryError("Input is not valid Base64 TagVault text.") from exc
    if len(envelope) > MAX_RECORD_BYTES:
        raise RecoveryError("Record is larger than this helper accepts.")
    return envelope


def decrypt_envelope(envelope: bytes, passphrase: str) -> str:
    if len(envelope) < HEADER_BYTES + NONCE_BYTES + TAG_BYTES + 1:
        raise RecoveryError("Record is too short to be a valid TagVault v2 record.")
    if envelope[: len(MAGIC)] != MAGIC:
        raise RecoveryError("This is not a TagVault v2 record.")
    if envelope[5] != KDF_ID_PBKDF2_SHA256:
        raise RecoveryError("This record uses an unsupported key-derivation method.")

    rounds = int.from_bytes(envelope[6:10], "big", signed=False)
    if rounds != ITERATIONS:
        raise RecoveryError("This record uses unsupported KDF parameters.")

    header = envelope[:HEADER_BYTES]
    salt = envelope[10:HEADER_BYTES]
    combined = envelope[HEADER_BYTES:]
    if len(combined) < NONCE_BYTES + TAG_BYTES + 1:
        raise RecoveryError("Record is missing encrypted data.")

    nonce = combined[:NONCE_BYTES]
    ciphertext_and_tag = combined[NONCE_BYTES:]
    canonical = unicodedata.normalize("NFC", passphrase.strip())
    candidates = [canonical]
    if canonical != passphrase:
        candidates.append(passphrase)
    plaintext = None
    for candidate in candidates:
        key = hashlib.pbkdf2_hmac(
            "sha256", candidate.encode("utf-8"), salt, ITERATIONS, dklen=32
        )
        try:
            plaintext = ChaCha20Poly1305(key).decrypt(nonce, ciphertext_and_tag, header)
            break
        except (InvalidTag, ValueError):
            continue
    if plaintext is None:
        raise RecoveryError("Passphrase is incorrect or the record is damaged.")

    try:
        vault = json.loads(plaintext.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RecoveryError("Decrypted data is not valid TagVault JSON.") from exc
    if not isinstance(vault, dict) or not isinstance(vault.get("secret"), str):
        raise RecoveryError("Decrypted JSON does not contain a TagVault secret.")
    return vault["secret"]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Decrypt a Base64 value copied from a TagVault v2 NDEF Text record."
    )
    parser.add_argument(
        "record_file",
        type=Path,
        help="text file containing only the record's Base64 Text value",
    )
    args = parser.parse_args()

    try:
        encoded = args.record_file.read_bytes()
        if len(encoded) > MAX_RECORD_BYTES * 2:
            raise RecoveryError("Input file is larger than this helper accepts.")
        envelope = decode_text_record_value(encoded)
        passphrase = getpass.getpass("TagVault passphrase (input hidden): ")
        secret = decrypt_envelope(envelope, passphrase)
    except (OSError, RecoveryError) as exc:
        print(f"Recovery failed: {exc}", file=sys.stderr)
        return 1

    print(secret)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())