#!/usr/bin/env python3
"""FLARE-On 13 #2 GhostStream - static solver (nothing from the challenge is executed).

Usage: python3 -I solve.py [GhostStream.wim]
Needs: pefile (pip install pefile) and the `7z` CLI (p7zip) to unpack the WIM with its ADS.

Chain (all recovered with Ghidra, see writeup.md):
  GhostStream.exe (parent)
    - reads <dir>\\PolyjuicePotion.txt, or, only if the debugger-only global
      DAT_140006668 == 1, the ADS "%sPolyjuicePotion.txt:LordVoldemort"
    - creates a window; its WndProc (WM_CREATE) transforms that buffer only if
      the process age from GetProcessTimes is > 0x2760 = 10080 minutes (7 days):
          buf[i] ^= WM_CREATE ^ previous_msg (WM_NCCALCSIZE) = 0x01 ^ 0x83
          buf[0:15] ^= 8b 8f 92 85 88 96 98 9b 94 8b 95 e6 ed f0 e7
    - drops resource BIN/101 as %TEMP%\\Partner_CTF.exe and writes the buffer to
      \\\\.\\pipe\\FlareChallenge
  Partner_CTF.exe (child)
    - key = pipe data cut at the first CR/LF; RC4 KSA(key)
    - hot-patches ucrtbase!strlen with: buf[k] = (buf[k]+k+1)^0x13; return 19
      and calls it on a 19-byte table D (decodes it in place to
      "Piertotum Locomotor") and on "Wingardium Leviosa" (result = L)
    - PRGA with i += L instead of i += 1:
          out[k] = S[(S[i]+S[j])&0xff] ^ D[k % len(D)] ^ ((L + ct[k]) & 0xff)
      Statically L reads as 19, but only L = 18 (the true length of
      "Wingardium Leviosa") decrypts, so both are tried and the binary's own
      "@flare-on.com" check picks the right one.
    - prints the flag if the output contains "@flare-on.com"
"""
import os
import subprocess
import sys
import tempfile

import pefile

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_WIM = os.path.join(HERE, "loot", "extracted", "GhostStream.wim")
SUFFIX = b"@flare-on.com"

# ---------------------------------------------------------------- parent side
PARENT_NAME = "GhostStream.exe"
WM_CREATE, WM_NCCALCSIZE = 0x01, 0x83
WNDPROC_XOR = bytes.fromhex("8b8f92858896989b948b95e6edf0e7")


def parent_strdec(enc, name=PARENT_NAME):
    """FUN_140001130: string decoder keyed on the running image's file name."""
    k = ord(name[len(name) >> 1]) ^ ord(name[0])
    mode = {0x1c: 0, 0xe5: 0, 0x24: 2, 0x62: 2, 0xe1: 2,
            0x53: 1, 0x60: 1, 0x6a: 1, 0x74: 1, 0x7b: 1}.get(sum(enc) & 0xff)
    if mode == 0:
        return bytes((b - k) & 0xff for b in enc)
    if mode == 1:
        return bytes(((b + 1) & 0xff) ^ k for b in enc)
    if mode == 2:
        return bytes((b - 2 + k) & 0xff for b in enc)
    return None


def le(*parts):
    return b"".join(v.to_bytes(n, "little") for v, n in parts)


def wndproc_transform(buf, prev_msg=WM_NCCALCSIZE):
    """FUN_140001b20, WM_CREATE branch taken when process age > 7 days."""
    b = bytearray(buf)
    x = WM_CREATE ^ prev_msg
    i = 0
    while i < len(b.split(b"\0")[0]):     # strlen re-evaluated every iteration
        b[i] ^= x
        i += 1
    for i, c in enumerate(WNDPROC_XOR):
        b[i] ^= c
    return bytes(b)


# --------------------------------------------------------------- partner side
def va_read(pe, va, n):
    return pe.get_data(va - pe.OPTIONAL_HEADER.ImageBase, n)


def patched_strlen(buf):
    """Code written over ucrtbase!strlen: 31 c0 04 13 b2 01 75 01 [e8] 00 11 30 01 fe c2 48 ff c1 38 21 75 f3 c3
         xor eax,eax; add al,0x13; mov dl,1
       L: add [rcx],dl; xor [rcx],al; inc dl; inc rcx; cmp [rcx],ah; jne L; ret"""
    b = bytearray(buf)
    k = 0
    while True:
        b[k] = ((b[k] + (k + 1)) & 0xff) ^ 0x13
        k += 1
        if b[k] == 0:
            break
    return 0x13, bytes(b)


