#!/usr/bin/env python3
# Copyright (C) 2026 Douglas Quigg (dstroy0) <dquigg123@gmail.com>
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# ============================================================================
#  LEGAL / AUTHORIZED-USE-ONLY NOTICE  --  READ BEFORE RUNNING
# ----------------------------------------------------------------------------
#  This is an ACTIVE network attack tool. It sends malformed, oversized, and
#  high-volume traffic to a target you name on the command line. Run it ONLY
#  against a device that YOU OWN or that you have EXPLICIT, WRITTEN PERMISSION
#  to test.
#
#  Unauthorized access to or interference with computer systems is a crime in
#  most jurisdictions (e.g. the U.S. Computer Fraud and Abuse Act, the U.K.
#  Computer Misuse Act, and equivalents elsewhere). You - the operator - are
#  solely responsible for how you use this tool and for staying within the law
#  and the scope of your authorization.
#
#  This software is provided "AS IS", WITHOUT WARRANTY OF ANY KIND. The authors
#  and contributors accept NO LIABILITY for any damage, disruption, data loss,
#  or legal consequence arising from its use or misuse. We build defensive
#  tooling; use it for good.
#
#  By passing --authorized (or answering the prompt) you affirm that you have
#  permission to test the target. See pentesting/README.md and docs/PENTEST.md.
# ============================================================================
"""
detws_pentest - an adversarial test driver for a *live* DeterministicESPAsyncWebServer.

It complements the host fuzz harness (env:native_pentest, which fuzzes the parsers
off-device). This one drives a flashed board over the network and confirms the
live defenses hold and the device stays up and fails closed - within its fixed,
zero-heap footprint.

Highlights
  * Feature-aware: auto-detects the build's enabled features from the device's own
    /diag endpoint (DETWS_ENABLE_DIAG) and runs only the applicable attacks - or
    select them by hand with --flags / --only / --skip.
  * Determinism oracle: samples the device's free heap (via /health, /stats,
    /metrics, or /diag) before and after every attack and flags any drift - the
    core promise is "no heap after begin()", so drift is a finding.
  * Liveness oracle: after every attack a legitimate request must still succeed;
    a wedged server is a finding.
  * Multi-protocol: HTTP/HTTPS, plus UDP CoAP/SNMP and TCP Modbus/OPC UA fuzzing.
  * CI-friendly: --json report, deterministic --seed, and an exit code that is
    non-zero when findings exist.

Stdlib only (Python 3.8+); no third-party packages.

Quick start
  python3 detws_pentest.py --host 192.168.1.85 --diag --authorized
  python3 detws_pentest.py --host 192.168.1.85 --flags DETWS_ENABLE_AUTH,DETWS_ENABLE_ACCEPT_THROTTLE --authorized
  python3 detws_pentest.py --host 192.168.1.85 --all --intensity high --json report.json --authorized
  python3 detws_pentest.py --list
"""

import argparse
import json
import random
import socket
import ssl
import struct
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Dict, List, Optional, Set, Tuple

# --------------------------------------------------------------------------- #
# Console / formatting
# --------------------------------------------------------------------------- #

_USE_COLOR = sys.stdout.isatty()


def _c(code: str, s: str) -> str:
    if not _USE_COLOR:
        return s
    return f"\033[{code}m{s}\033[0m"


def bold(s):
    return _c("1", s)


def red(s):
    return _c("31", s)


def green(s):
    return _c("32", s)


def yellow(s):
    return _c("33", s)


def blue(s):
    return _c("36", s)


def grey(s):
    return _c("90", s)


# --------------------------------------------------------------------------- #
# Result model
# --------------------------------------------------------------------------- #


class Status(Enum):
    PASS = "PASS"  # the device behaved safely (failed closed / defense engaged / stayed up)
    FAIL = "FAIL"  # a finding - unsafe behavior
    WARN = "WARN"  # ran, but the outcome could not be fully confirmed
    SKIP = "SKIP"  # not applicable to this build, or the port was unreachable
    ERROR = "ERROR"  # the test itself errored


class Severity(Enum):
    INFO = 0
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4


@dataclass
class Finding:
    severity: Severity
    message: str


@dataclass
class TestResult:
    name: str
    status: Status = Status.PASS
    detail: str = ""
    findings: List[Finding] = field(default_factory=list)
    heap_before: Optional[int] = None
    heap_after: Optional[int] = None
    duration_s: float = 0.0

    def fail(self, severity: Severity, message: str):
        self.status = Status.FAIL
        self.findings.append(Finding(severity, message))

    def add_finding(self, severity: Severity, message: str):
        self.findings.append(Finding(severity, message))
        if severity.value >= Severity.MEDIUM.value:
            self.status = Status.FAIL


# --------------------------------------------------------------------------- #
# Target / context
# --------------------------------------------------------------------------- #


@dataclass
class Target:
    host: str
    http_port: int = 80
    https_port: int = 443
    use_tls: bool = False
    coap_port: int = 5683
    snmp_port: int = 161
    modbus_port: int = 502
    opcua_port: int = 4840
    username: str = "admin"
    password: str = "admin"
    secure_path: str = "/secure"  # a route expected to require auth
    upload_path: str = "/upload"
    timeout: float = 4.0


INTENSITY = {
    #            connections  iterations  duration_s
    "low": dict(conns=20, iters=200, duration=3),
    "medium": dict(conns=64, iters=1000, duration=6),
    "high": dict(conns=160, iters=5000, duration=12),
}


