#!/usr/bin/env python3
"""FLARE-On 13 #3 FlareOn13.doc - static solver.

Usage: python3 -I solve.py [path/to/flareon13.doc]
Needs: pikepdf (pip install pikepdf), and `mutool` (mupdf-tools) to dump the JBIG2 image.

flareon13.doc is a polyglot:
  EICAR string | PDF 1.7 (RC4-40 encrypted, empty user password) wrapping
  a raw 2352-byte/sector Mode-1 CD image with a UDF FS (README.TXT, click_me
  Mach-O decoy) | encrypted ZIP (ZipCrypto + method 2 "Reduce, factor 1")
  holding flag.txt | VHD footer ("conectix", doubling as the ZIP comment).
Six layer phrases feed the Crystal checker in the Mach-O (see writeup.md).
Neither 7-Zip nor Python's zipfile supports Reduce, so it is implemented here.
"""
import os
import struct
import subprocess
import sys
import zlib

import pikepdf

HERE = os.path.dirname(os.path.abspath(__file__))
LOOT = os.path.join(HERE, "loot")
DEFAULT_DOC = os.path.join(LOOT, "extracted", "flareon13.doc")
PASSWORD = b"jbig2_fax_geometry"   # read from the decrypted PDF's JBIG2 image (part 2)
# ZipCrypto internal keys recovered with bkcrack (known plaintext = the all-zero
# Reduce follower-set counts for bytes 0xFF..0xD8 at the start of the stream):
#   bkcrack -C flareon13.doc -c flag.txt -p <40 zero bytes> -o 0
ZIP_KEYS = (0x0917ED61, 0xCFC8F696, 0x5D3E9072)
CRYSTAL = os.path.join(LOOT, "carved", "macho", "click_me~.arm64-3")


# ------------------------------------------------------------------ ZipCrypto
CRC_TABLE = [0] * 256
for _n in range(256):
    _c = _n
    for _ in range(8):
        _c = (_c >> 1) ^ 0xEDB88320 if _c & 1 else _c >> 1
    CRC_TABLE[_n] = _c


class ZipCrypto:
    def __init__(self, pw):
        self.k = [0x12345678, 0x23456789, 0x34567890]
        for b in pw:
            self._update(b)

    def _update(self, b):
        k0, k1, k2 = self.k
        k0 = (k0 >> 8) ^ CRC_TABLE[(k0 ^ b) & 0xff]
        k1 = ((k1 + (k0 & 0xff)) * 134775813 + 1) & 0xffffffff
        k2 = (k2 >> 8) ^ CRC_TABLE[(k2 ^ (k1 >> 24)) & 0xff]
        self.k = [k0, k1, k2]

    def decrypt(self, data):
        out = bytearray()
        for b in data:
            t = (self.k[2] | 2) & 0xffff
            c = b ^ (((t * (t ^ 1)) >> 8) & 0xff)
            self._update(c)
            out.append(c)
        return bytes(out)


# --------------------------------------------------------------- Unreduce (2-5)
class BitReader:
    """PKWARE bit order: least-significant bit first."""
    def __init__(self, data):
        self.data, self.pos, self.buf, self.n = data, 0, 0, 0

    def bits(self, n):
        while self.n < n:
            byte = self.data[self.pos] if self.pos < len(self.data) else 0
            self.pos += 1
            self.buf |= byte << self.n
            self.n += 8
        v = self.buf & ((1 << n) - 1)
        self.buf >>= n
        self.n -= n
        return v


def b_table(n):
    """Bits needed to index a follower set of size n (Info-ZIP B_table)."""
    return max(1, (n - 1).bit_length())


def unreduce(data, out_size, factor):
    br = BitReader(data)
    followers = [[] for _ in range(256)]
    for j in range(255, -1, -1):
        n = br.bits(6)
        followers[j] = [br.bits(8) for _ in range(n)]

    out = bytearray()
    mask = (1 << (8 - factor)) - 1          # L(X) = X & mask
    last = 0
    state = 0
    V = length = 0
    DLE = 144
    while len(out) < out_size:
        # 1) probabilistic decode of the next byte
        if not followers[last]:
            c = br.bits(8)
        elif br.bits(1):
            c = br.bits(8)
        else:
            c = followers[last][br.bits(b_table(len(followers[last])))]
        last = c
        # 2) DLE run-length / back-reference expansion
        if state == 0:
            if c != DLE:
                out.append(c)
            else:
                state = 1
        elif state == 1:
            if c:
                V = c
                length = V & mask
                state = 2 if length == mask else 3
            else:
                out.append(DLE)
                state = 0
        elif state == 2:
            length += c
            state = 3
        else:  # state 3
            dist = ((V >> (8 - factor)) << 8) + c + 1
            for _ in range(length + 3):
                out.append(out[-dist] if dist <= len(out) else 0)
            state = 0
    return bytes(out[:out_size])


