#!/usr/bin/env python3
"""FLARE-On 13 #7 FlareCalc - static solver.

Usage: python3 -I solve.py [FlareCalc.exe]

Chain (all from Ghidra, see writeup.md):
  * A float divide-by-zero exception handler (FUN_140002a90) reads the last
    three calculator results d1 (latest), d2, d3 and only proceeds when
        int(d1) == 8675309, int(d2) == 88224646, int(1e7 * d3) == 31428571 (22/7)
    It then runs stage 2 = the bytes of the PE's own certificate table
    (file offset of DataDirectory[4], entry +0x6d3a) with
        key string = "8675309-88224646-31428571".
  * Stage 2 reads the WIN_CERTIFICATE PKCS#7, finds custom OID
    1.3.6.1.4.1.1337.1 -> OCTET STRING -> SEQUENCE of {INTEGER id, OCTET STRING}
    and takes id 0xeee, then
        k   = MD5(key[:0x80] || u32le(id))
        b   = fold(k, b*31 + c);  data[i] ^= (b = b*0x6d + 0x39)
        keep one byte, skip (r%3)+1 bytes, r from LCG s=s*0x41c64e6d+0x3039,
        seed s = fold(key, 0x2437, h*129 + c)  (LCG also advanced per skip)
  * The same container holds 1035 modules (ids -> code). Command '3' feeds one
    input byte at a time: id = module[id](byte, state[32]) starting at 0xcb7a;
    reaching 0x1185 means the input is accepted. walk_modules.py finds that
    input (emulating the small decrypted modules with Unicorn):
        FLAREON_ASN1_Labyrinth_Key_2026!
  * Module 0x1185 then XTEA-decrypts item 0x2706 (key = state[0:16]^state[16:32],
    big-endian blocks, delta 0x12345678^0x8c032fc1) -> the flag.
"""
import hashlib
import os
import struct
import sys

import pefile

HERE = os.path.dirname(os.path.abspath(__file__))
EXE = os.path.join(HERE, "loot", "extracted", "FlareCalc.exe")
KEY = b"8675309-88224646-31428571"
OID = bytes.fromhex("2b060104018a3901")      # 1.3.6.1.4.1.1337.1
ITEM_ID = 0x01761E6E ^ 0x01761080            # 0xeee
M32 = 0xFFFFFFFF


def der_len(b, p):
    n = b[p]
    p += 1
    if n & 0x80 == 0:
        return n, p
    k = n & 0x7F
    return int.from_bytes(b[p:p + k], "big"), p + k


def find_item(blob, item_id):
    o = blob.find(OID)
    p = o + 8
    if blob[p] == 1:                          # mirrors FUN_0001159e
        p = o + 0xB
    assert blob[p] == 0x04, "expected OCTET STRING after OID"
    ln, p = der_len(blob, p + 1)
    seq = blob[p:p + ln]
    assert seq[0] == 0x30
    ln, q = der_len(seq, 1)
    end = q + ln
    while q < end:
        assert seq[q] == 0x30
        ln, r = der_len(seq, q + 1)
        nxt = r + ln
        assert seq[r] == 0x02
        il, r = der_len(seq, r + 1)
        iid = int.from_bytes(seq[r:r + il], "big")
        r += il
        assert seq[r] == 0x04
        dl, r = der_len(seq, r + 1)
        if iid == item_id:
            return seq[r:r + dl]
        q = nxt
    raise KeyError(item_id)


def layer1(data, key, item_id):
    k = hashlib.md5(key[:0x80] + struct.pack("<I", item_id)).digest()
    b = 0
    for c in k:
        b = (b * 31 + c) & 0xFF
    out = bytearray()
    for c in data:
        b = (b * 0x6D + 0x39) & 0xFF
        out.append(c ^ b)
    return bytes(out)


