#!/usr/bin/env python3
"""Build minimal PE analysis containers for Binary Ninja Cloud.

Binary Ninja Cloud accepts PE/ELF/Mach-O files, but not raw blobs.  These
containers add only a PE loader header around either the byte-exact original
blob or the reproducibly decrypted blob.  The payload is mapped at 0x402000 so
all addresses match chall.exe.
"""

from hashlib import sha256
from pathlib import Path
from struct import pack_into


HERE = Path(__file__).resolve().parent
OUT = HERE / "analysis-containers"
IMAGE_BASE = 0x400000
SECTION_RVA = 0x2000
FILE_ALIGNMENT = 0x200
SECTION_ALIGNMENT = 0x1000


def align(value: int, alignment: int) -> int:
    return (value + alignment - 1) & ~(alignment - 1)


def build_pe(payload: bytes, *, bits: int, entry_offset: int) -> bytes:
    """Wrap payload in a one-section PE without changing any payload byte."""
    if bits not in (32, 64):
        raise ValueError("bits must be 32 or 64")
    if not 0 <= entry_offset < len(payload):
        raise ValueError("entry offset outside payload")

    optional_size = 0xE0 if bits == 32 else 0xF0
    raw_size = align(len(payload), FILE_ALIGNMENT)
    headers_size = FILE_ALIGNMENT
    image_size = align(SECTION_RVA + len(payload), SECTION_ALIGNMENT)

    image = bytearray(headers_size + raw_size)
    image[:2] = b"MZ"
    pack_into("<I", image, 0x3C, 0x80)

    pe = 0x80
    image[pe : pe + 4] = b"PE\0\0"
    machine = 0x14C if bits == 32 else 0x8664
    characteristics = 0x0102 if bits == 32 else 0x0022
    pack_into(
        "<HHIIIHH",
        image,
        pe + 4,
        machine,
        1,
        0,
        0,
        0,
        optional_size,
        characteristics,
    )

    opt = pe + 24
    if bits == 32:
        pack_into("<HBB", image, opt, 0x10B, 14, 0)
        pack_into("<III", image, opt + 4, raw_size, 0, 0)
        pack_into(
            "<III",
            image,
            opt + 16,
            SECTION_RVA + entry_offset,
            SECTION_RVA,
            SECTION_RVA,
        )
        pack_into("<I", image, opt + 28, IMAGE_BASE)
        common = opt + 32
    else:
        pack_into("<HBB", image, opt, 0x20B, 14, 0)
        pack_into("<III", image, opt + 4, raw_size, 0, 0)
        pack_into("<II", image, opt + 16, SECTION_RVA + entry_offset, SECTION_RVA)
        pack_into("<Q", image, opt + 24, IMAGE_BASE)
        common = opt + 32

    pack_into("<II", image, common, SECTION_ALIGNMENT, FILE_ALIGNMENT)
    pack_into("<HHHHHH", image, common + 8, 6, 0, 0, 0, 6, 0)
    pack_into("<I", image, common + 20, 0)
    pack_into("<II", image, common + 24, image_size, headers_size)
    pack_into("<IHH", image, common + 32, 0, 3, 0x8160)

    if bits == 32:
        pack_into("<IIIIII", image, common + 40, 0x100000, 0x1000, 0x100000, 0x1000, 0, 16)
    else:
        pack_into("<QQQQII", image, common + 40, 0x100000, 0x1000, 0x100000, 0x1000, 0, 16)

    section = opt + optional_size
    image[section : section + 8] = b".text\0\0\0"
    pack_into(
        "<IIIIIIHHI",
        image,
        section + 8,
        len(payload),
        SECTION_RVA,
        raw_size,
        headers_size,
        0,
        0,
        0,
        0,
        0xE0000020,  # code | execute | read | write
    )
    image[headers_size : headers_size + len(payload)] = payload
    assert image[headers_size : headers_size + len(payload)] == payload
    return bytes(image)


def main() -> None:
    original = (HERE / "bitbender.bin").read_bytes()
    decrypted = (HERE / "bitbender.decrypted.bin").read_bytes()
    assert len(original) == len(decrypted) == 0x2A2
    assert (HERE / "chall.exe").read_bytes()[0x600 : 0x600 + 0x2A2] == original
    assert decrypted[0xD4:0x196] == (HERE / "stage2.bin").read_bytes()
    assert decrypted[0x196:0x230] == (HERE / "stage3.bin").read_bytes()
    assert decrypted[0x230:0x282] == (HERE / "stage4.bin").read_bytes()

    OUT.mkdir(exist_ok=True)
    targets = {
        "02-original-blob-x86-bootstrap.exe": (original, 32, 0x000),
        "03-original-blob-x64-stage1.exe": (original, 64, 0x017),
        "04-decrypted-blob-x86-stage2.exe": (decrypted, 32, 0x0D4),
        "05-decrypted-blob-x64-stage3.exe": (decrypted, 64, 0x196),
        "06-decrypted-blob-x86-stage4.exe": (decrypted, 32, 0x230),
    }

    for name, (payload, bits, entry) in targets.items():
        path = OUT / name
        data = build_pe(payload, bits=bits, entry_offset=entry)
        path.write_bytes(data)
        assert data[FILE_ALIGNMENT : FILE_ALIGNMENT + len(payload)] == payload
        print(f"{sha256(data).hexdigest()}  {path.relative_to(HERE)}")


if __name__ == "__main__":
    main()