"""Too Much Crypto (BlackAlps 2026) - key recovery from the first ciphertext block."""
import sys
from base64 import b64decode
from blake3 import blake3
sys.path.insert(0, "files")
import ascon_light as A  # only reused for encrypt (verification) + getFlag

pt = b'Ascon-AEAD128, with less crypto than in NIST SP 800-232.'
ct = b64decode(b'HcRLfNb6GrW9VIyxcMtcXt6coXZOKKu33xZrEOtzb7FBYJUnOfFy6hYE2Gl80ILp8OqWd6BUIoo=')
nonce = b64decode(b'rNo1NnaqG3EQhTlnVcPaTA==')
tag = b64decode(b'a0PWORRlAlq872NDeLLMtg==')

w = lambda b, n: [int.from_bytes(b[8*i:8*i+8], "little") for i in range(n)]
IV = A.IV_AEAD128
N0, N1 = w(nonce, 2)
# blake3 only sees S0..S2 = IV, N0, N1 -> fully public
h = w(blake3(A._words_to_bytes([IV, N0, N1])).digest()[:24], 3)
# After init: S0 = K1^h0^N1, S1 = IV^h1^K0 ; C_0 = P_0 ^ (S0||S1)
S0, S1 = w(bytes(a ^ b for a, b in zip(ct[:16], pt[:16])), 2)
K1 = S0 ^ h[0] ^ N1
K0 = S1 ^ h[1] ^ IV
key = A._words_to_bytes([K0, K1])
c2, t2 = A.ascon_aead128_encrypt(key, nonce, b"", pt)
assert (c2, t2) == (ct, tag), "verification failed"
print("key =", key.hex())
print(A.getFlag(key))