#!/usr/bin/env python3
"""Recover the SAND .wbt XOR key after a game update changes it — no RE needed.

Known-plaintext crib: the walker icon (textureRawData) is RGBA8888 and its empty
space is the solid background pixel BG = 02 05 0D 00 (RGBA 2,5,13,0). The BSON
header before the pixel data is a CONSTANT size, so pixel byte 0 always lands at
decoded offset 0x2A:
    4  (int32 BSON doc length)
  + 17 (textureSize field:    0x10 + "textureSize\0"(12) + int32(4))
  + 21 (textureRawData header: 0x05 + "textureRawData\0"(15) + int32 len(4) + subtype(1))
  = 42 = 0x2A
Pixel (0,0) is the bottom-left of the image (Unity bottom-up), which is background
on a normal walker, so a long run of BG sits right at offset 0x2A — inside the
first 0xA000 cipher chunk, where the keystream residue is just i % keylen.

Recovery: for each i in [0x2A, 0xA000),
    key[i % keylen] = enc[i] XOR BG[(i - 0x2A) % 4]
Majority-vote each residue (non-background pixels are the minority and wash out),
try keylen 6 first (current), then 1..16, and accept the key whose full decode
parses as BSON. Independent of whatever the new key bytes are.

Usage:
  recover_key.py <wbt> [<wbt> ...]     # print recovered key(s)
"""
import sys, gzip, struct
from collections import Counter

BG = bytes((0x02, 0x05, 0x0D, 0x00))   # background pixel, RGBA (2,5,13,0)
PIXOFF = 0x2A                          # fixed decoded offset of pixel byte 0
CHUNK = 0xA000

def recover(path, keylens=(6,) + tuple(range(1, 17))):
    raw = gzip.decompress(open(path, 'rb').read())
    end = min(CHUNK, len(raw))
    tried = set()
    for L in keylens:
        if L in tried:
            continue
        tried.add(L)
        votes = [Counter() for _ in range(L)]
        for i in range(PIXOFF, end):
            kb = raw[i] ^ BG[(i - PIXOFF) % 4]
            votes[i % L][kb] += 1
        if any(not v for v in votes):
            continue
        key = bytes(v.most_common(1)[0][0] for v in votes)
        if _verifies(raw, key):
            return key, L
    return None, None

def _verifies(raw, key):
    L = len(key)
    dec = bytes(raw[i] ^ key[(i % CHUNK) % L] for i in range(len(raw)))
    # cheap structural check: BSON doc length == total, and known field names present
    if len(dec) < 8:
        return False
    if struct.unpack_from('<i', dec, 0)[0] != len(dec):
        return False
    return dec[5:16] == b'textureSize' and b'textureRawData' in dec[:64] and b'walker' in dec

if __name__ == '__main__':
    files = sys.argv[1:] or sys.exit("usage: recover_key.py <wbt> ...")
    for f in files:
        key, L = recover(f)
        if key:
            print(f"{f.split('/')[-1][:8]}  KEY = {key.hex().upper()}  (keylen {L})")
        else:
            print(f"{f.split('/')[-1][:8]}  FAILED (no BG run / unexpected layout)")