@dataclass
class Context:
    target: Target
    intensity: str = "medium"
    rng: random.Random = field(default_factory=random.Random)
    enabled_features: Set[str] = field(default_factory=set)
    config: Dict[str, int] = field(default_factory=dict)
    verbose: bool = False

    @property
    def knob(self) -> dict:
        return INTENSITY[self.intensity]

    def log(self, msg: str):
        if self.verbose:
            print(grey("    . " + msg))


# --------------------------------------------------------------------------- #
# Low-level transport helpers
# --------------------------------------------------------------------------- #


def _wrap_tls(sock: socket.socket, host: str) -> socket.socket:
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE  # device certs are typically self-signed / pinned
    return ctx.wrap_socket(sock, server_hostname=host)


def tcp_send_recv(
    tgt: Target,
    data: bytes,
    read: bool = True,
    tls: Optional[bool] = None,
    port: Optional[int] = None,
    timeout: Optional[float] = None,
) -> bytes:
    """Open a TCP (optionally TLS) socket, send `data`, read until close/timeout."""
    use_tls = tgt.use_tls if tls is None else tls
    p = port if port is not None else (tgt.https_port if use_tls else tgt.http_port)
    to = timeout if timeout is not None else tgt.timeout
    s = socket.create_connection((tgt.host, p), timeout=to)
    try:
        if use_tls:
            s = _wrap_tls(s, tgt.host)
        s.settimeout(to)
        if data:
            s.sendall(data)
        if not read:
            return b""
        chunks = []
        try:
            while True:
                b = s.recv(4096)
                if not b:
                    break
                chunks.append(b)
                if sum(len(c) for c in chunks) > 1 << 20:
                    break
        except socket.timeout:
            pass
        return b"".join(chunks)
    finally:
        try:
            s.close()
        except OSError:
            pass


def udp_send_recv(host: str, port: int, data: bytes, timeout: float = 3.0) -> Optional[bytes]:
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(timeout)
    try:
        s.sendto(data, (host, port))
        try:
            resp, _ = s.recvfrom(4096)
            return resp
        except socket.timeout:
            return None
    finally:
        s.close()


@dataclass
class HttpResponse:
    status: int
    headers: Dict[str, str]
    body: bytes
    raw: bytes


def http_request(
    tgt: Target,
    method: str = "GET",
    path: str = "/",
    headers: Optional[Dict[str, str]] = None,
    body: bytes = b"",
    tls: Optional[bool] = None,
    auth: Optional[Tuple[str, str]] = None,
    raw_request: Optional[bytes] = None,
    timeout: Optional[float] = None,
) -> Optional[HttpResponse]:
    """Send one HTTP/1.1 request (or a fully raw byte blob) and parse the response."""
    if raw_request is None:
        hdrs = {"Host": tgt.host, "Connection": "close"}
        if headers:
            hdrs.update(headers)
        if auth:
            import base64

            tok = base64.b64encode(f"{auth[0]}:{auth[1]}".encode()).decode()
            hdrs["Authorization"] = "Basic " + tok
        if body:
            hdrs.setdefault("Content-Length", str(len(body)))
        lines = [f"{method} {path} HTTP/1.1"]
        lines += [f"{k}: {v}" for k, v in hdrs.items()]
        req = ("\r\n".join(lines) + "\r\n\r\n").encode() + body
    else:
        req = raw_request
    try:
        data = tcp_send_recv(tgt, req, read=True, tls=tls, timeout=timeout)
    except (OSError, ssl.SSLError):
        return None
    if not data:
        return None
    return _parse_http(data)


def _parse_http(data: bytes) -> Optional[HttpResponse]:
    try:
        head, _, body = data.partition(b"\r\n\r\n")
        lines = head.split(b"\r\n")
        status = int(lines[0].split(b" ")[1])
        headers = {}
        for ln in lines[1:]:
            if b":" in ln:
                k, _, v = ln.partition(b":")
                headers[k.decode(errors="replace").strip().lower()] = v.decode(errors="replace").strip()
        return HttpResponse(status, headers, body, data)
    except (ValueError, IndexError):
        return HttpResponse(0, {}, data, data)


# --------------------------------------------------------------------------- #
# Oracles: liveness + heap drift
# --------------------------------------------------------------------------- #


def is_alive(tgt: Target) -> bool:
    """A legitimate GET / must still get a response."""
    r = http_request(tgt, "GET", "/", timeout=tgt.timeout)
    return r is not None and r.status > 0


def _extract_heap(body: bytes) -> Optional[int]:
    text = body.decode(errors="replace")
    # JSON-ish: "free_heap": 211234  /  "free_heap_bytes":...  /  "heap":...
    for key in ("free_heap", "free_heap_bytes", "freeheap", "heap_free", "heap"):
        idx = text.find(key)
        if idx >= 0:
            j = idx + len(key)
            num = ""
            seen = False
            for ch in text[j : j + 32]:
                if ch.isdigit():
                    num += ch
                    seen = True
                elif seen:
                    break
            if num:
                return int(num)
    # Prometheus: detws_free_heap_bytes <n>
    for ln in text.splitlines():
        ls = ln.strip()
        if "heap" in ls and not ls.startswith("#"):
            parts = ls.replace("\t", " ").split(" ")
            for p in reversed(parts):
                if p.replace(".", "").isdigit():
                    try:
                        return int(float(p))
                    except ValueError:
                        pass
    return None


