#!/usr/bin/env python3
"""
demo_fuzzer.py — boofuzz script for VulnStore (class demo)
boofuzz 0.4.2

VulnStore protocol (port 4141)
──────────────────────────────
  Banner:  +VulnStore 1.0 (demo target — port 4141)\r\n
  SET <key> <value>   →  +OK | -ERR
  GET <key>           →  +<value> | -ERR
  DEL <key>           →  +OK | -ERR
  LEN <key>           →  +<n> | -ERR
  APPEND <key> <data> →  +OK | -ERR
  KEYS                →  +key1,key2,...
  RESET               →  +OK
  INFO                →  +VulnStore/1.0 ...
  QUIT                →  +BYE

Vulnerabilities to find
───────────────────────
  SET    — stack buffer overflow (value_buf[64] + strcpy)
  LEN    — signed/unsigned confusion (atoi result compared to size_t)
  APPEND — integer truncation → heap overflow (existing+new length overflows int)

Demo script features
────────────────────
  s_group()    — cycles through a fixed list of values (boundary lengths for
                 the SET overflow — reaches the crash quickly in class)
  s_string()   — generates hundreds of mutations (long strings, format
                 specifiers, null bytes, special characters)
  pre/post-send callbacks passed to Session() constructor (0.4.2 API)
  restart_callbacks — automatic server restart after crash detection

Usage
─────
  python3 demo_fuzzer.py

  Web UI: http://localhost:26001  (different port from fuzzer.py)
  DB:     boofuzz-results/
"""

import subprocess
import time
import os
import sys

from boofuzz import (
    Session,
    Target,
    TCPSocketConnection,
    s_initialize,
    s_static,
    s_string,
    s_group,
    s_get,
)

TARGET_HOST = "127.0.0.1"
TARGET_PORT = 4141
SERVER_BIN  = os.path.join(os.path.dirname(__file__), "demo_server")
SLEEP_TIME  = 0.04


# ── target management ──────────────────────────────────────────────────────

