#!/usr/bin/env python3
"""Agatha seal check: open one sealed pick and check it.

Python 3, standard library only. The offline part needs no internet, so it
runs in any AI assistant's code tool as well as on your own computer:

  1. it tests its own AES-256-GCM on the GCM specification's AES-256 test
     case (McGrew and Viega, test case 15) before it is used;
  2. it base64-decodes the sealed bytes at the end of the memo:
     nonce (12 bytes) | ciphertext | tag (16 bytes);
  3. it decrypts them with the day key (AES-256-GCM, no associated data) -
     a wrong key or a changed byte fails the tag check;
  4. it compares the SHA-256 of the opened record with the memo's
     fingerprint, and prints the record.

The memo must also be the one on the Solana chain. With --tx and internet
access the script reads the transaction itself from a Solana RPC (status,
fee payer, memo, block time); without them, compare the memo with the text
of the transaction's Memo instruction on https://explorer.solana.com.

The day key is a published AES-256 decryption key (64 hex characters), not a
wallet key. A copying slip - a character of the memo or the key dropped,
added or changed - is said as such (exit 2, or a tag that does not verify):
copy both again, character for character, before calling a seal false.
Exit 0: everything that ran passed; 1: a check failed; 2: an input could
not be read.

Usage:
  python3 seal-check.py MEMO DAY_KEY [--tx SIGNATURE]
  python3 seal-check.py MEMO --tx SIGNATURE     (day key not released yet)
Or paste the values into MEMO, DAY_KEY and TX below and run it with no
arguments. In a notebook or an AI assistant's code tool, paste this whole
file into one cell with MEMO and DAY_KEY filled in (or call
main([memo, day_key]) afterwards): the kernel's own command line is never
read, and the kernel is never exited. How sealing works:
https://agatha.gg/verify
"""

import argparse
import base64
import hashlib
import hmac
import json
import re
import sys
import urllib.request
from datetime import datetime, timezone

MEMO = ""      # agatha:v1:p:<pick>:<date>:<fingerprint>:<base64>
DAY_KEY = ""   # 64 hex characters, published after the day's games are graded
TX = ""        # the transaction signature, for the optional chain read

AGATHA_ADDRESS = "AGATHAhqdW1ojytKypnrm1phNNNXUwNGBXQRovYHaoFN"
RPC = "https://api.mainnet-beta.solana.com"
MEMO_PROGRAMS = ("Memo1UhkJRfHyvLMcVucJwxXeuD728EqVDDwQDxFMNo",
                 "MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr")
MEMO_RE = re.compile(r"agatha:v1:p:([1-9][0-9]{0,8}):([0-9]{4}-[0-9]{2}-[0-9]{2}):"
                     r"([0-9a-f]{64}):([A-Za-z0-9+/]+={0,2})")

# ---- AES-256 encryption (FIPS-197), the only direction GCM needs ----------


def _xt(a):
    return ((a << 1) ^ 0x11B) & 0xFF if a & 0x80 else a << 1


_EXP, _LOG, _x = [0] * 255, [0] * 256, 1
for _i in range(255):
    _EXP[_i], _LOG[_x] = _x, _i
    _x ^= _xt(_x)                      # times 3, a generator of GF(2^8)
SBOX = []
for _a in range(256):
    _b = 0 if _a == 0 else _EXP[(255 - _LOG[_a]) % 255]   # the inverse
    _s = _b
    for _k in range(1, 5):
        _s ^= ((_b << _k) | (_b >> (8 - _k))) & 0xFF
    SBOX.append(_s ^ 0x63)


def _round_keys(key):
    w = [list(key[i:i + 4]) for i in range(0, 32, 4)]
    rcon = 1
    for i in range(8, 60):
        t = list(w[i - 1])
        if i % 8 == 0:
            t = [SBOX[b] for b in t[1:] + t[:1]]
            t[0] ^= rcon
            rcon = _xt(rcon)
        elif i % 8 == 4:
            t = [SBOX[b] for b in t]
        w.append([a ^ b for a, b in zip(w[i - 8], t)])
    return [sum(w[4 * r:4 * r + 4], []) for r in range(15)]


def _encrypt_block(rk, block):
    s = [b ^ k for b, k in zip(block, rk[0])]
    for rnd in range(1, 15):
        s = [SBOX[b] for b in s]
        s = [s[(i + 4 * (i % 4)) % 16] for i in range(16)]          # ShiftRows
        if rnd < 14:                                                 # MixColumns
            m = []
            for c in range(0, 16, 4):
                a = s[c:c + 4]
                t = a[0] ^ a[1] ^ a[2] ^ a[3]
                m += [a[j] ^ t ^ _xt(a[j] ^ a[(j + 1) % 4]) for j in range(4)]
            s = m
        s = [b ^ k for b, k in zip(s, rk[rnd])]
    return bytes(s)

# ---- GCM decryption (NIST SP 800-38D), 96-bit nonce, no associated data ----


def _gmul(x, y):
    z, r = 0, 0xE1 << 120
    for i in range(127, -1, -1):
        if (y >> i) & 1:
            z ^= x
        x = (x >> 1) ^ r if x & 1 else x >> 1
    return z