def layer2(data, key):
    s = 0x2437
    for c in key:
        s = (s * 0x81 + (c - 256 if c > 127 else c)) & M32
    nxt = lambda v: (v * 0x41C64E6D + 0x3039) & M32
    out, i = bytearray(), 0
    while i < len(data):
        out.append(data[i])
        s = nxt(s)
        skip = s % 3 + 1
        for _ in range(skip):
            s = nxt(s)
        i += 1 + skip
    return bytes(out)


def all_items(pkcs7):
    o = pkcs7.find(OID)
    p = o + 0xB if pkcs7[o + 8] == 1 else o + 8
    ln, p = der_len(pkcs7, p + 1)
    seq = pkcs7[p:p + ln]
    ln, q = der_len(seq, 1)
    end, items = q + ln, {}
    while q < end:
        ln, r = der_len(seq, q + 1)
        nxt = r + ln
        il, r = der_len(seq, r + 1)
        iid = int.from_bytes(seq[r:r + il], "big")
        r += il
        dl, r = der_len(seq, r + 1)
        items[iid] = seq[r:r + dl]
        q = nxt
    return items


def xtea_decrypt_be(data, key16, n_blocks):
    k = struct.unpack("<4I", key16)
    delta, out = 0x12345678 ^ 0x8C032FC1, b""
    for i in range(0, n_blocks * 8, 8):
        v0, v1 = struct.unpack(">2I", data[i:i + 8])
        s = (delta * 32) & M32
        for _ in range(32):
            v1 = (v1 - ((((v0 << 4) ^ (v0 >> 5)) + v0) ^ (s + k[(s >> 11) & 3]))) & M32
            s = (s - delta) & M32
            v0 = (v0 - ((((v1 << 4) ^ (v1 >> 5)) + v1) ^ (s + k[s & 3]))) & M32
        out += struct.pack(">2I", v0, v1)
    return out


def main():
    exe = sys.argv[1] if len(sys.argv) > 1 else EXE
    pe = pefile.PE(exe, fast_load=True)
    sec = pe.OPTIONAL_HEADER.DATA_DIRECTORY[4]
    raw = open(exe, "rb").read()
    certtab = raw[sec.VirtualAddress:sec.VirtualAddress + sec.Size]
    dwlen = struct.unpack_from("<I", certtab)[0]
    pkcs7 = certtab[8:dwlen]
    open(os.path.join(HERE, "loot", "stage2_payload.bin"), "wb").write(certtab[:0x8149D])
    item = find_item(pkcs7, ITEM_ID)
    print(f"[*] cert table @0x{sec.VirtualAddress:x} ({sec.Size} B); item 0x{ITEM_ID:x}: {len(item)} B")
    stage3 = layer2(layer1(item, KEY, ITEM_ID), KEY)
    open(os.path.join(HERE, "loot", "stage3.bin"), "wb").write(stage3)
    print(f"[*] stage3: {len(stage3)} B -> loot/stage3.bin; head {stage3[:16].hex()}")

    idir = os.path.join(HERE, "loot", "items")
    os.makedirs(idir, exist_ok=True)
    items = all_items(pkcs7)
    for iid, d in items.items():
        open(os.path.join(idir, f"item_{iid:x}.bin"), "wb").write(layer2(layer1(d, KEY, iid), KEY))
    print(f"[*] decrypted {len(items)} modules -> loot/items/")

    import importlib.util
    spec = importlib.util.spec_from_file_location("walk", os.path.join(HERE, "walk_modules.py"))
    walk = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(walk)
    path, state = walk.walk()
    print(f"[*] accepted input: {path.decode()}")
    text = open(os.path.join(idir, "item_2706.bin"), "rb").read()
    flag = xtea_decrypt_be(text, bytes(a ^ b for a, b in zip(state[:16], state[16:32])), 5)
    flag = flag.rstrip(b"\0").decode()
    print(f"[+] FLAG: {flag}")
    with open(os.path.join(HERE, "flag.txt"), "w") as f:
        f.write(flag + "\n")


if __name__ == "__main__":
    main()