HEALTH_PATHS = ("/health", "/stats", "/metrics", "/diag")


def sample_heap(tgt: Target) -> Optional[int]:
    for path in HEALTH_PATHS:
        r = http_request(tgt, "GET", path, timeout=tgt.timeout)
        if r and r.status == 200 and r.body:
            h = _extract_heap(r.body)
            if h is not None:
                return h
    return None


# --------------------------------------------------------------------------- #
# Feature detection (the /diag endpoint)
# --------------------------------------------------------------------------- #

# Map the diag "features" booleans to the canonical DETWS_ENABLE_* flags.
DIAG_FEATURE_FLAG = {
    "websocket": "DETWS_ENABLE_WEBSOCKET",
    "sse": "DETWS_ENABLE_SSE",
    "multipart": "DETWS_ENABLE_MULTIPART",
    "file_serving": "DETWS_ENABLE_FILE_SERVING",
    "auth": "DETWS_ENABLE_AUTH",
}


def detect_features(tgt: Target) -> Tuple[Set[str], Dict[str, int]]:
    """Read /diag (DETWS_ENABLE_DIAG) and return (enabled flags, config sizes)."""
    flags: Set[str] = set()
    cfg: Dict[str, int] = {}
    r = http_request(tgt, "GET", "/diag", timeout=tgt.timeout)
    if not r or r.status != 200:
        return flags, cfg
    try:
        doc = json.loads(r.body.decode(errors="replace"))
    except json.JSONDecodeError:
        return flags, cfg
    for name, on in (doc.get("features") or {}).items():
        if on and name in DIAG_FEATURE_FLAG:
            flags.add(DIAG_FEATURE_FLAG[name])
    for k, v in (doc.get("config") or {}).items():
        if isinstance(v, int):
            cfg[k] = v
    return flags, cfg


# --------------------------------------------------------------------------- #
# Attack module registry
# --------------------------------------------------------------------------- #


@dataclass
class Attack:
    name: str
    tags: Set[str]  # category tags for --only / --skip
    features: Set[str]  # required DETWS_ENABLE_* (empty = always applicable)
    proto: str  # http / https / udp-coap / udp-snmp / tcp-modbus / tcp-opcua
    desc: str
    run: Callable[[Context, TestResult], None]


REGISTRY: List[Attack] = []


def attack(name, tags, features, proto, desc):
    def deco(fn):
        REGISTRY.append(Attack(name, set(tags), set(features), proto, desc, fn))
        return fn

    return deco


# ----- HTTP core (no feature gate; the HTTP server is always present) ------- #


@attack("http_baseline", ["http", "sanity"], [], "http", "Confirm the server answers a normal request (baseline).")
def a_baseline(ctx, res):
    r = http_request(ctx.target, "GET", "/")
    if r is None:
        res.fail(Severity.HIGH, "no response to a normal GET / (is the host/port right?)")
    else:
        res.detail = f"GET / -> {r.status}"


@attack(
    "http_oversized_request_line",
    ["http", "limits"],
    [],
    "http",
    "A megabyte-long request target must yield 414, not an overflow.",
)
def a_oversized(ctx, res):
    path = "/" + "A" * 100000
    r = http_request(ctx.target, "GET", path)
    if r is None:
        res.add_finding(Severity.LOW, "connection dropped with no status line (acceptable if slot was freed)")
        res.status = Status.WARN
        return
    res.detail = f"status={r.status}"
    if r.status not in (414, 400, 431):
        res.add_finding(Severity.MEDIUM, f"oversized target returned {r.status}, expected 414/400/431")


@attack(
    "http_huge_content_length",
    ["http", "limits"],
    [],
    "http",
    "A giant Content-Length must yield 413, never reserve heap.",
)
def a_huge_cl(ctx, res):
    r = http_request(ctx.target, "POST", "/", headers={"Content-Length": "99999999"})
    if r is None:
        res.status = Status.WARN
        res.detail = "no status line"
        return
    res.detail = f"status={r.status}"
    if r.status not in (413, 400):
        res.add_finding(Severity.MEDIUM, f"huge Content-Length returned {r.status}, expected 413/400")


@attack(
    "http_header_flood",
    ["http", "limits"],
    [],
    "http",
    "Hundreds of headers must be clamped to MAX_HEADERS, no overflow.",
)
def a_header_flood(ctx, res):
    hdrs = {f"X-Pad-{i}": "1" for i in range(300)}
    r = http_request(ctx.target, "GET", "/", headers=hdrs)
    res.detail = f"status={r.status if r else 'none'}"
    if r is None:
        res.status = Status.WARN


@attack(
    "http_smuggling",
    ["http", "smuggling"],
    [],
    "http",
    "Ambiguous TE/CL framing must be rejected or normalized, never desync.",
)
def a_smuggling(ctx, res):
    vectors = [
        b"POST / HTTP/1.1\r\nHost: x\r\nContent-Length: 5\r\nContent-Length: 6\r\n\r\nHELLO",
        b"POST / HTTP/1.1\r\nHost: x\r\nContent-Length: 4\r\nTransfer-Encoding: chunked\r\n\r\n0\r\n\r\n",
        b"POST / HTTP/1.1\r\nHost: x\r\nTransfer-Encoding: chunked\r\nContent-Length: 4\r\n\r\n5\r\nHELLO\r\n0\r\n\r\n",
        b"POST / HTTP/1.1\r\nHost: x\r\nContent-Length: -1\r\n\r\n",
    ]
    statuses = []
    for v in vectors:
        r = http_request(ctx.target, raw_request=v)
        statuses.append(r.status if r else 0)
    res.detail = "statuses=" + ",".join(str(s) for s in statuses)
    # A 2xx on an ambiguous frame is the classic smuggling tell.
    for v, s in zip(vectors, statuses):
        if 200 <= s < 300:
            res.add_finding(Severity.MEDIUM, f"ambiguous frame accepted with {s}: {v[:40]!r}")