def start_server():
    proc = subprocess.Popen(
        [SERVER_BIN],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    time.sleep(0.3)
    return proc


def restart_server(target=None, fuzz_data_logger=None, session=None, **kwargs):
    """
    Called by boofuzz when crash_threshold_request consecutive test cases
    against the same request receive no response.  Kill the old instance
    and start a fresh one.
    """
    print("\n[!] Crash detected — restarting demo_server...")
    subprocess.run(["pkill", "-f", os.path.basename(SERVER_BIN)],
                   capture_output=True)
    time.sleep(0.3)
    subprocess.Popen(
        [SERVER_BIN],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    time.sleep(0.5)
    print("[!] Restarted.\n")


# ── protocol block definitions ─────────────────────────────────────────────
#
# boofuzz 0.4.2 API notes:
#
#   s_string(value, ...)  — `value` must be str, not bytes.  Generates
#       hundreds of mutations internally.  Does NOT accept a fuzz_list
#       or fuzz_values parameter; use s_group for a fixed list of values.
#
#   s_group(name=, values=, default_value=)  — all keyword args; the first
#       positional parameter is `name` (not `default_value`).
#
#   s_static(value)  — accepts bytes or str; never mutated.

def define_protocol():
    # ── SET <key> <value>: boundary sweep ─────────────────────────────────
    # Cycles through specific lengths bracketing value_buf[64] so the
    # class demo reaches the overflow without waiting for the full mutation
    # suite.  boofuzz's s_group steps through values[] in order.
    s_initialize("SET_boundary")
    s_static(b"SET ")
    s_string("demo_key", name="boundary_key", fuzzable=False)
    s_static(b" ")
    s_group(
        name="boundary_value",
        default_value=b"hello",
        values=[
            b"A" * 63,    # one short of value_buf — safe
            b"A" * 64,    # fills value_buf; NUL overwrites next byte
            b"A" * 65,    # one byte past buffer
            b"A" * 72,    # overwrites 8 bytes of adjacent stack memory
            b"A" * 128,   # almost certainly overwrites saved return address
            b"A" * 256,
            b"A" * 512,
        ],
    )
    s_static(b"\r\n")

    # ── SET <key> <value>: full string mutation suite ─────────────────────
    # boofuzz generates ~850 string mutations: long repeating strings,
    # format specifiers (%s %x %n), null bytes, CRLF injections, etc.
    s_initialize("SET")
    s_static(b"SET ")
    s_string("demo_key", name="set_key", fuzzable=False)
    s_static(b" ")
    s_string("hello", name="set_value", max_len=512)
    s_static(b"\r\n")

    # ── GET <key> ─────────────────────────────────────────────────────────
    s_initialize("GET")
    s_static(b"GET ")
    s_string("demo_key", name="get_key", max_len=256)
    s_static(b"\r\n")

    # ── SET_NUMERIC + LEN: signed/unsigned confusion ───────────────────────
    # SET_NUMERIC stores a numeric string in "num_key".
    # LEN then calls atoi(vals[idx]) on it; a negative stored value wraps
    # to SIZE_MAX when cast to size_t.
    #
    # s_group cycles through the interesting numeric values; LEN is
    # not fuzzed (the key is fixed — we are fuzzing the stored value path).
    s_initialize("SET_NUMERIC")
    s_static(b"SET ")
    s_string("num_key", name="num_key_field", fuzzable=False)
    s_static(b" ")
    s_group(
        name="num_value",
        default_value=b"-1",
        values=[
            b"-1",
            b"-2147483648",          # INT_MIN
            b"0",
            b"4294967295",           # UINT_MAX on 32-bit
            b"9999999999999999999",  # larger than INT64_MAX
        ],
    )
    s_static(b"\r\n")

    s_initialize("LEN")
    s_static(b"LEN ")
    s_string("num_key", name="len_key", fuzzable=False)
    s_static(b"\r\n")

    # ── APPEND <key> <data>: integer truncation ───────────────────────────
    # existing_len + new_len computed as int; wraps near INT_MAX.
    # boofuzz's built-in string mutations include lengths up to max_len,
    # so very long APPEND values exercise the integer-overflow path.
    s_initialize("APPEND")
    s_static(b"APPEND ")
    s_string("append_key", name="append_key_field", fuzzable=False)
    s_static(b" ")
    s_string("B" * 200, name="append_data", max_len=2048)
    s_static(b"\r\n")

    # ── KEYS, INFO, RESET, QUIT ───────────────────────────────────────────
    s_initialize("KEYS")
    s_static(b"KEYS\r\n")

    s_initialize("INFO")
    s_static(b"INFO")
    s_string("", name="info_trailing", max_len=64)   # trailing garbage
    s_static(b"\r\n")

    s_initialize("RESET")
    s_static(b"RESET\r\n")

    s_initialize("QUIT")
    s_static(b"QUIT\r\n")


# ── session graph ─────────────────────────────────────────────────────────

def build_session_graph(session):
    session.connect(s_get("SET_boundary"))             # fast boundary sweep

    session.connect(s_get("SET"))                      # full mutation suite
    session.connect(s_get("SET"),  s_get("GET"))
    session.connect(s_get("SET"),  s_get("KEYS"))
    session.connect(s_get("SET"),  s_get("QUIT"))

    session.connect(s_get("SET_NUMERIC"))              # int confusion path
    session.connect(s_get("SET_NUMERIC"), s_get("LEN"))

    session.connect(s_get("RESET"))                    # APPEND path
    session.connect(s_get("RESET"), s_get("APPEND"))

    session.connect(s_get("INFO"))


# ── callbacks ─────────────────────────────────────────────────────────────
#
# boofuzz 0.4.2 CallbackMonitor calls each function as:
#   f(target=target, fuzz_data_logger=fuzz_data_logger,
#     session=session, sock=target)
#
# Signatures must not include positional parameters beyond these four.
# Use **kwargs to absorb any extras without breaking the script.

def pre_send(target=None, fuzz_data_logger=None, session=None, **kwargs):
    """Drain the server banner at the start of each session."""
    try:
        target.recv(256)
    except Exception:
        pass


def post_send(target=None, fuzz_data_logger=None, session=None, **kwargs):
    """Drain the server response so the receive buffer stays clean."""
    try:
        resp = target.recv(512)
        if resp and resp.startswith(b"-ERR"):
            fuzz_data_logger.log_info(f"Server -ERR: {resp!r}")
    except Exception:
        pass


# ── main ──────────────────────────────────────────────────────────────────

def main():
    if not os.path.isfile(SERVER_BIN):
        print(f"[!] Server binary not found: {SERVER_BIN}")
        print("    Compile with: make   (or: gcc -o demo_server demo_server.c "
              "-fno-stack-protector -no-pie -g)")
        sys.exit(1)

    start_server()
    define_protocol()

    session = Session(
        target=Target(
            connection=TCPSocketConnection(
                TARGET_HOST,
                TARGET_PORT,
                recv_timeout=1.0,
            ),
        ),
        restart_callbacks=[restart_server],
        crash_threshold_request=3,
        crash_threshold_element=5,
        sleep_time=SLEEP_TIME,
        # In boofuzz 0.4.2 these are constructor parameters, not instance
        # attributes — appending to session.pre_send_callbacks after
        # construction does not work.  The post-send hook uses
        # post_test_case_callbacks (not post_send_callbacks).
        pre_send_callbacks=[pre_send],
        post_test_case_callbacks=[post_send],
        web_port=26001,
        keep_web_open=False,
    )

    build_session_graph(session)

    print(f"[*] Fuzzing VulnStore at {TARGET_HOST}:{TARGET_PORT}")
    print(f"[*] Web UI:  http://localhost:26001")
    print(f"[*] Vulns:   SET_boundary(overflow) SET(fmtstr) LEN(int) APPEND(truncation)")
    print(f"[*] Press Ctrl-C to stop.\n")

    session.fuzz()


if __name__ == "__main__":
    main()