def partner_decrypt(key, partner_pe, L=None):
    ct = le((0x8b397ad735875c96, 8), (0xc85856bb286d0475, 8),
            (0x37a956a4e1daee35, 8), (0x0d340a27f9561581, 8),
            (0x262c82ee, 4), (0xf8a5, 2))                      # 38 bytes
    key = key.split(b"\r")[0].split(b"\n")[0]               # strcspn(buf, "\r\n")
    key = key.split(b"\0")[0]
    if not key:
        return None
    S = list(range(256))
    j = 0
    for i in range(256):
        j = (j + key[i % len(key)] + S[i]) & 0xff
        S[i], S[j] = S[j], S[i]

    table = va_read(partner_pe, 0x140005078, 0x40)
    _, table = patched_strlen(table)                         # strlen(&DAT_140005078)
    if L is None:
        L, _ = patched_strlen(b"Wingardium Leviosa\0")      # static reading: 19
    D = table.split(b"\0")[0]

    out = bytearray()
    i = j = 0
    for k in range(len(ct)):
        i = (i + L) & 0xff
        t = S[i]
        j = (j + t) & 0xff
        S[i], S[j] = S[j], t
        out.append(S[(S[i] + t) & 0xff] ^ D[k % len(D)] ^ ((L + ct[k]) & 0xff))
    return bytes(out)


# ----------------------------------------------------------------------- main
def main():
    wim = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_WIM
    out_dir = os.path.join(HERE, "loot", "wim_contents")
    if not os.path.isdir(out_dir):
        subprocess.run(["7z", "x", "-sns", "-y", "-o" + out_dir, wim],
                       check=True, stdout=subprocess.DEVNULL)
    gdir = os.path.join(out_dir, "GhostStream")
    parent_path = os.path.join(gdir, "GhostStream.exe")
    main_stream = open(os.path.join(gdir, "PolyjuicePotion.txt"), "rb").read()
    ads = open(os.path.join(gdir, "PolyjuicePotion.txt:LordVoldemort"), "rb").read()

    # show the path string the debugger-only branch selects
    enc = le((0x3c1d40f2, 4), (0x42374639, 4), (0x1d323036, 4), (0x3c36413c, 4),
             (0x4541fb3b, 4), (0x3c190741, 4), (0x3c23313f, 4), (0x3a323139, 4),
             (0x3f3c, 2), (0x41, 1))
    print("[*] DAT_140006668==1 path format :", parent_strdec(enc).decode())

    # carve the child from the parent's resources (BIN / 101)
    ppe = pefile.PE(parent_path)
    child = None
    for t in ppe.DIRECTORY_ENTRY_RESOURCE.entries:
        if str(t.name) == "BIN":
            d = t.directory.entries[0].directory.entries[0].data.struct
            child = ppe.get_data(d.OffsetToData, d.Size)
    child_path = os.path.join(HERE, "loot", "Partner_CTF.bin")
    with open(child_path, "wb") as f:
        f.write(child)
    print(f"[*] carved Partner_CTF.exe ({len(child)} bytes) -> loot/Partner_CTF.bin")
    cpe = pefile.PE(data=child)

    candidates = {
        "main stream, no 7-day transform": main_stream[:0xff],
        "main stream, 7-day transform": wndproc_transform(main_stream[:0xff]),
        "ADS LordVoldemort, no 7-day transform": ads,
        "ADS LordVoldemort, 7-day transform": wndproc_transform(ads),
    }
    flag = None
    # L (PRGA i-step and ciphertext addend) is the return value of the strlen
    # call on "Wingardium Leviosa": 19 if the hook is hit, 18 (true length) if not.
    for label, key in candidates.items():
        for L in (19, len("Wingardium Leviosa")):
            pt = partner_decrypt(key, cpe, L)
            ok = pt is not None and SUFFIX in pt
            shown = key.split(b"\r")[0]
            print(f"[{'+' if ok else '-'}] {label:40} L={L} key={shown!r}")
            if ok:
                print(f"      -> {pt!r}")
                flag = pt.split(b"\0")[0].decode()
    if not flag:
        sys.exit("[!] no candidate produced a flag")
    print(f"[+] FLAG: {flag}")
    with open(os.path.join(HERE, "flag.txt"), "w") as f:
        f.write(flag + "\n")


if __name__ == "__main__":
    main()