Flare-on 12

# -*- coding: utf-8 -*-
# Recursive Ghidra state analyzer with path tracking and file output
# Only prints & saves discovered paths of EXACTLY 16 characters

from ghidra.util.task import ConsoleTaskMonitor
from ghidra.program.model.scalar import Scalar
from collections import deque
import os

# ---------------- CONFIG ----------------
table_addr         = 0x140C687B8   # jump table base (4-byte entries)
start_state        = 0
max_state_value    = 0x100000
scan_limit_bytes   = 0xA00
branch_scan        = 400
jump_lookahead     = 80
stub_jump_window   = 0x40
near_jump_max      = 0x200
ascii_only         = True
DEST_SLOT_DISP     = 0x58D30      # expected RSP displacement for the state MOV
terminal_func_names = ("FUN_140c685ee", "thunk_FUN_140c685ee")
output_filename    = "chal5_paths_16.txt"   # saved in user's home directory
TARGET_LENGTH      = 16               # only output/save paths of this exact length
# ---------------------------------------

monitor   = ConsoleTaskMonitor()
listing   = currentProgram.getListing()
memory    = currentProgram.getMemory()
base_img  = currentProgram.getImageBase().getOffset()
base_high = base_img & 0xFFFFFFFF00000000

# ---------- helpers ----------
def to_printable(b):
    try:
        return "'%s'" % chr(b) if 32 <= b <= 126 else ''
    except:
        return ''

def read_u32(addr):
    try:
        return getInt(toAddr(addr)) & 0xFFFFFFFF
    except:
        return None

def is_mapped_addr(va):
    try:
        return memory.contains(toAddr(va))
    except:
        return False

def ensure_code_at(va):
    """Disassemble at 'va' if nothing is present there yet (no bulk disassembly)."""
    if not is_mapped_addr(va):
        return
    if listing.getInstructionAt(toAddr(va)):
        return
    try:
        disassemble(toAddr(va))
    except:
        pass

def next_ins(addr):
    ins = listing.getInstructionAt(addr)
    return ins if ins else listing.getInstructionAfter(addr)

def has_rsp(opobjs):
    for o in opobjs:
        try:
            if o.__class__.__name__ == "Register" and o.getName().upper() == "RSP":
                return True
        except:
            pass
    return False

def scalar_sum(opobjs):
    total = 0
    for o in opobjs:
        if isinstance(o, Scalar):
            total += o.getValue()
        else:
            try:
                total += o.getDisplacement()
            except:
                pass
    return total

def mem_disp(ins, idx):
    """Return displacement if operand idx is [RSP+disp], else None."""
    try:
        objs = ins.getOpObjects(idx)
        if not objs or not has_rsp(objs):
            return None
        return scalar_sum(objs)
    except:
        return None

# ---------- dest decoding ----------
def choose_dest_from_entry(rel32, entry_addr):
    cand1 = (base_high | (rel32 & 0xFFFFFFFF))
    if is_mapped_addr(cand1):
        return (cand1, "baseHigh|rel32")
    rels = rel32 - 0x100000000 if (rel32 & 0x80000000) else rel32
    cand2 = base_img + rels
    if is_mapped_addr(cand2):
        return (cand2, "base+rel32(signed)")
    return (None, None)

# ---------- stub follower ----------
def follow_stub(dest):
    ensure_code_at(dest)
    ins = next_ins(toAddr(dest))
    if not ins:
        return dest
    scanned = 0
    saw_real = False
    while ins and scanned < stub_jump_window:
        m = ins.getMnemonicString().upper()
        if m in ("CMP", "MOV", "LEA", "CALL", "TEST", "XOR", "ADD", "SUB"):
            saw_real = True
        if m == "JMP" and not saw_real:
            for r in ins.getReferencesFrom():
                tgt = r.getToAddress().getOffset()
                if dest < tgt <= dest + near_jump_max and is_mapped_addr(tgt):
                    print("  [stub] following JMP 0x%X -> 0x%X" % (dest, tgt))
                    return tgt
        scanned += ins.getLength()
        ins = ins.getNext()
    return dest

# ---------- terminal jmp detection ----------
def is_terminal_jmp(ins):
    if not ins or ins.getMnemonicString().upper() != "JMP":
        return False
    refs = list(ins.getReferencesFrom())
    for r in refs:
        fn = getFunctionContaining(r.getToAddress())
        if fn:
            name = fn.getName().lower()
            if any(t.lower() in name for t in terminal_func_names):
                return True
    try:
        op = ins.getOpObjects(0)
        if op and isinstance(op[0], Scalar):
            rel = op[0].getSignedValue()
            tgt = ins.getAddress().getOffset() + ins.getLength() + rel
            fn = getFunctionContaining(toAddr(tgt))
            if fn:
                name = fn.getName().lower()
                if any(t.lower() in name for t in terminal_func_names):
                    return True
    except:
        pass
    return False