@attack(
    "http_path_traversal",
    ["http", "traversal"],
    ["DETWS_ENABLE_FILE_SERVING"],
    "http",
    "../ and encoded traversal must never escape the served root.",
)
def a_traversal(ctx, res):
    probes = [
        "/../../../../etc/passwd",
        "/%2e%2e/%2e%2e/%2e%2e/etc/passwd",
        "/..%2f..%2f..%2fetc%2fpasswd",
        "/static/../../../../etc/passwd",
        "/....//....//etc/passwd",
    ]
    leaked = []
    statuses = []
    for p in probes:
        r = http_request(ctx.target, "GET", p)
        if r is None:
            statuses.append(0)
            continue
        statuses.append(r.status)
        if r.status == 200 and (b"root:" in r.body or b"/bin/" in r.body):
            leaked.append(p)
    res.detail = "statuses=" + ",".join(str(s) for s in statuses)
    if leaked:
        res.fail(Severity.CRITICAL, "path traversal leaked filesystem content: " + ", ".join(leaked))


@attack(
    "http_slowloris",
    ["http", "dos", "liveness"],
    [],
    "http",
    "Many half-open dribbling connections must not wedge the server.",
)
def a_slowloris(ctx, res):
    n = ctx.knob["conns"]
    tgt = ctx.target
    socks = []
    try:
        for _ in range(n):
            try:
                s = socket.create_connection((tgt.host, tgt.http_port), timeout=tgt.timeout)
                s.sendall(b"GET / HTTP/1.1\r\nHost: x\r\n")  # headers, never terminated
                socks.append(s)
            except OSError:
                break
        res.detail = f"opened {len(socks)}/{n} half-open connections"
        time.sleep(min(ctx.knob["duration"], 8))
        # Drip a byte into each to keep them "active".
        for s in socks:
            try:
                s.sendall(b"X")
            except OSError:
                pass
        if not is_alive(tgt):
            res.fail(Severity.HIGH, "server stopped answering legitimate requests during slowloris")
    finally:
        for s in socks:
            try:
                s.close()
            except OSError:
                pass


@attack("http_conn_saturation", ["http", "dos"], [], "http", "Opening more than MAX_CONNS sockets must cap, not crash.")
def a_saturation(ctx, res):
    tgt = ctx.target
    cap = ctx.config.get("MAX_CONNS", 8)
    want = cap * 3 + 8
    socks = []
    try:
        for _ in range(want):
            try:
                s = socket.create_connection((tgt.host, tgt.http_port), timeout=1.5)
                socks.append(s)
            except OSError:
                break
        res.detail = f"established {len(socks)} concurrent sockets (MAX_CONNS={cap})"
        # The server should refuse beyond its pool, not fall over.
        if not is_alive(tgt):
            res.add_finding(Severity.MEDIUM, "server unresponsive while connection pool saturated")
    finally:
        for s in socks:
            try:
                s.close()
            except OSError:
                pass


@attack("http_random_fuzz", ["http", "fuzz"], [], "http", "Random byte streams to :80 must never crash the server.")
def a_http_fuzz(ctx, res):
    tgt = ctx.target
    iters = min(ctx.knob["iters"], 1500)
    sent = 0
    for _ in range(iters):
        n = ctx.rng.randint(0, 400)
        blob = bytes(ctx.rng.randint(0, 255) for _ in range(n))
        try:
            tcp_send_recv(tgt, blob, read=False, timeout=1.0)
            sent += 1
        except OSError:
            pass
    res.detail = f"sent {sent} random streams"
    if not is_alive(tgt):
        res.fail(Severity.HIGH, "server unresponsive after random HTTP fuzz")


# ----- Defenses ------------------------------------------------------------- #


@attack(
    "accept_throttle_flood",
    ["defense", "dos"],
    ["DETWS_ENABLE_ACCEPT_THROTTLE"],
    "http",
    "A rapid accept flood must be throttled (drops/RST), server stays up.",
)
def a_accept_throttle(ctx, res):
    tgt = ctx.target
    n = ctx.knob["conns"] * 4
    ok = refused = 0

    def hit(_):
        nonlocal ok, refused
        try:
            r = http_request(tgt, "GET", "/", timeout=1.5)
            if r:
                ok += 1
        except OSError:
            refused += 1

    with ThreadPoolExecutor(max_workers=32) as ex:
        list(ex.map(hit, range(n)))
    res.detail = f"{ok} served, {refused} refused of {n}"
    if not is_alive(tgt):
        res.fail(Severity.HIGH, "server down after accept flood (throttle did not protect it)")
    elif refused == 0:
        res.add_finding(Severity.LOW, "no connections were refused; throttle may be set high or disabled")


