#!/usr/bin/env python3
"""Balance notes: an encrypted snapshot of a player's bank balances, carried in the memo of every
DC payout, so DC's own transaction history is an independent record of what the bank said each
player held (disaster recovery, disputes, reconciliation).

Standard library only: this one file is both the server's codec and the operators' decoder. Copy
it anywhere with Python 3.10+ and run:

  ./notecodec.py decode  --key <BANK_NOTE_KEY> MEMO [MEMO ...]   # parts of one note, any order
  ./notecodec.py recover --key <BANK_NOTE_KEY> feed.json         # latest balance per account

feed.json is DC transaction items (`python -m app.cli notes-export`, or the raw getTransactions
pages). The key is also readable from the BANK_NOTE_KEY environment variable.

Format v2 (all integers big-endian), 21 + 10 bytes per account before base64:
  memo part = base64url, no padding, of: note id u16 | part index u4 | part count u4 | chunk
  blob      = nonce 4B | plaintext XOR keystream | tag 8B           (the chunks joined)
  plaintext = version u4 | key id u4 | time u32 (unix seconds) | recipient Minecraft UUID, first 4B
              | per account: flags u2 + (number - 4000000000) u30 | balance i48 in 1/10000
  keystream = HMAC-SHA256(Kenc, note id | nonce | counter u32) blocks
  tag       = HMAC-SHA256(Kmac, note id | nonce | ciphertext)[:8]     (encrypt-then-MAC)
  Kenc, Kmac = HMAC-SHA256(master, b"enc" / b"mac")
Every field is fixed width, so the number of parts depends only on the number of accounts.
Account numbers are always 4 followed by 9 digits. v1 notes (longer fields) still decode.
"""

import argparse
import base64
import hashlib
import hmac
import json
import os
import re
import secrets
import struct
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from decimal import Decimal

VERSION = 2
NONCE, TAG, HEAD, REC, PART_HEAD = 4, 8, 9, 10, 3
NUMBER_BASE, BALANCE_MAX = 4_000_000_000, 1 << 47  # i48 units: up to ~14 billion
V1 = {"nonce": 8, "tag": 12, "head": 26, "rec": 14}  # notes sent before v2, still readable
MAX_PARTS = 15
FLAG_ORG, FLAG_SAVINGS = 1, 2
UNIT = Decimal("0.0001")
B64 = re.compile(r"[A-Za-z0-9_-]+")


class NoteError(ValueError):
    pass


@dataclass
class Entry:
    number: str
    balance: Decimal
    org: bool = False
    savings: bool = False


@dataclass
class Snapshot:
    ts: datetime
    uuid: str  # undashed hex: v2 notes carry the first 8 characters, v1 all 32
    entries: list[Entry] = field(default_factory=list)
    key_id: int = 1

    def as_dict(self) -> dict:
        return {
            "ts": self.ts.isoformat(),
            "minecraft_uuid": self.uuid,
            "key_id": self.key_id,
            "accounts": [
                {
                    "number": e.number,
                    "balance": f"{e.balance:.4f}",
                    "org": e.org,
                    "savings": e.savings,
                }
                for e in self.entries
            ],
        }


def parse_key(key: str | bytes) -> bytes:
    raw = key if isinstance(key, bytes) else base64.urlsafe_b64decode(key + "=" * (-len(key) % 4))
    if len(raw) < 32:
        raise NoteError("note key must be at least 32 bytes (base64)")
    return raw


def new_key() -> str:
    return base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode()


def _keys(master: bytes) -> tuple[bytes, bytes]:
    return (
        hmac.new(master, b"enc", hashlib.sha256).digest(),
        hmac.new(master, b"mac", hashlib.sha256).digest(),
    )


def _stream(kenc: bytes, nonce: bytes, n: int) -> bytes:
    out = bytearray()
    ctr = 0
    while len(out) < n:
        out += hmac.new(kenc, nonce + struct.pack(">I", ctr), hashlib.sha256).digest()
        ctr += 1
    return bytes(out[:n])


def _b64(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def _unb64(s: str) -> bytes:
    s = s.strip()
    if not s or not B64.fullmatch(s):  # b64decode would silently drop stray characters
        raise NoteError("not a bank note")
    try:
        return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))
    except (ValueError, TypeError) as e:
        raise NoteError("not a bank note") from e


def part_capacity(memo_max: int) -> int:
    """Raw bytes of blob that fit in one memo of memo_max characters."""
    cap = memo_max * 3 // 4 - PART_HEAD
    if cap < 8:
        raise NoteError("memo limit too small for notes")
    return cap