# ---------- analyze a single state ----------
def analyze_state_body(dest):
    results = []
    ensure_code_at(dest)
    start = next_ins(toAddr(dest))
    if not start:
        return results
    end_addr = dest + scan_limit_bytes
    ins = start
    while ins and ins.getAddress().getOffset() < end_addr:
        mnem = ins.getMnemonicString().upper()
        if mnem == "CMP":
            cmp_disp = mem_disp(ins, 0)
            if cmp_disp is None:
                ins = ins.getNext()
                continue
            imm_objs = ins.getOpObjects(1)
            if not imm_objs or not isinstance(imm_objs[0], Scalar):
                ins = ins.getNext()
                continue
            imm8 = imm_objs[0].getValue() & 0xFF
            if ascii_only and not (32 <= imm8 <= 126):
                ins = ins.getNext()
                continue
            look = ins.getNext()
            jz_refs = []
            steps = 0
            while look and steps < jump_lookahead:
                lm = look.getMnemonicString().upper()
                if lm in ("JZ", "JE"):
                    jz_refs.extend(list(look.getReferencesFrom()))
                if lm == "RET" or is_terminal_jmp(look):
                    break
                look = look.getNext()
                steps += 1
            for ref in jz_refs:
                tgt_va = ref.getToAddress().getOffset()
                ensure_code_at(tgt_va)
                cur = listing.getInstructionAt(toAddr(tgt_va)) or listing.getInstructionAfter(toAddr(tgt_va))
                inner = 0
                while cur and inner < branch_scan:
                    mm = cur.getMnemonicString().upper()
                    if mm == "MOV":
                        mov_disp = mem_disp(cur, 0)
                        if mov_disp == DEST_SLOT_DISP:
                            op1 = cur.getOpObjects(1)
                            if op1 and isinstance(op1[0], Scalar):
                                imm64 = op1[0].getValue() & 0xFFFFFFFFFFFFFFFF
                                results.append({
                                    "cmp_addr": ins.getAddress(),
                                    "cmp_disp": cmp_disp,
                                    "ch": imm8,
                                    "mov_addr": cur.getAddress(),
                                    "mov_disp": mov_disp,
                                    "val": imm64
                                })
                                break
                    if mm == "RET" or is_terminal_jmp(cur):
                        break
                    cur = cur.getNext()
                    inner += 1
        if mnem == "RET" or is_terminal_jmp(ins):
            break
        ins = ins.getNext()
    uniq = {}
    for r in results:
        key = (r["cmp_addr"], r["ch"])
        if key not in uniq:
            uniq[key] = r
    return list(uniq.values())

# ---------- BFS traversal with path tracking ----------
visited = set()
queue   = deque([start_state])
edges   = {}
state_to_dest = {}
paths   = {start_state: [""]}  # state -> list of prefix strings

print("[*] Starting BFS with path tracking from state %d" % start_state)

while queue:
    s = queue.popleft()
    if s in visited:
        continue
    visited.add(s)

    entry = table_addr + (s * 4)
    raw   = read_u32(entry)
    if raw is None or raw == 0:
        print("State %d: invalid or empty table entry at 0x%X" % (s, entry))
        continue

    dest, method = choose_dest_from_entry(raw, entry)
    if dest is None:
        print("State %d: could not resolve entry (raw=0x%X)" % (s, raw))
        continue

    ensure_code_at(dest)
    real_dest = follow_stub(dest)
    state_to_dest[s] = real_dest

    print("\n==============================")
    print("State %d  entry @ 0x%X (raw=0x%X) -> dest 0x%X  [%s]" %
          (s, entry, raw, real_dest, method))

    found = analyze_state_body(real_dest)
    if not found:
        print("  (no transitions found)")
        edges[s] = []
        continue

    edges[s] = []
    for r in found:
        ch = r["ch"]
        val = r["val"]
        next_state = None
        if 0 <= val <= max_state_value:
            next_state = int(val)
        edges[s].append((ch, next_state))
        print("  '%s' (0x%02X) -> state %s" % (chr(ch), ch, str(next_state)))

        if next_state is not None:
            prev_paths = paths.get(s, [""])
            new_paths = []
            for prefix in prev_paths:
                # build new path by appending this character
                new_paths.append(prefix + chr(ch))
            # append new paths to next_state's list (allow duplicates here; we'll dedupe later)
            paths.setdefault(next_state, []).extend(new_paths)
            if next_state not in visited and next_state not in queue:
                queue.append(next_state)

# ---------- extract 16-char paths, dedupe, print & save ----------
# gather all sequences from the 'paths' mapping
all_sequences = []
for seq_list in paths.values():
    all_sequences.extend(seq_list)

# filter exact length and unique
final_paths = sorted(set([s for s in all_sequences if len(s) == TARGET_LENGTH]))

print("\n\n=== Discovered paths of length %d ===" % TARGET_LENGTH)
if not final_paths:
    print("(no %d-character paths found)" % TARGET_LENGTH)
else:
    for p in final_paths:
        print(p)

# Save results to file in user's home directory
out_path = os.path.join(os.path.expanduser("~"), output_filename)
try:
    with open(out_path, "w", encoding="utf-8") as f:
        if not final_paths:
            f.write("(no %d-character paths found)\n" % TARGET_LENGTH)
        else:
            for p in final_paths:
                f.write(p + "\n")
    print("[+] Saved %d paths to: %s" % (len(final_paths), out_path))
except Exception as e:
    print("[-] Failed to write output file: %s" % str(e))

print("\nDone. Visited %d states." % len(visited))