@attack(
    "per_ip_throttle",
    ["defense"],
    ["DETWS_ENABLE_PER_IP_THROTTLE"],
    "http",
    "Rapid requests from one IP should start returning 429.",
)
def a_per_ip(ctx, res):
    tgt = ctx.target
    saw_429 = False
    statuses = []
    for _ in range(min(ctx.knob["iters"], 120)):
        r = http_request(tgt, "GET", "/", timeout=1.5)
        if r:
            statuses.append(r.status)
            if r.status == 429:
                saw_429 = True
                break
    res.detail = f"{len(statuses)} requests, saw 429={saw_429}"
    if not saw_429:
        res.add_finding(Severity.LOW, "no 429 seen; per-IP limit may be high or not wired to this route")


@attack(
    "auth_bruteforce",
    ["defense", "auth"],
    ["DETWS_ENABLE_AUTH"],
    "http",
    "Bad credentials must get 401; protected route must not open.",
)
def a_auth_brute(ctx, res):
    tgt = ctx.target
    opened = False
    last = None
    for i in range(min(ctx.knob["iters"], 60)):
        r = http_request(tgt, "GET", tgt.secure_path, auth=(tgt.username, f"wrong{i}"))
        if r is None:
            continue
        last = r.status
        if r.status == 200:
            opened = True
            break
    res.detail = f"last status={last} on {tgt.secure_path}"
    if opened:
        res.fail(Severity.CRITICAL, f"protected route {tgt.secure_path} served 200 to a wrong password")
    elif last is None:
        res.status = Status.WARN
        res.detail += " (no response; check --secure-path)"


@attack(
    "auth_lockout",
    ["defense", "auth"],
    ["DETWS_ENABLE_AUTH_LOCKOUT"],
    "http",
    "Repeated bad logins should trip a lockout (429 + Retry-After).",
)
def a_auth_lockout(ctx, res):
    tgt = ctx.target
    locked = False
    retry_after = False
    for i in range(min(ctx.knob["iters"], 60)):
        r = http_request(tgt, "GET", tgt.secure_path, auth=(tgt.username, f"bad{i}"))
        if r and r.status == 429:
            locked = True
            retry_after = "retry-after" in r.headers
            break
    res.detail = f"locked={locked} retry_after={retry_after}"
    if not locked:
        res.add_finding(Severity.MEDIUM, "no lockout (429) after many bad logins; lockout may be disabled/high")
    elif not retry_after:
        res.add_finding(Severity.LOW, "lockout returned 429 without a Retry-After header")


@attack(
    "csrf_guard",
    ["defense", "auth"],
    ["DETWS_ENABLE_CSRF"],
    "http",
    "A state-changing request without a valid token must get 403.",
)
def a_csrf(ctx, res):
    tgt = ctx.target
    r = http_request(tgt, "POST", "/", body=b"x=1", headers={"Content-Type": "application/x-www-form-urlencoded"})
    res.detail = f"POST without token -> {r.status if r else 'none'}"
    if r and 200 <= r.status < 300:
        res.add_finding(Severity.MEDIUM, "unprotected POST accepted without a CSRF token")


@attack(
    "range_fuzz",
    ["http", "fuzz"],
    ["DETWS_ENABLE_RANGE"],
    "http",
    "Malformed Range headers must yield 416/200, never an overflow.",
)
def a_range(ctx, res):
    tgt = ctx.target
    ranges = [
        "bytes=999999999-",
        "bytes=-99999999999",
        "bytes=0-0,1-1,2-2",
        "bytes=abc-def",
        "bytes=10-5",
        "bytes=" + "9" * 40 + "-",
    ]
    statuses = []
    for rg in ranges:
        r = http_request(tgt, "GET", "/", headers={"Range": rg})
        statuses.append(r.status if r else 0)
    res.detail = "statuses=" + ",".join(str(s) for s in statuses)
    if not is_alive(tgt):
        res.fail(Severity.HIGH, "server unresponsive after Range fuzz")


@attack(
    "websocket_fuzz",
    ["ws", "fuzz"],
    ["DETWS_ENABLE_WEBSOCKET"],
    "http",
    "Malformed WebSocket upgrades/frames must not crash the server.",
)
def a_ws(ctx, res):
    tgt = ctx.target
    # 1) malformed upgrade requests
    bad_upgrades = [
        b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n",  # no key
        b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: !!!\r\nSec-WebSocket-Version: 13\r\n\r\n",
        b"GET /ws HTTP/1.1\r\nHost: x\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: "
        + b"A" * 4000
        + b"\r\n\r\n",
    ]
    for u in bad_upgrades:
        http_request(tgt, raw_request=u)
    # 2) a real-ish handshake then a bogus oversized frame
    import base64

    key = base64.b64encode(b"0123456789abcdef").decode()
    up = (
        f"GET /ws HTTP/1.1\r\nHost: {tgt.host}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
        f"Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n"
    ).encode()
    try:
        s = socket.create_connection((tgt.host, tgt.http_port), timeout=tgt.timeout)
        s.sendall(up)
        s.recv(2048)
        # client frame, FIN+binary, 127 => 8-byte length claiming a huge payload
        s.sendall(b"\x82\xff" + struct.pack(">Q", 1 << 40) + b"\x00\x00\x00\x00")
        s.close()
    except OSError:
        pass
    if not is_alive(tgt):
        res.fail(Severity.HIGH, "server unresponsive after WebSocket fuzz")
    else:
        res.detail = "server survived malformed upgrades + oversized frame"


