#!/usr/bin/env python3
"""FLARE-On 13 #5 catthief - static solver.

Usage: python3 -I solve.py [capture.pcapng] [catthief.exe]
Needs: pefile, tshark (for extract_http.py).

Protocol (recovered from catthief.exe, Rust, src\\compress.rs):
  * every message (both directions) is RC4-encrypted with a fresh RC4 state
    keyed by 25 bytes = (DAT_1400c9010[16] || 0x0407576c41004a6c LE || 0x07) ^ 0x33
    = "i11_b3_p1und3r1n_y3r_d474"  (same keystream for every message)
  * server -> client: plaintext file path to steal
  * client -> server: canonical-Huffman compressed data, as a bitstring
        [11-bit N][32-bit M][N bytes code-length table][M bytes code bits]
    packed MSB-first; first upload is the compressed word "READY".
"""
import os
import struct
import subprocess
import sys

import pefile

HERE = os.path.dirname(os.path.abspath(__file__))
LOOT = os.path.join(HERE, "loot")
PCAP = os.path.join(LOOT, "extracted", "catthief", "capture.pcapng")
EXE = os.path.join(LOOT, "extracted", "catthief", "catthief.exe")
FLAG = "th3r3_b3_tr34sure_1n_th4t_pc4p@flare-on.com"   # transcribed from catcatcat.jpg


def get_key(exe):
    pe = pefile.PE(exe)
    ib = pe.OPTIONAL_HEADER.ImageBase
    raw = pe.get_data(0x1400C9010 - ib, 16) + struct.pack("<Q", 0x0407576C41004A6C) + b"\x07"
    return bytes(b ^ 0x33 for b in raw)               # main @0x140003470 loop: ^= 0x33 x 0x19


def rc4(key, data):
    S, j = list(range(256)), 0
    for i in range(256):                             # FUN_1400052e2 (key index i % 25)
        j = (j + S[i] + key[i % len(key)]) & 0xff
        S[i], S[j] = S[j], S[i]
    out, i, j = bytearray(), 0, 0
    for c in data:                                   # FUN_140005281
        i = (i + 1) & 0xff
        j = (j + S[i]) & 0xff
        S[i], S[j] = S[j], S[i]
        out.append(c ^ S[(S[i] + S[j]) & 0xff])
    return bytes(out)


class Bits:
    def __init__(self, data):
        self.d, self.p = data, 0

    def get(self, n):
        v = 0
        for _ in range(n):
            v = (v << 1) | ((self.d[self.p >> 3] >> (7 - (self.p & 7))) & 1)
            self.p += 1
        return v


def huffman_decode(blob):
    br = Bits(blob)
    n = br.get(11)                                    # FUN_1400053fe(.., N, 0xb)
    m = br.get(32)                                    # FUN_1400053fe(.., M, 0x20)
    lengths = [br.get(8) for _ in range(n)]           # FUN_140005457(table, N)
    payload = bytes(br.get(8) for _ in range(m))      # FUN_140005457(bits, M)
    # canonical code: sort (len, symbol); code = (code + 1) << (len - prevlen)
    syms = sorted((l, s) for s, l in enumerate(lengths) if l)
    table, code, prev = {}, 0, None
    for l, s in syms:
        if prev is not None:
            code = (code + 1) << (l - prev)
        else:
            code = 0
        table[(l, code)] = s
        prev = l
    out = bytearray()
    bits = Bits(payload)
    cur, cl, total = 0, 0, len(payload) * 8
    maxlen = max(l for l, _ in syms)
    while bits.p < total:
        cur = (cur << 1) | bits.get(1)
        cl += 1
        if (cl, cur) in table:
            out.append(table[(cl, cur)])
            cur, cl = 0, 0
        elif cl > maxlen:
            raise ValueError("bad code")
    return n, m, lengths, bytes(out)


def main():
    pcap = sys.argv[1] if len(sys.argv) > 1 else PCAP
    exe = sys.argv[2] if len(sys.argv) > 2 else EXE
    key = get_key(exe)
    print(f"[*] RC4 key: {key.decode()}")
    hdir = os.path.join(LOOT, "http")
    subprocess.run([sys.executable, "-I", os.path.join(HERE, "extract_http.py"), pcap, hdir],
                   check=True, stdout=subprocess.DEVNULL)
    outdir = os.path.join(LOOT, "recovered")
    os.makedirs(outdir, exist_ok=True)
    s, requested = 0, "READY.txt"          # stream 0 is the beacon; stream k uploads reply k-1
    while os.path.exists(os.path.join(hdir, f"stream{s}_req.bin")):
        req = rc4(key, open(os.path.join(hdir, f"stream{s}_req.bin"), "rb").read())
        resp = rc4(key, open(os.path.join(hdir, f"stream{s}_resp.bin"), "rb").read())
        n, m, lengths, data = huffman_decode(req)
        name = requested.split("\\")[-1]
        if name.endswith(".jpg"):                     # drop Huffman pad symbols after EOI
            data = data[:data.rfind(b"\xff\xd9") + 2]
        open(os.path.join(outdir, name), "wb").write(data)
        print(f"[*] stream {s}: upload N={n} M={m} -> {name} ({len(data)} bytes)")
        print(f"             server reply (next file to steal): {resp.decode(errors='replace')!r}")
        requested = resp.decode(errors="replace")
        s += 1
    # the flag is written on the scroll in catcatcat.jpg
    try:
        from PIL import Image
        im = Image.open(os.path.join(outdir, "catcatcat.jpg"))
        k = im.size[0] / 1400
        im.crop((int(860 * k), int(530 * k), int(1200 * k), int(690 * k))).save(
            os.path.join(LOOT, "scroll_crop.png"))
        print("[+] flag text is on the scroll in loot/recovered/catcatcat.jpg (crop: loot/scroll_crop.png)")
    except ImportError:
        print("[+] open loot/recovered/catcatcat.jpg - the flag is written on the scroll")
    print(f"[+] FLAG (read from the image): {FLAG}")
    with open(os.path.join(HERE, "flag.txt"), "w") as f:
        f.write(FLAG + "\n")


if __name__ == "__main__":
    main()