#!/usr/bin/env python3
"""Find the byte sequence that drives FlareCalc's stage-2 module chain from
id 0xcb7a to 0x1185, then decrypt item 0x2706 with module 0x1185.

Usage: python3 -I walk_modules.py
Each module is a small pure function  ret_id = f(cl=byte, rdx=state[32]).
They are *emulated* with Unicorn (sandboxed CPU emulator; nothing runs natively
and the challenge binary itself is never executed). The junk busy loop in each
module (sum 0..999999 = 0x6a4ae6e0, cancelled by "-0x6a4ae6e0") is skipped
by setting its counter/sum to their final values.
"""
import os
import re
import sys

import capstone
from unicorn import Uc, UC_ARCH_X86, UC_MODE_64, UC_HOOK_CODE, UC_HOOK_INSN
from unicorn.x86_const import (UC_X86_REG_RAX, UC_X86_REG_RBX, UC_X86_REG_RCX, UC_X86_REG_RDX,
                               UC_X86_REG_RIP, UC_X86_REG_RSP, UC_X86_REG_R8, UC_X86_REG_R9,
                               UC_X86_INS_CPUID)

HERE = os.path.dirname(os.path.abspath(__file__))
ITEMS = os.path.join(HERE, "loot", "items")
START, FINAL = 0xCB7A, 0x1185
CODE, STATE, STACK, RET = 0x100000, 0x200000, 0x300000, 0x0DEAD000
MD = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)

mods = {int(f[5:-4], 16): open(os.path.join(ITEMS, f), "rb").read()
        for f in os.listdir(ITEMS) if f.startswith("item_")}


def loop_info(code):
    """Return (head, exit, counter_slot, sum_slot, limit) of the busy loop, if any."""
    def slot(s):
        m = re.search(r"\[rsp(?: \+ (0x[0-9a-f]+|\d+))?\]", s)
        return (int(m.group(1), 0) if m.group(1) else 0) if m else None
    ins = list(MD.disasm(code, 0))
    for k, i in enumerate(ins):
        if i.mnemonic == "jb" and int(i.op_str, 16) < i.address:
            head, exit_ = int(i.op_str, 16), ins[k + 1].address
            body = [x for x in ins if head <= x.address <= i.address]
            # counter: the slot loaded right before the cmp; sum: the other slot stored in the body
            ctr = next((slot(x.op_str) for x in reversed(body[:-2]) if x.mnemonic == "mov"
                        and slot(x.op_str) is not None), None)
            stores = [slot(x.op_str.split(",")[0]) for x in body if x.mnemonic == "mov"
                      and "[rsp" in x.op_str.split(",")[0]]
            sm = next((s for s in stores if s is not None and s != ctr), None)
            if ctr is None or sm is None:
                return None
            return head, exit_, ctr, sm
    return None


INFO = {k: loop_info(v) for k, v in mods.items()}


def run(mid, b, state):
    code = mods[mid]
    uc = Uc(UC_ARCH_X86, UC_MODE_64)
    uc.mem_map(CODE, 0x10000)
    uc.mem_map(STATE, 0x1000)
    uc.mem_map(STACK, 0x10000)
    uc.mem_map(RET, 0x1000)
    uc.mem_write(CODE, code)
    uc.mem_write(STATE, bytes(state))
    sp = STACK + 0x8000
    uc.mem_write(sp, RET.to_bytes(8, "little"))
    uc.reg_write(UC_X86_REG_RSP, sp)
    uc.reg_write(UC_X86_REG_RCX, b)
    uc.reg_write(UC_X86_REG_RDX, STATE)
    uc.reg_write(UC_X86_REG_R8, 0)
    info = INFO[mid]

    def on_code(u, addr, size, _):
        if info and addr == CODE + info[0]:
            rsp = u.reg_read(UC_X86_REG_RSP)
            u.mem_write(rsp + info[2], (1000000).to_bytes(4, "little"))
            u.mem_write(rsp + info[3], (sum(range(1000000)) & 0xFFFFFFFF).to_bytes(4, "little"))
            u.reg_write(UC_X86_REG_RIP, CODE + info[1])

    def on_cpuid(u, _):
        u.reg_write(UC_X86_REG_RAX, 0x000306A9)
        u.reg_write(UC_X86_REG_RBX, 0)
        u.reg_write(UC_X86_REG_RCX, 0)
        u.reg_write(UC_X86_REG_RDX, 0x1FEBFBFF)          # SSE2 (bit 26) set, as on any x64 CPU
        return 1                                          # instruction handled

    if info:
        uc.hook_add(UC_HOOK_CODE, on_code, begin=CODE + info[0], end=CODE + info[0])
    uc.hook_add(UC_HOOK_INSN, on_cpuid, None, 1, 0, UC_X86_INS_CPUID)
    uc.emu_start(CODE, RET, count=20000)
    return uc.reg_read(UC_X86_REG_RAX) & 0xFFFFFFFF, bytearray(uc.mem_read(STATE, 32))


def walk(max_len=120):
    """Depth-first search over bytes; printable ASCII first, then the rest."""
    order = list(range(0x20, 0x7F)) + [b for b in range(256) if not 0x20 <= b < 0x7F]
    best = []

    def dfs(mid, state, path):
        if mid == FINAL:
            return path, state
        if len(path) >= max_len:
            return None
        for b in order:
            r, st = run(mid, b, state)
            if r in mods:
                if len(path) + 1 > len(best):
                    best[:] = path + [b]
                    print(f"  {len(path)+1:3}: {bytes(best)!r} -> 0x{r:x}", flush=True)
                res = dfs(r, st, path + [b])
                if res:
                    return res
        return None

    res = dfs(START, bytearray(32), [])
    return (bytes(res[0]), res[1]) if res else (bytes(best), None)


if __name__ == "__main__":
    path, state = walk()
    print("[+] input:", path)
    open(os.path.join(HERE, "loot", "module_path.txt"), "wb").write(path)
    if state is not None:
        open(os.path.join(HERE, "loot", "final_state.bin"), "wb").write(state)