def gcm_decrypt(key, nonce, ct, tag):
    """The plaintext, or ValueError when the tag does not verify."""
    if len(key) != 32 or len(nonce) != 12 or len(tag) != 16:
        raise ValueError("wrong key, nonce or tag length")
    rk = _round_keys(key)
    h = int.from_bytes(_encrypt_block(rk, bytes(16)), "big")
    y = 0
    for i in range(0, len(ct), 16):
        y = _gmul(y ^ int.from_bytes(ct[i:i + 16].ljust(16, b"\0"), "big"), h)
    y = _gmul(y ^ (len(ct) * 8), h)                      # len(A) = 0, len(C)
    want = y ^ int.from_bytes(_encrypt_block(rk, nonce + b"\0\0\0\1"), "big")
    if not hmac.compare_digest(want.to_bytes(16, "big"), tag):
        raise ValueError("the tag does not verify: wrong key or changed bytes")
    out = bytearray()
    for n, i in enumerate(range(0, len(ct), 16)):
        ks = _encrypt_block(rk, nonce + (n + 2).to_bytes(4, "big"))
        out += bytes(a ^ b for a, b in zip(ct[i:i + 16], ks))
    return bytes(out)


def self_test():
    """GCM spec test case 15 (AES-256, 96-bit nonce), and a forged tag."""
    k = bytes.fromhex("feffe9928665731c6d6a8f9467308308" * 2)
    n = bytes.fromhex("cafebabefacedbaddecaf888")
    p = bytes.fromhex(
        "d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a72"
        "1c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b391aafd255")
    c = bytes.fromhex(
        "522dc1f099567d07f47f37a32a84427d643a8cdcbfe5c0c97598a2bd2555d1aa"
        "8cb08e48590dbb3da7b08b1056828838c5f61e6393ba7a0abcc9f662898015ad")
    t = bytes.fromhex("b094dac5d93471bdec1a502270e3cc6c")
    if gcm_decrypt(k, n, c, t) != p:
        return False
    try:
        gcm_decrypt(k, n, c, bytes(16))
    except ValueError:
        return True
    return False

# ---- the seal ----------------------------------------------------------------


def _clean(text):
    """A pasted value: line breaks and spaces removed (explorers wrap a long
    memo onto lines), and any quotes or backticks around it."""
    return re.sub(r"\s+", "", text or "").strip("'\"`")


def parse_memo(memo):
    """(pick, date, fingerprint, sealed bytes) from a memo, or ValueError
    saying what is wrong with the copy."""
    m = MEMO_RE.fullmatch(_clean(memo))
    if not m:
        raise ValueError("not a whole Agatha seal memo (agatha:v1:p:<pick>:<date>:<fingerprint>:<base64>)")
    try:
        blob = base64.b64decode(m.group(4), validate=True)
    except ValueError as e:  # binascii.Error: a character dropped or added
        raise ValueError(f"its base64 does not decode ({e}): a character is missing, extra or changed") from None
    return int(m.group(1)), m.group(2), m.group(3), blob


def read_key(day_key):
    """The day key's 32 bytes, or ValueError saying what is wrong with the copy."""
    k = _clean(day_key)
    if not re.fullmatch(r"[0-9a-fA-F]{64}", k):
        raise ValueError(f"a day key is 64 hex characters (0-9, a-f); this one has {len(k)} characters")
    return bytes.fromhex(k)


def open_seal(memo, day_key):
    """(record bytes, the SHA-256 hex, whether it equals the fingerprint)."""
    _pick, _day, fp, blob = parse_memo(memo)
    if len(blob) < 29:
        raise ValueError("too few sealed bytes: the memo was cut short")
    record = gcm_decrypt(read_key(day_key), blob[:12], blob[12:-16], blob[-16:])
    digest = hashlib.sha256(record).hexdigest()
    return record, digest, digest == fp


def read_chain(sig, rpc=RPC):
    """The getTransaction result (jsonParsed) from a Solana RPC, or None."""
    body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "getTransaction", "params": [
        sig, {"encoding": "jsonParsed", "commitment": "confirmed",
              "maxSupportedTransactionVersion": 0}]}).encode("ascii")
    req = urllib.request.Request(rpc, data=body, headers={
        "Content-Type": "application/json", "User-Agent": "agatha-seal-check/1"})
    with urllib.request.urlopen(req, timeout=20) as r:
        return json.loads(r.read().decode("utf-8")).get("result")