# ------------------------------------------------------------------- helpers
def raw_cd_to_iso(blob):
    """Strip 2352-byte raw Mode-1 sectors down to their 2048-byte user data."""
    sync = b"\x00" + b"\xff" * 10 + b"\x00"
    iso = bytearray()
    for s in range(len(blob) // 2352):
        sec = blob[s * 2352:(s + 1) * 2352]
        iso += sec[16:16 + 2048] if sec[:12] == sync else b"\0" * 2048
    return bytes(iso)


def read_zip_entry(blob):
    eocd = blob.rfind(b"PK\x05\x06")
    cd_size, cd_off = struct.unpack("<II", blob[eocd + 12:eocd + 20])
    cd = blob[cd_off:cd_off + cd_size]
    (_, _, _, flags, method, mtime, _, crc, csize, usize, nlen, xlen, clen,
     _, _, _, lho) = struct.unpack("<IHHHHHHIIIHHHHHII", cd[:46])
    name = cd[46:46 + nlen].decode()
    nlen2, xlen2 = struct.unpack("<HH", blob[lho + 26:lho + 30])
    data = blob[lho + 30 + nlen2 + xlen2: lho + 30 + nlen2 + xlen2 + csize]
    return dict(name=name, flags=flags, method=method, mtime=mtime, crc=crc,
                csize=csize, usize=usize, data=data, comment=blob[eocd + 22:])


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_DOC
    blob = open(path, "rb").read()
    os.makedirs(os.path.join(LOOT, "pdf"), exist_ok=True)
    os.makedirs(os.path.join(LOOT, "carved"), exist_ok=True)

    print("[*] header         :", blob[:0x44].decode("latin1"))
    print("[*] VHD footer     :", blob[-512:-504], "creator", blob[-512 + 28:-512 + 32])

    # Layer: raw CD image -> UDF (shown for completeness; it only holds decoys)
    iso = raw_cd_to_iso(blob)
    with open(os.path.join(LOOT, "carved", "cooked.iso"), "wb") as f:
        f.write(iso)
    print(f"[*] raw CD image   : {len(blob)//2352} sectors -> loot/carved/cooked.iso (UDF)")

    # Layer: PDF (RC4-40, empty user password)
    pdf = pikepdf.open(path, password="")
    print(f"[*] PDF encryption : R={pdf.encryption.R} V={pdf.encryption.V} "
          f"bits={pdf.encryption.bits} key={pdf.encryption.encryption_key.hex()}")
    pdf.save(os.path.join(LOOT, "pdf", "decrypted.pdf"))
    im = pdf.pages[0].obj.Resources.XObject.Im
    print(f"[*] PDF image      : {int(im.Width)}x{int(im.Height)} {im.Filter}")
    try:
        subprocess.run(["mutool", "extract", os.path.join(LOOT, "pdf", "decrypted.pdf")],
                       cwd=os.path.join(LOOT, "pdf"), check=True,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        print("[*] image dumped to loot/pdf/image-*.png -> reads:", PASSWORD.decode())
    except (OSError, subprocess.CalledProcessError):
        print("[!] mutool not available; password taken from write-up:", PASSWORD.decode())

    # Layer: ZIP (ZipCrypto + Reduce factor 1)
    z = read_zip_entry(blob)
    print(f"[*] ZIP entry      : {z['name']} method={z['method']} flags={z['flags']:#x} "
          f"csize={z['csize']} usize={z['usize']} crc={z['crc']:08x} comment={z['comment']!r}")
    zc = ZipCrypto(b"")
    zc.k = list(ZIP_KEYS)
    plain = zc.decrypt(z["data"])
    expect = (z["mtime"] >> 8) if z["flags"] & 8 else (z["crc"] >> 24)
    print(f"[*] ZipCrypto hdr  : {plain[:12].hex()} (check {plain[11]:#04x}, expected {expect:#04x})")
    assert plain[11] == expect
    part4 = unreduce(plain[12:], z["usize"], z["method"] - 1)
    assert zlib.crc32(part4) == z["crc"], "CRC mismatch"
    print(f"[+] flag.txt (part 4): {part4.decode()}")

    parts = [
        com_stub_output(blob),                                   # 0 DOS COM / EICAR
        bytes.fromhex(HIDDEN_TEXT_HEX).decode("cp037"),          # 1 PDF invisible text
        PASSWORD.decode(),                                       # 2 PDF JBIG2 image
        udf_lvinfo(iso),                                         # 3 UDF IUVD LVInfo
        part4.decode(),                                          # 4 ZIP flag.txt
        odin_string(os.path.join(LOOT, "carved", "macho", "click_me~.arm64")),  # 5 Odin
    ]
    flag = crystal_decrypt(parts, CRYSTAL)
    print(f"[+] FLAG: {flag}")
    with open(os.path.join(HERE, "flag.txt"), "w") as f:
        f.write(flag + "\n")


# ------------------------------------------------------------ phrase extractors
HIDDEN_TEXT_HEX = "9481899586998194856d8582838489836d878896a2a3"   # 3 Tr text in page stream


def com_stub_output(blob):
    """The alphanumeric COM stub XORs 'EICAR' (0x160..0x164) with 03 05 02 13 17,
    patches int 21h/int 20h, then prints the $-terminated string at 0x160."""
    b = bytearray(blob[:0x100])
    for off, x in zip(range(0x60, 0x65), (0x03, 0x05, 0x02, 0x13, 0x17)):
        b[off] ^= x
    return bytes(b[0x60:b.index(b"$", 0x60)]).decode()


def udf_lvinfo(iso):
    """Implementation Use Volume Descriptor (sector 21), LVInfo1 dstring."""
    sec = iso[21 * 2048:22 * 2048]
    i = sec.index(b"udf_")
    return sec[i:sec.index(b"\0", i)].decode()


def odin_string(path):
    """Emulates the Odin slice: SHA-0-style compression of a constant block, hex
    digits of a,b,c XORed with a mask -> 19-byte UTF-8 string."""
    d = open(path, "rb").read()
    rd = lambda a, n: d[a - 0x100000000:a - 0x100000000 + n]
    blk = (rd(0x100001270, 8) + rd(0x100001278, 8) + rd(0x100001260, 8) +
           rd(0x100001268, 8) + b"\0" * 16 + rd(0x100001250, 8) + rd(0x100001258, 8))
    W = list(struct.unpack("<16I", blk))
    for t in range(16, 80):
        W.append(W[t - 8] ^ W[t - 3] ^ W[t - 14] ^ W[t - 16])        # no rotate
    M = 0xffffffff
    rol = lambda x, n: ((x << n) | (x >> (32 - n))) & M
    a, b, c, dd, e = 0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0
    for t in range(80):
        f, k = [((b & c) | (~b & dd), 0x5a827999), (b ^ c ^ dd, 0x6ed9eba1),
                ((b & c) | (b & dd) | (c & dd), 0x8f1bbcdc), (b ^ c ^ dd, 0xca62c1d6)][t // 20]
        a, b, c, dd, e = (rol(a, 5) + f + e + k + W[t]) & M, a, rol(b, 30), c, dd
    a, b, c = (a + 0x67452301) & M, (b + 0xefcdab89) & M, (c + 0x98badcfe) & M
    hx, mask = "0123456789abcdef", rd(0x100001280, 16)
    out = [ord(hx[(a >> (28 - 4 * i)) & 15]) ^ mask[i] for i in range(8)]
    out += [ord(hx[(b >> (28 - 4 * i)) & 15]) ^ mask[8 + i] for i in range(8)]
    out += [ord(hx[(c >> 28) & 15]) ^ 0x50, ord(hx[(c >> 24) & 15]) ^ 0x4a,
            ord(hx[(c >> 20) & 15]) ^ 0x07]
    return bytes(out).decode()


def crystal_decrypt(parts, path):
    """Crystal slice: each line's MD5 must be one of six stored digests; key =
    XOR of SHA-256(part); RC4(key) over the hex-decoded 80-nibble blob."""
    import hashlib
    d = open(path, "rb").read()
    rd = lambda a, n: d[a - 0x100000000:a - 0x100000000 + n]
    nib = lambda bs: "".join("0123456789abcdef"[x & 15] for x in bs)
    digests = [nib(rd(0x100035e30 + 0x20 * i, 32)) for i in range(6)]
    for p in parts:
        h = hashlib.md5(p.encode()).hexdigest()
        print(f"    part {digests.index(h) if h in digests else '?'}: {p}")
        assert h in digests
    key = bytes(32)
    for p in parts:
        key = bytes(x ^ y for x, y in zip(key, hashlib.sha256(p.encode()).digest()))
    S, j = list(range(256)), 0
    for i in range(256):
        j = (j + S[i] + key[i % 32]) & 0xff
        S[i], S[j] = S[j], S[i]
    ct = bytes.fromhex(nib(rd(0x100035ef0, 0x50)))
    i = j = 0
    out = bytearray()
    for c in ct:
        i = (i + 1) & 0xff
        j = (j + S[i]) & 0xff
        S[i], S[j] = S[j], S[i]
        out.append(S[(S[i] + S[j]) & 0xff] ^ c)
    return out.decode()


if __name__ == "__main__":
    main()