@attack(
    "upload_oversize",
    ["http", "limits"],
    ["DETWS_ENABLE_UPLOAD"],
    "http",
    "An oversized streamed upload must be bounded, not exhaust memory.",
)
def a_upload(ctx, res):
    tgt = ctx.target
    body = b"A" * 200000
    r = http_request(tgt, "POST", tgt.upload_path, body=body, timeout=tgt.timeout + 4)
    res.detail = f"POST {tgt.upload_path} {len(body)}B -> {r.status if r else 'none'}"
    if not is_alive(tgt):
        res.fail(Severity.HIGH, "server unresponsive after oversized upload")


# ----- TLS ------------------------------------------------------------------ #


@attack(
    "tls_handshake",
    ["tls"],
    ["DETWS_ENABLE_TLS"],
    "https",
    "TLS handshake completes; report negotiated version/cipher.",
)
def a_tls(ctx, res):
    tgt = ctx.target
    try:
        raw = socket.create_connection((tgt.host, tgt.https_port), timeout=tgt.timeout)
        c = ssl.create_default_context()
        c.check_hostname = False
        c.verify_mode = ssl.CERT_NONE
        s = c.wrap_socket(raw, server_hostname=tgt.host)
        ver = s.version()
        cipher = s.cipher()
        s.close()
        res.detail = f"{ver}, cipher={cipher[0] if cipher else '?'}"
        if ver in ("SSLv3", "TLSv1", "TLSv1.1"):
            res.add_finding(Severity.MEDIUM, f"negotiated a deprecated TLS version: {ver}")
    except (OSError, ssl.SSLError) as e:
        res.status = Status.SKIP
        res.detail = f"no TLS on :{tgt.https_port} ({e.__class__.__name__})"


# ----- UDP / binary protocols ----------------------------------------------- #


@attack(
    "coap_fuzz",
    ["coap", "udp", "fuzz"],
    ["DETWS_ENABLE_COAP"],
    "udp-coap",
    "Malformed/oversized CoAP datagrams must not crash the agent.",
)
def a_coap(ctx, res):
    tgt = ctx.target
    sent = answered = 0
    # structured-ish: valid 4-byte header + random options
    for _ in range(min(ctx.knob["iters"], 400)):
        tkl = ctx.rng.randint(0, 8)
        hdr = bytes([0x40 | tkl, ctx.rng.randint(0, 0x7F), ctx.rng.randint(0, 255), ctx.rng.randint(0, 255)])
        rest = bytes(ctx.rng.randint(0, 255) for _ in range(ctx.rng.randint(0, 120)))
        try:
            r = udp_send_recv(tgt.host, tgt.coap_port, hdr + rest, timeout=0.4)
            sent += 1
            if r:
                answered += 1
        except OSError:
            pass
    res.detail = f"sent {sent}, {answered} answered"
    if sent == 0:
        res.status = Status.SKIP
        res.detail = f"no UDP socket to :{tgt.coap_port}"


@attack(
    "snmp_fuzz",
    ["snmp", "udp", "fuzz"],
    ["DETWS_ENABLE_SNMP"],
    "udp-snmp",
    "Malformed SNMP/BER datagrams must not crash the agent.",
)
def a_snmp(ctx, res):
    tgt = ctx.target
    sent = answered = 0
    for _ in range(min(ctx.knob["iters"], 400)):
        choice = ctx.rng.random()
        if choice < 0.5:
            inner = bytes(ctx.rng.randint(0, 255) for _ in range(ctx.rng.randint(0, 60)))
            pkt = bytes([0x30, len(inner) & 0xFF]) + inner  # SEQUENCE wrapper
        else:
            pkt = bytes(ctx.rng.randint(0, 255) for _ in range(ctx.rng.randint(0, 80)))
        try:
            r = udp_send_recv(tgt.host, tgt.snmp_port, pkt, timeout=0.4)
            sent += 1
            if r:
                answered += 1
        except OSError:
            pass
    res.detail = f"sent {sent}, {answered} answered"
    if sent == 0:
        res.status = Status.SKIP


@attack(
    "modbus_fuzz",
    ["modbus", "tcp", "fuzz"],
    ["DETWS_ENABLE_MODBUS"],
    "tcp-modbus",
    "Malformed Modbus ADUs must yield exceptions, not crashes.",
)
def a_modbus(ctx, res):
    tgt = ctx.target
    sent = 0
    for _ in range(min(ctx.knob["iters"], 400)):
        n = ctx.rng.randint(0, 270)
        adu = bytes(ctx.rng.randint(0, 255) for _ in range(n))
        try:
            tcp_send_recv(tgt, adu, read=False, port=tgt.modbus_port, tls=False, timeout=0.8)
            sent += 1
        except OSError:
            break
    res.detail = f"sent {sent} ADUs to :{tgt.modbus_port}"
    if sent == 0:
        res.status = Status.SKIP
        res.detail = f"no TCP to :{tgt.modbus_port}"