def check_chain(tx, memo, payer):
    """Lines and an ok flag for a getTransaction result against the memo."""
    if not tx:
        return False, ["chain: no confirmed transaction with that signature"], None
    lines, ok = [], True
    bt = tx.get("blockTime")
    when = datetime.fromtimestamp(bt, timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") if bt else None
    lines.append(f"chain: transaction found in slot {tx.get('slot')}, block time {when}")
    if (tx.get("meta") or {}).get("err") is None and tx.get("meta") is not None:
        lines.append("chain: it succeeded: PASS")
    else:
        ok = False
        lines.append("chain: it did not succeed: FAIL")
    msg = (tx.get("transaction") or {}).get("message") or {}
    keys = msg.get("accountKeys") or [None]
    k0 = keys[0] if isinstance(keys[0], dict) else {"pubkey": keys[0], "signer": True}
    if k0.get("pubkey") == payer and k0.get("signer") is True:
        lines.append(f"chain: fee payer and signer is {payer}: PASS")
    else:
        ok = False
        lines.append(f"chain: fee payer is {k0.get('pubkey')}, not {payer}: FAIL")
    memos = [ix.get("parsed") for ix in msg.get("instructions") or []
             if isinstance(ix, dict) and ix.get("programId") in MEMO_PROGRAMS]
    if _clean(memo) in memos:
        lines.append("chain: its memo is identical to this memo: PASS")
    else:
        ok = False
        lines.append("chain: its memo differs from this memo: FAIL")
    return ok, lines, bt


def _in_kernel():
    """Inside a notebook kernel (Jupyter, an AI assistant's code tool): its
    command line is the kernel's own (ipykernel_launcher.py -f <file>), never
    this script's, and exiting would end the cell with a traceback."""
    return ("ipykernel" in sys.modules or "IPython" in sys.modules
            or (len(sys.argv) > 2 and sys.argv[1] == "-f" and sys.argv[2].endswith(".json")))


def main(argv=None):
    if argv is None:
        argv = [] if _in_kernel() else sys.argv[1:]
    ap = argparse.ArgumentParser(description="Open and check one Agatha seal.")
    ap.add_argument("memo", nargs="?", default=MEMO)
    ap.add_argument("day_key", nargs="?", default=DAY_KEY)
    ap.add_argument("--tx", default=TX, help="transaction signature: read it from a Solana RPC")
    ap.add_argument("--payer", default=AGATHA_ADDRESS)
    ap.add_argument("--rpc", default=RPC)
    a = ap.parse_args(argv)
    if not a.memo:
        print("memo: NOT GIVEN - run: python3 seal-check.py MEMO DAY_KEY [--tx SIGNATURE], or paste the values "
              "into MEMO, DAY_KEY and TX at the top of this file (in a notebook: main([memo, day_key]))")
        return 2
    ok = self_test()
    print("self-test, GCM spec test case 15 and a forged tag:", "PASS" if ok else "FAIL")
    if not ok:
        return 1
    try:
        pick, day, fp, _blob = parse_memo(a.memo)
    except ValueError as e:
        print(f"memo: NOT READ - {e}. Copy the memo again, whole and character for character "
              "(it starts agatha:v1:p: and ends with the base64), and run this again; a copying slip is "
              "not a failed seal.")
        return 2
    print(f"memo: pick #{pick}, board day {day}, fingerprint {fp} ({len(_clean(a.memo))} characters)")
    start = None
    if a.day_key:
        try:
            read_key(a.day_key)
        except ValueError as e:
            print(f"day key: NOT READ - {e}. Copy it again from https://agatha.gg/keys.json and run this again.")
            return 2
        try:
            record, digest, same = open_seal(a.memo, a.day_key)
        except ValueError as e:
            print("open the seal with the day key: FAIL -", e)
            print("  a wrong key (another day's?) or one changed character in the memo or the key fails the tag: "
                  "copy both again, character for character, and run this again before calling the seal false")
            return 1
        print("open the seal with the day key: PASS")
        print("SHA-256 of the record:", digest)
        print("equals the fingerprint:", "PASS" if same else "FAIL")
        print("record:", record.decode("ascii", "replace"))
        if not same:
            print("  the seal opened, so its bytes are whole: check that the memo's fingerprint field was copied "
                  "exactly before calling the seal false")
            return 1
        try:
            opened = json.loads(record)
            start = opened.get("start") if isinstance(opened, dict) else None
        except ValueError:
            start = None
    else:
        print("open the seal: NOT RUN - the day key is not released yet")
    if not a.tx:
        print("chain: NOT RUN - compare this memo with the transaction's memo on "
              "https://explorer.solana.com, or run again with --tx SIGNATURE")
        return 0
    try:
        tx = read_chain(a.tx, a.rpc)
    except Exception as e:  # noqa: BLE001 - no network, a refusal: said, not raised
        print(f"chain: NOT RUN - could not reach {a.rpc} ({e.__class__.__name__}); compare "
              f"the memo on https://explorer.solana.com/tx/{a.tx} instead")
        return 0
    chain_ok, lines, bt = check_chain(tx, a.memo, a.payer)
    for line in lines:
        print(line)
    if bt and start:
        try:
            begins = datetime.fromisoformat(str(start).replace("Z", "+00:00")).timestamp()
            print("timing: sealed before the start", start + ":", "PASS" if bt < begins else "FAIL")
            chain_ok = chain_ok and bt < begins
        except ValueError:
            print("timing: NOT RUN - the record's start is not a timestamp")
    return 0 if chain_ok else 1


if __name__ == "__main__":
    _code = main()
    if not _in_kernel():      # a notebook kernel is not exited: the lines above are the answer
        sys.exit(_code)