def part_count(n_accounts: int, memo_max: int) -> int:
    size = NONCE + HEAD + REC * n_accounts + TAG
    cap = part_capacity(memo_max)
    n = -(-size // cap)
    if n > MAX_PARTS:
        raise NoteError(f"note would need {n} parts (max {MAX_PARTS})")
    return n


def encode(key, snap: Snapshot, memo_max: int) -> list[str]:
    kenc, kmac = _keys(parse_key(key))
    if not 0 <= snap.key_id < 16:
        raise NoteError("key id must be 0-15")
    plain = bytearray(
        struct.pack(">BI", VERSION << 4 | snap.key_id, int(snap.ts.timestamp()))
        + bytes.fromhex(snap.uuid[:8])
    )
    for e in snap.entries:
        flags = (FLAG_ORG if e.org else 0) | (FLAG_SAVINGS if e.savings else 0)
        n = int(e.number) - NUMBER_BASE
        units = int((Decimal(e.balance) / UNIT).to_integral_value())
        if not 0 <= n < 1 << 30:
            raise NoteError(f"account number {e.number} is not 4 followed by 9 digits")
        if not -BALANCE_MAX <= units < BALANCE_MAX:
            raise NoteError(f"balance of {e.number} is too large for a note")
        plain += struct.pack(">I", flags << 30 | n) + units.to_bytes(6, "big", signed=True)
    sid = struct.pack(">H", secrets.randbelow(1 << 16))
    nonce = secrets.token_bytes(NONCE)
    ct = bytes(a ^ b for a, b in zip(plain, _stream(kenc, sid + nonce, len(plain)), strict=True))
    blob = nonce + ct + hmac.new(kmac, sid + nonce + ct, hashlib.sha256).digest()[:TAG]
    cap, count = part_capacity(memo_max), part_count(len(snap.entries), memo_max)
    return [
        _b64(sid + bytes([i << 4 | count]) + blob[i * cap : (i + 1) * cap]) for i in range(count)
    ]


def part_info(memo: str) -> tuple[int, int, int, bytes]:
    """(snapshot id, index, count, chunk) of one memo part; NoteError if it can't be one."""
    raw = _unb64(memo)
    if len(raw) <= PART_HEAD:
        raise NoteError("not a bank note")
    sid, ic = struct.unpack(">HB", raw[:PART_HEAD])
    i, n = ic >> 4, ic & 15
    if n == 0 or i >= n:
        raise NoteError("not a bank note")
    return sid, i, n, raw[PART_HEAD:]


def decode(key, parts: list[str]) -> Snapshot:
    infos = [part_info(p) for p in parts]
    sids, counts = {x[0] for x in infos}, {x[2] for x in infos}
    if len(sids) != 1 or len(counts) != 1:
        raise NoteError("parts belong to different notes")
    chunks = {x[1]: x[3] for x in infos}
    n = counts.pop()
    if sorted(chunks) != list(range(n)):
        raise NoteError(f"note has {n} parts; got {sorted(i + 1 for i in chunks)}")
    blob = b"".join(chunks[i] for i in range(n))
    kenc, kmac = _keys(parse_key(key))
    sid = struct.pack(">H", sids.pop())
    if len(blob) >= NONCE + HEAD + TAG and (len(blob) - NONCE - HEAD - TAG) % REC == 0:
        nonce, ct, tag = blob[:NONCE], blob[NONCE:-TAG], blob[-TAG:]
        mac = hmac.new(kmac, sid + nonce + ct, hashlib.sha256).digest()[:TAG]
        if hmac.compare_digest(tag, mac):
            plain = bytes(
                a ^ b for a, b in zip(ct, _stream(kenc, sid + nonce, len(ct)), strict=True)
            )
            return _parse_v2(plain)
    return _decode_v1(kenc, kmac, blob)


def _ts(seconds: float) -> datetime:
    # timezone.utc, not datetime.UTC: this file must also run on Python 3.10
    return datetime.fromtimestamp(seconds, timezone.utc)  # noqa: UP017


def _parse_v2(plain: bytes) -> Snapshot:
    vk, secs = struct.unpack(">BI", plain[:5])
    if vk >> 4 != VERSION:
        raise NoteError(f"unknown note version {vk >> 4}")
    snap = Snapshot(ts=_ts(secs), uuid=plain[5:HEAD].hex(), key_id=vk & 15)
    for o in range(HEAD, len(plain), REC):
        (word,) = struct.unpack(">I", plain[o : o + 4])
        flags = word >> 30
        snap.entries.append(
            Entry(
                number=str(NUMBER_BASE + (word & ((1 << 30) - 1))),
                balance=Decimal(int.from_bytes(plain[o + 4 : o + REC], "big", signed=True)) * UNIT,
                org=bool(flags & FLAG_ORG),
                savings=bool(flags & FLAG_SAVINGS),
            )
        )
    return snap


def _decode_v1(kenc: bytes, kmac: bytes, blob: bytes) -> Snapshot:
    """v1: nonce 8 | ct | tag 12; plaintext version u8 | key id u8 | time µs u64 | full UUID 16B |
    per account number u40 | flags u8 | balance i64. The note id was not authenticated."""
    nn, tl, hd, rc = V1["nonce"], V1["tag"], V1["head"], V1["rec"]
    if len(blob) < nn + hd + tl or (len(blob) - nn - hd - tl) % rc:
        raise NoteError("note does not verify (wrong key, or edited)")
    nonce, ct, tag = blob[:nn], blob[nn:-tl], blob[-tl:]
    if not hmac.compare_digest(tag, hmac.new(kmac, nonce + ct, hashlib.sha256).digest()[:tl]):
        raise NoteError("note does not verify (wrong key, or edited)")
    plain = bytes(a ^ b for a, b in zip(ct, _stream(kenc, nonce, len(ct)), strict=True))
    version, key_id, us = struct.unpack(">BBQ", plain[:10])
    if version != 1:
        raise NoteError(f"unknown note version {version}")
    snap = Snapshot(ts=_ts(us / 1_000_000), uuid=plain[10:hd].hex(), key_id=key_id)
    for o in range(hd, len(plain), rc):
        flags, units = struct.unpack(">Bq", plain[o + 5 : o + rc])
        snap.entries.append(
            Entry(
                number=str(int.from_bytes(plain[o : o + 5], "big")),
                balance=Decimal(units) * UNIT,
                org=bool(flags & FLAG_ORG),
                savings=bool(flags & FLAG_SAVINGS),
            )
        )
    return snap


def scan(key, memos: list[tuple[str, dict]]) -> list[tuple[Snapshot, list[dict]]]:
    """Find every decodable note among (memo, item) pairs, joining multi-part notes. Memos that
    are not ours (wrong key, plain text, someone else's traffic) are skipped."""
    groups: dict[tuple[int, int], dict[int, tuple[str, dict]]] = {}
    for memo, item in memos:
        try:
            sid, i, n, _ = part_info(memo)
        except NoteError:
            continue
        groups.setdefault((sid, n), {})[i] = (memo, item)
    out = []
    for (_, n), parts in groups.items():
        if len(parts) != n:
            continue
        try:
            snap = decode(key, [parts[i][0] for i in range(n)])
        except NoteError:
            continue
        out.append((snap, [parts[i][1] for i in range(n)]))
    out.sort(key=lambda x: x[0].ts)
    return out


def feed_items(doc) -> list[dict]:
    """Accept a list of items, one getTransactions page ({"items": [...]}) or a list of pages."""
    if isinstance(doc, dict):
        doc = doc.get("items") or []
    items = []
    for x in doc:
        if isinstance(x, dict) and "items" in x:
            items += x["items"] or []
        elif isinstance(x, dict):
            items.append(x)
    return items


def recover(key, items: list[dict]) -> dict:
    """Latest known balance per account number, and which Minecraft account reported it."""
    memos = [(str(it.get("memo") or it.get("message") or ""), it) for it in items]
    latest: dict[str, dict] = {}
    snaps = scan(key, memos)
    for snap, _ in snaps:  # oldest first, so later notes overwrite
        for e in snap.entries:
            latest[e.number] = {
                "balance": f"{e.balance:.4f}",
                "as_of": snap.ts.isoformat(),
                "org": e.org,
                "savings": e.savings,
                "reported_to": snap.uuid,
            }
    return {"notes_found": len(snaps), "accounts": dict(sorted(latest.items()))}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Decode bank balance notes from DC payout memos.")
    sub = ap.add_subparsers(dest="cmd", required=True)
    d = sub.add_parser("decode", help="decode the memo part(s) of one note")
    d.add_argument("memos", nargs="+")
    r = sub.add_parser("recover", help="latest balance per account from a DC feed export")
    r.add_argument("feed", help="JSON file, or - for stdin")
    sub.add_parser("new-key", help="print a fresh BANK_NOTE_KEY")
    for p in (d, r):
        p.add_argument("--key", default=os.environ.get("BANK_NOTE_KEY"), help="BANK_NOTE_KEY")
    a = ap.parse_args(argv)
    if a.cmd == "new-key":
        print(new_key())
        return 0
    if not a.key:
        ap.error("--key (or BANK_NOTE_KEY) is required")
    try:
        if a.cmd == "decode":
            print(json.dumps(decode(a.key, a.memos).as_dict(), indent=2))
        else:
            if a.feed == "-":
                doc = json.load(sys.stdin)
            else:
                with open(a.feed, encoding="utf-8") as f:
                    doc = json.load(f)
            print(json.dumps(recover(a.key, feed_items(doc)), indent=2))
    except NoteError as e:
        print(f"error: {e}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())