@attack(
    "opcua_fuzz",
    ["opcua", "tcp", "fuzz"],
    ["DETWS_ENABLE_OPCUA"],
    "tcp-opcua",
    "Malformed OPC UA (UACP) frames must not crash the endpoint.",
)
def a_opcua(ctx, res):
    tgt = ctx.target
    sent = 0
    for _ in range(min(ctx.knob["iters"], 300)):
        msgtype = ctx.rng.choice([b"HEL", b"OPN", b"MSG", b"CLO", b"XXX"])
        size = ctx.rng.randint(0, 200)
        body = bytes(ctx.rng.randint(0, 255) for _ in range(ctx.rng.randint(0, 120)))
        frame = msgtype + b"F" + struct.pack("<I", size) + body
        try:
            tcp_send_recv(tgt, frame, read=False, port=tgt.opcua_port, tls=False, timeout=0.8)
            sent += 1
        except OSError:
            break
    res.detail = f"sent {sent} frames to :{tgt.opcua_port}"
    if sent == 0:
        res.status = Status.SKIP
        res.detail = f"no TCP to :{tgt.opcua_port}"


# --------------------------------------------------------------------------- #
# Runner
# --------------------------------------------------------------------------- #


def select_attacks(ctx: Context, only: Set[str], skip: Set[str], run_all: bool) -> List[Attack]:
    chosen = []
    for atk in REGISTRY:
        if only and not (atk.tags & only) and atk.name not in only:
            continue
        if skip and (atk.tags & skip or atk.name in skip):
            continue
        if not run_all and atk.features and not atk.features.issubset(ctx.enabled_features):
            continue
        chosen.append(atk)
    return chosen


def run_attack(ctx: Context, atk: Attack) -> TestResult:
    res = TestResult(atk.name)
    res.heap_before = sample_heap(ctx.target)
    t0 = time.time()
    try:
        atk.run(ctx, res)
    except Exception as e:  # a module bug must not abort the run
        res.status = Status.ERROR
        res.detail = f"{e.__class__.__name__}: {e}"
    res.duration_s = time.time() - t0
    res.heap_after = sample_heap(ctx.target)
    # Heap-drift oracle (the determinism promise).
    if res.heap_before and res.heap_after:
        drift = res.heap_before - res.heap_after
        thresh = max(512, res.heap_before // 100)  # 1% or 512B noise floor
        if drift > thresh:
            res.add_finding(
                Severity.HIGH,
                f"free heap dropped {drift} bytes ({res.heap_before}->{res.heap_after}); "
                f"determinism promise is 'no heap after begin()'",
            )
    # Liveness oracle (unless the test already concluded the server is down).
    if res.status not in (Status.SKIP,) and not is_alive(ctx.target):
        res.fail(Severity.CRITICAL, "server is not answering after this test (possible crash/wedge)")
    return res


# --------------------------------------------------------------------------- #
# Reporting
# --------------------------------------------------------------------------- #


def print_banner():
    line = "=" * 78
    print(bold(line))
    print(bold("  detws_pentest - live adversarial tester for DeterministicESPAsyncWebServer"))
    print(line)
    print(
        yellow(
            "  AUTHORIZED USE ONLY. This tool sends real attack traffic. Use it only\n"
            "  against devices you own or are explicitly permitted to test. You are\n"
            "  responsible for staying within the law and your authorization. No\n"
            "  warranty; no liability. We use our powers for good."
        )
    )
    print(line)


STATUS_COLOR = {
    Status.PASS: green,
    Status.FAIL: red,
    Status.WARN: yellow,
    Status.SKIP: grey,
    Status.ERROR: red,
}


def print_results(results: List[TestResult]):
    print()
    print(bold("  Results"))
    print("  " + "-" * 76)
    name_w = max((len(r.name) for r in results), default=12)
    for r in results:
        col = STATUS_COLOR[r.status]
        print(
            f"  {col(r.status.value.ljust(5))}  {r.name.ljust(name_w)}  " f"{grey(f'{r.duration_s:5.1f}s')}  {r.detail}"
        )
        for f in r.findings:
            sev = f.severity.name
            paint = red if f.severity.value >= Severity.HIGH.value else yellow
            print(f"         {paint('* ' + sev)}: {f.message}")
    print("  " + "-" * 76)
    counts = {s: sum(1 for r in results if r.status == s) for s in Status}
    findings = sum(len(r.findings) for r in results)
    summary = "  ".join(STATUS_COLOR[s](f"{counts[s]} {s.value}") for s in Status if counts[s])
    print(f"  {summary}   |   {findings} finding(s)")


def to_json(results: List[TestResult], ctx: Context) -> str:
    return json.dumps(
        {
            "target": ctx.target.host,
            "intensity": ctx.intensity,
            "enabled_features": sorted(ctx.enabled_features),
            "results": [
                {
                    "name": r.name,
                    "status": r.status.value,
                    "detail": r.detail,
                    "duration_s": round(r.duration_s, 3),
                    "heap_before": r.heap_before,
                    "heap_after": r.heap_after,
                    "findings": [{"severity": f.severity.name, "message": f.message} for f in r.findings],
                }
                for r in results
            ],
        },
        indent=2,
    )


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


def parse_args(argv):
    p = argparse.ArgumentParser(
        prog="detws_pentest",
        description="Adversarial tester for a live DeterministicESPAsyncWebServer device.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="AUTHORIZED USE ONLY - test only what you own or may legally test.",
    )
    p.add_argument("--host", help="target IP / hostname")
    p.add_argument("--port", type=int, default=80, help="HTTP port (default 80)")
    p.add_argument("--https-port", type=int, default=443)
    p.add_argument("--tls", action="store_true", help="run HTTP attacks over TLS")
    p.add_argument("--coap-port", type=int, default=5683)
    p.add_argument("--snmp-port", type=int, default=161)
    p.add_argument("--modbus-port", type=int, default=502)
    p.add_argument("--opcua-port", type=int, default=4840)
    p.add_argument("--username", default="admin")
    p.add_argument("--password", default="admin")
    p.add_argument("--secure-path", default="/secure", help="a route expected to require auth")
    p.add_argument("--upload-path", default="/upload")
    p.add_argument("--timeout", type=float, default=4.0)

    p.add_argument("--diag", action="store_true", help="auto-detect enabled features from the device /diag endpoint")
    p.add_argument("--flags", default="", help="comma-separated DETWS_ENABLE_* flags your build has")
    p.add_argument("--all", action="store_true", help="run every applicable attack regardless of detected flags")
    p.add_argument("--only", default="", help="comma-separated tags/names to include")
    p.add_argument("--skip", default="", help="comma-separated tags/names to exclude")

    p.add_argument("--intensity", choices=list(INTENSITY), default="medium")
    p.add_argument("--seed", type=int, default=0xC0FFEE, help="PRNG seed (reproducibility)")
    p.add_argument("--json", metavar="FILE", help="write a JSON report to FILE")
    p.add_argument("-v", "--verbose", action="store_true")
    p.add_argument("--list", action="store_true", help="list attacks and exit")
    p.add_argument(
        "--authorized", action="store_true", help="affirm you are permitted to test the target (skips the prompt)"
    )
    return p.parse_args(argv)


def cmd_list():
    print(bold("Available attacks (tag(s) / required feature):"))
    for atk in REGISTRY:
        feat = ",".join(sorted(atk.features)) or "always"
        print(f"  {atk.name.ljust(28)} [{','.join(sorted(atk.tags))}]  " f"{grey('needs ' + feat)}")
        print(grey(f"      {atk.desc}"))


def confirm_authorization(args) -> bool:
    if args.authorized:
        return True
    print(yellow("\n  Type 'I HAVE PERMISSION' to confirm you are authorized to test " f"{args.host}:"))
    try:
        ans = input("  > ").strip()
    except (EOFError, KeyboardInterrupt):
        return False
    return ans == "I HAVE PERMISSION"


def main(argv=None):
    args = parse_args(sys.argv[1:] if argv is None else argv)

    print_banner()

    if args.list:
        cmd_list()
        return 0

    if not args.host:
        print(red("\n  error: --host is required (or use --list).\n"))
        return 2

    if not confirm_authorization(args):
        print(red("\n  Authorization not confirmed. Aborting.\n"))
        return 3

    tgt = Target(
        host=args.host,
        http_port=args.port,
        https_port=args.https_port,
        use_tls=args.tls,
        coap_port=args.coap_port,
        snmp_port=args.snmp_port,
        modbus_port=args.modbus_port,
        opcua_port=args.opcua_port,
        username=args.username,
        password=args.password,
        secure_path=args.secure_path,
        upload_path=args.upload_path,
        timeout=args.timeout,
    )
    ctx = Context(target=tgt, intensity=args.intensity, rng=random.Random(args.seed), verbose=args.verbose)

    # Feature set: from flags, from /diag, or --all.
    if args.flags:
        ctx.enabled_features |= {f.strip() for f in args.flags.split(",") if f.strip()}
    if args.diag or args.all:
        feats, cfg = detect_features(tgt)
        ctx.enabled_features |= feats
        ctx.config.update(cfg)
        if feats:
            print(blue(f"\n  /diag detected: {', '.join(sorted(feats))}"))
            if cfg:
                print(grey(f"  config: {cfg}"))
        elif args.diag:
            print(
                yellow(
                    "\n  /diag not available or returned no features " "(enable DETWS_ENABLE_DIAG, or pass --flags)."
                )
            )

    only = {x.strip() for x in args.only.split(",") if x.strip()}
    skip = {x.strip() for x in args.skip.split(",") if x.strip()}
    attacks = select_attacks(ctx, only, skip, args.all)

    if not attacks:
        print(yellow("\n  No attacks selected for this feature set. " "Try --all, --diag, or --flags.\n"))
        return 0

    # Pre-flight reachability.
    base = sample_heap(tgt)
    print(
        blue(
            f"\n  Target {tgt.host}:{tgt.http_port}  intensity={ctx.intensity}  "
            f"attacks={len(attacks)}  baseline_heap={base if base is not None else 'n/a'}"
        )
    )
    if not is_alive(tgt):
        print(
            red(
                "\n  warning: baseline GET / got no response - the host may be down "
                "or the port wrong. Continuing anyway.\n"
            )
        )

    results = []
    for atk in attacks:
        print(grey(f"\n  >> {atk.name} ({atk.proto})"))
        res = run_attack(ctx, atk)
        results.append(res)

    print_results(results)

    if args.json:
        with open(args.json, "w", encoding="utf-8") as fh:
            fh.write(to_json(results, ctx))
        print(grey(f"\n  JSON report written to {args.json}"))

    # Exit code: non-zero if any finding (CI gate).
    findings = sum(len(r.findings) for r in results)
    fails = sum(1 for r in results if r.status in (Status.FAIL, Status.ERROR))
    print()
    if fails or findings:
        print(red(f"  {fails} test(s) flagged, {findings} finding(s). Review above."))
        return 1
    print(green("  No findings. The device held every invariant under test."))
    return 0


if __name__ == "__main__":
    sys.exit(main())
