#!/usr/bin/env python3
"""Velox hostile HTTP test server.

Every failure mode a real-world download hits, on purpose, so CORE's resume / retry /
segmentation logic can be tested without hunting for a broken CDN. Standard library only
(Python >= 3.11); no pip install, ever.

    tools/testserver/testserver.py --port 8080
    tools/testserver/testserver.py --port 0        # ephemeral; prints the port to stdout

Request shape
-------------
    GET /<mode>[+<mode>...]/file/<size>

`<size>` is a byte count with an optional K/M/G suffix (KiB/MiB/GiB): 1048576, 64K, 512M,
5G. The body is deterministic — byte i is `prng(seed, path, i)` — so a client can request
any range and verify it, and check the whole file against:

    GET /<mode>/sha256/<size>      -> {"sha256": "<hex>", "size": <n>}

`<mode>` (omit, or use `plain`, for the well-behaved server):

    no-range                  no Accept-Ranges, Range ignored, always 200 full body
    lies-about-accept-ranges  advertises Accept-Ranges: bytes but ignores Range (200)
    etag-changes              ETag differs every response; If-Range mismatch -> 200 full
    flaky-reset               TCP RST partway through the body; succeeds on the 3rd try
    slow-loris                headers/body dribbled out for --loris-seconds, then normal
    redirect-chain            302 x --redirect-depth before the real resource
    401-basic                 Basic auth, credentials test:test
    401-digest                Digest auth (qop=auth), credentials test:test
    403-without-referer       403 unless Referer names this server's origin
    416-always                Range -> 416; plain GET -> 200 (exercises the re-probe path)
    content-length-mismatch   Content-Length lies; connection closes short
    expiring-signed-url       needs ?exp=&sig=; past exp -> 403. See /<mode>/sign/<size>
    throttled                 body rate-limited to --throttle-bps
    chunked-no-length         Transfer-Encoding: chunked, no Content-Length
    utf8-content-disposition  Content-Disposition filename*=UTF-8''... (RFC 5987)
    legacy-content-disposition  MIME encoded-word + raw latin-1 filename (mojibake bait)

Control
-------
    GET  /__health            -> 200 "ok"
    POST /__control  {"reset": true}   clears per-path attempt counters (flaky-reset,
                                       redirect-chain retry state) between test cases
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import json
import os
import socket
import struct
import sys
import threading
import time
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import parse_qs, urlsplit

# --- deterministic synthetic content -------------------------------------------------

CHUNK = 64 * 1024


def _size_to_bytes(spec: str) -> int:
    spec = spec.strip()
    mult = 1
    if spec and spec[-1] in "kKmMgG":
        mult = {"k": 1024, "m": 1024**2, "g": 1024**3}[spec[-1].lower()]
        spec = spec[:-1]
    n = int(spec)
    if n < 0:
        raise ValueError("negative size")
    return n * mult


def _stream_body(seed: int, tag: str, start: int, end: int):
    """Yield the deterministic bytes for [start, end). Keyed by (seed, tag) so different
    paths have different content but a given path is stable across requests and ranges."""
    key = hashlib.sha256(f"{seed}:{tag}".encode()).digest()
    pos = start
    block = start // CHUNK
    while pos < end:
        h = hashlib.sha256(key + struct.pack("<Q", block)).digest()
        buf = (h * (CHUNK // len(h) + 1))[:CHUNK]
        off = pos - block * CHUNK
        take = min(CHUNK - off, end - pos)
        yield buf[off:off + take]
        pos += take
        block += 1


def _full_sha256(seed: int, tag: str, size: int) -> str:
    d = hashlib.sha256()
    for part in _stream_body(seed, tag, 0, size):
        d.update(part)
    return d.hexdigest()


# --- per-path attempt state (for flaky modes) --------------------------------------

class Attempts:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._n: dict[str, int] = {}

    def bump(self, key: str) -> int:
        with self._lock:
            self._n[key] = self._n.get(key, 0) + 1
            return self._n[key]

    def reset(self) -> None:
        with self._lock:
            self._n.clear()


# --- request handler --------------------------------------------------------------

class Handler(BaseHTTPRequestHandler):
    server_version = "veloxtestserver/1.0"
    protocol_version = "HTTP/1.1"

    # injected by make_server()
    seed: int = 0
    loris_seconds: float = 5.0
    redirect_depth: int = 5
    throttle_bps: int = 1024 * 1024
    attempts: Attempts = Attempts()
    verbose: bool = False

    # -- logging -------------------------------------------------------------------
    def log_message(self, fmt: str, *args) -> None:
        if self.verbose:
            sys.stderr.write("  %s - %s\n" % (self.address_string(), fmt % args))

    # -- helpers ------------------------------------------------------------------
    def _parts(self):
        u = urlsplit(self.path)
        segs = [s for s in u.path.split("/") if s != ""]
        query = parse_qs(u.query)
        return segs, query

    def _origin(self) -> str:
        host = self.headers.get("Host", f"127.0.0.1:{self.server.server_address[1]}")
        return f"http://{host}"

    def _send_simple(self, status: int, body: bytes = b"", ctype: str = "text/plain"):
        self.send_response(status)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if self.command != "HEAD":
            self.wfile.write(body)

    def _reset_connection(self) -> None:
        """Abort with a TCP RST rather than a clean FIN, the way a flaky CDN drops you."""
        try:
            self.connection.setsockopt(
                socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)
            )
            self.connection.close()
        except OSError:
            pass
        self.close_connection = True

    # -- entry points -----------------------------------------------------------
    def do_HEAD(self) -> None:
        self._dispatch()

    def do_GET(self) -> None:
        self._dispatch()

    def do_POST(self) -> None:
        segs, _ = self._parts()
        if segs == ["__control"]:
            length = int(self.headers.get("Content-Length", "0") or "0")
            raw = self.rfile.read(length) if length else b"{}"
            try:
                msg = json.loads(raw or b"{}")
            except json.JSONDecodeError:
                self._send_simple(400, b"bad json")
                return
            if msg.get("reset"):
                self.attempts.reset()
            self._send_simple(200, b'{"ok":true}', "application/json")
            return
        self._send_simple(405, b"method not allowed")

    # -- routing --------------------------------------------------------------
    def _dispatch(self) -> None:
        try:
            segs, query = self._parts()
        except ValueError:
            self._send_simple(400, b"bad path")
            return

        if segs == ["__health"]:
            self._send_simple(200, b"ok")
            return
        if not segs:
            self._send_simple(200, b"velox testserver: see --help / module docstring")
            return

        # /redirect/<n>/<rest...> — internal hop target for redirect-chain
        if segs[0] == "redirect":
            self._serve_redirect_hop(segs, query)
            return

        modes = set(segs[0].split("+")) if segs[0] not in ("file", "sha256", "sign") else {"plain"}
        rest = segs if segs[0] in ("file", "sha256", "sign") else segs[1:]
        if not rest:
            self._send_simple(400, b"expected /<mode>/file/<size>")
            return

        kind = rest[0]
        arg = rest[1] if len(rest) > 1 else ""

        if kind == "sha256":
            try:
                size = _size_to_bytes(arg)
            except ValueError:
                self._send_simple(400, b"bad size")
                return
            tag = f"{'+'.join(sorted(modes))}/file/{arg}"
            body = json.dumps({"sha256": _full_sha256(self.seed, tag, size), "size": size}).encode()
            self._send_simple(200, body, "application/json")
            return

        if kind == "sign":
            self._serve_sign(modes, arg, query)
            return

        if kind != "file":
            self._send_simple(404, b"unknown resource")
            return

        try:
            size = _size_to_bytes(arg)
        except ValueError:
            self._send_simple(400, b"bad size")
            return

        self._serve_file(modes, arg, size, query)

    # -- redirect-chain ------------------------------------------------------
    def _serve_redirect_hop(self, segs: list[str], query) -> None:
        # /redirect/<n>/<mode>/file/<size>
        try:
            n = int(segs[1])
        except (IndexError, ValueError):
            self._send_simple(400, b"bad redirect hop")
            return
        tail = "/".join(segs[2:])
        q = urlsplit(self.path).query
        qs = f"?{q}" if q else ""
        self.send_response(302)
        if n <= 0:
            self.send_header("Location", f"{self._origin()}/{tail}{qs}")
        else:
            self.send_header("Location", f"{self._origin()}/redirect/{n - 1}/{tail}{qs}")
        self.send_header("Content-Length", "0")
        self.end_headers()

    # -- signed URL helper -------------------------------------------------
    def _sign(self, path: str, exp: int) -> str:
        return hmac.new(
            struct.pack("<q", self.seed), f"{path}:{exp}".encode(), hashlib.sha256
        ).hexdigest()

    def _serve_sign(self, modes: set[str], arg: str, query) -> None:
        ttl = int((query.get("ttl", ["10"])[0]))
        exp = int(time.time()) + ttl
        path = f"/{'+'.join(sorted(modes))}/file/{arg}"
        sig = self._sign(path, exp)
        url = f"{self._origin()}{path}?exp={exp}&sig={sig}"
        self._send_simple(200, json.dumps({"url": url, "exp": exp}).encode(), "application/json")

    # -- the main file route ---------------------------------------------
    def _serve_file(self, modes: set[str], size_spec: str, size: int, query) -> None:
        tag = f"{'+'.join(sorted(modes))}/file/{size_spec}"
        path_key = urlsplit(self.path).path

        # --- auth gates (checked before anything else) ---
        if "401-basic" in modes and not self._basic_ok():
            self.send_response(401)
            self.send_header("WWW-Authenticate", 'Basic realm="velox-test"')
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        if "401-digest" in modes and not self._digest_ok():
            nonce = hashlib.md5(f"{time.time()}:{self.seed}".encode()).hexdigest()
            self.send_response(401)
            self.send_header(
                "WWW-Authenticate",
                f'Digest realm="velox-test", qop="auth", nonce="{nonce}", '
                f'opaque="{hashlib.md5(b"velox").hexdigest()}"',
            )
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        if "403-without-referer" in modes:
            ref = self.headers.get("Referer", "")
            if not ref or self._origin() not in ref:
                self._send_simple(403, b"referer required")
                return
        if "expiring-signed-url" in modes:
            exp = query.get("exp", [None])[0]
            sig = query.get("sig", [None])[0]
            if exp is None or sig is None:
                self._send_simple(403, b'{"error":"unsigned"}', "application/json")
                return
            want = self._sign(f"/{'+'.join(sorted(modes))}/file/{size_spec}", int(exp))
            if not hmac.compare_digest(sig, want):
                self._send_simple(403, b'{"error":"bad signature"}', "application/json")
                return
            if int(exp) < time.time():
                self._send_simple(403, b'{"error":"expired"}', "application/json")
                return

        # --- redirect-chain: bounce before serving ---
        if "redirect-chain" in modes and query.get("_r", ["0"])[0] != "done":
            depth = self.redirect_depth
            tail = f"{'+'.join(sorted(modes))}/file/{size_spec}"
            self.send_response(302)
            self.send_header("Location", f"{self._origin()}/redirect/{depth - 1}/{tail}?_r=done")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        # --- 416-always ---
        rng = self._parse_range(size)
        if "416-always" in modes and rng is not None:
            self.send_response(416)
            self.send_header("Content-Range", f"bytes */{size}")
            self.send_header("Content-Length", "0")
            self.end_headers()
            return

        # --- ranges honoured? ---
        honour_range = not ({"no-range", "lies-about-accept-ranges"} & modes)
        advertise_ar = "no-range" not in modes

        etag = self._etag(tag, changing="etag-changes" in modes)
        if rng is not None and honour_range and "etag-changes" in modes:
            if_range = self.headers.get("If-Range")
            if if_range and if_range != etag:
                rng = None  # validator failed -> full 200, mirrors a changed file

        start, end = (0, size)
        partial = False
        if rng is not None and honour_range:
            start, end = rng
            partial = True

        status = 206 if partial else 200
        body_len = end - start

        # --- content-length-mismatch: declare a wrong length ---
        declared_len = body_len
        short_by = 0
        if "content-length-mismatch" in modes:
            short_by = min(4096, body_len // 2 + 1)
            declared_len = body_len            # header says the real length...
            body_len_to_send = body_len - short_by  # ...but we send less and hang up

        # --- flaky-reset: fail the first two attempts ---
        if "flaky-reset" in modes:
            n = self.attempts.bump(f"reset:{path_key}:{start}-{end}")
            if n % 3 != 0:
                cut = max(1, (end - start) // 2)
                self.send_response(status)
                self._common_headers(size, start, end, partial, advertise_ar, etag,
                                     declared_len)
                self.end_headers()
                if self.command != "HEAD":
                    self._write_body(tag, start, start + cut, throttle=False)
                self._reset_connection()
                return

        # --- chunked-no-length ---
        if "chunked-no-length" in modes and not partial:
            self.send_response(200)
            self.send_header("Content-Type", "application/octet-stream")
            self.send_header("Transfer-Encoding", "chunked")
            self._content_disposition(modes)
            self.end_headers()
            if self.command != "HEAD":
                for part in self._body_iter(tag, start, end,
                                            throttle="throttled" in modes):
                    self.wfile.write(b"%X\r\n%s\r\n" % (len(part), part))
                self.wfile.write(b"0\r\n\r\n")
            return

        # --- slow-loris: dribble the response ---
        if "slow-loris" in modes:
            self._serve_loris(tag, size, start, end, partial, advertise_ar, etag)
            return

        # --- normal (well-behaved, or throttled, or lying-about-length) path ---
        self.send_response(status)
        self._common_headers(size, start, end, partial, advertise_ar, etag, declared_len)
        self._content_disposition(modes)
        self.end_headers()
        if self.command == "HEAD":
            return
        if "content-length-mismatch" in modes:
            self._write_body(tag, start, start + body_len_to_send, throttle=False)
            self._reset_connection()
            return
        self._write_body(tag, start, end, throttle="throttled" in modes)

    # -- header / body plumbing ----------------------------------------
    def _common_headers(self, size, start, end, partial, advertise_ar, etag, declared_len):
        self.send_header("Content-Type", "application/octet-stream")
        if advertise_ar:
            self.send_header("Accept-Ranges", "bytes")
        self.send_header("ETag", etag)
        self.send_header("Last-Modified", "Wed, 01 Jan 2025 00:00:00 GMT")
        if partial:
            self.send_header("Content-Range", f"bytes {start}-{end - 1}/{size}")
        self.send_header("Content-Length", str(declared_len))

    def _content_disposition(self, modes: set[str]) -> None:
        if "utf8-content-disposition" in modes:
            # € rates.pdf
            self.send_header(
                "Content-Disposition",
                "attachment; filename=\"rates.pdf\"; "
                "filename*=UTF-8''%E2%82%AC%20rates.pdf",
            )
        elif "legacy-content-disposition" in modes:
            # MIME encoded-word (RFC 2047) + a raw Latin-1 fallback: classic mojibake bait
            self.send_header(
                "Content-Disposition",
                'attachment; filename="=?UTF-8?B?xI1lc2vDoS1zbcOsxJlz.pdf?="',
            )

    def _body_iter(self, tag: str, start: int, end: int, throttle: bool):
        budget = self.throttle_bps
        window_start = time.monotonic()
        sent_in_window = 0
        for part in _stream_body(self.seed, tag, start, end):
            if throttle:
                sent_in_window += len(part)
                if sent_in_window >= budget:
                    elapsed = time.monotonic() - window_start
                    if elapsed < 1.0:
                        time.sleep(1.0 - elapsed)
                    window_start = time.monotonic()
                    sent_in_window = 0
            yield part

    def _write_body(self, tag: str, start: int, end: int, throttle: bool) -> None:
        try:
            for part in self._body_iter(tag, start, end, throttle):
                self.wfile.write(part)
        except (BrokenPipeError, ConnectionResetError):
            self.close_connection = True

    def _serve_loris(self, tag, size, start, end, partial, advertise_ar, etag) -> None:
        deadline = time.monotonic() + self.loris_seconds
        status_line = f"HTTP/1.1 {206 if partial else 200} X\r\n"
        self.wfile.write(status_line.encode())
        headers = [
            ("Content-Type", "application/octet-stream"),
            ("ETag", etag),
        ]
        if advertise_ar:
            headers.append(("Accept-Ranges", "bytes"))
        if partial:
            headers.append(("Content-Range", f"bytes {start}-{end - 1}/{size}"))
        headers.append(("Content-Length", str(end - start)))
        headers.append(("Connection", "close"))
        self.close_connection = True
        for k, v in headers:
            line = f"{k}: {v}\r\n".encode()
            for b in line:
                self.wfile.write(bytes([b]))
                self.wfile.flush()
                if time.monotonic() < deadline:
                    time.sleep(0.2)
        self.wfile.write(b"\r\n")
        if self.command == "HEAD":
            return
        slow = True
        for part in _stream_body(self.seed, tag, start, end):
            if slow and time.monotonic() < deadline:
                for b in part:
                    self.wfile.write(bytes([b]))
                    self.wfile.flush()
                    time.sleep(0.05)
                    if time.monotonic() >= deadline:
                        slow = False
                        break
            else:
                try:
                    self.wfile.write(part)
                except (BrokenPipeError, ConnectionResetError):
                    return

    # -- auth --------------------------------------------------------------
    def _basic_ok(self) -> bool:
        h = self.headers.get("Authorization", "")
        if not h.startswith("Basic "):
            return False
        try:
            user, _, pw = base64.b64decode(h[6:]).decode().partition(":")
        except Exception:
            return False
        return user == "test" and pw == "test"

    def _digest_ok(self) -> bool:
        h = self.headers.get("Authorization", "")
        if not h.startswith("Digest "):
            return False
        params = {}
        for item in h[7:].split(","):
            if "=" not in item:
                continue
            k, v = item.strip().split("=", 1)
            params[k] = v.strip('"')
        need = {"username", "realm", "nonce", "uri", "response"}
        if not need.issubset(params) or params["username"] != "test":
            return False
        ha1 = hashlib.md5(f"test:{params['realm']}:test".encode()).hexdigest()
        ha2 = hashlib.md5(f"{self.command}:{params['uri']}".encode()).hexdigest()
        if params.get("qop") == "auth":
            resp = hashlib.md5(
                f"{ha1}:{params['nonce']}:{params.get('nc','')}:"
                f"{params.get('cnonce','')}:auth:{ha2}".encode()
            ).hexdigest()
        else:
            resp = hashlib.md5(f"{ha1}:{params['nonce']}:{ha2}".encode()).hexdigest()
        return hmac.compare_digest(resp, params["response"])

    # -- misc ------------------------------------------------------------
    def _etag(self, tag: str, changing: bool) -> str:
        if changing:
            return '"' + hashlib.md5(f"{tag}:{time.time_ns()}".encode()).hexdigest() + '"'
        return '"' + hashlib.md5(tag.encode()).hexdigest() + '"'

    def _parse_range(self, size: int):
        h = self.headers.get("Range")
        if not h or not h.startswith("bytes="):
            return None
        spec = h[6:].split(",")[0].strip()
        try:
            if spec.startswith("-"):
                n = int(spec[1:])
                return (max(0, size - n), size)
            lo_s, _, hi_s = spec.partition("-")
            lo = int(lo_s)
            hi = int(hi_s) + 1 if hi_s else size
            if lo >= size or lo < 0 or hi > size or lo >= hi:
                return None
            return (lo, hi)
        except ValueError:
            return None


def make_server(host: str, port: int, args) -> ThreadingHTTPServer:
    attempts = Attempts()

    class Bound(Handler):
        pass

    Bound.seed = args.seed
    Bound.loris_seconds = args.loris_seconds
    Bound.redirect_depth = args.redirect_depth
    Bound.throttle_bps = args.throttle_bps
    Bound.attempts = attempts
    Bound.verbose = args.verbose

    httpd = ThreadingHTTPServer((host, port), Bound)
    httpd.daemon_threads = True
    return httpd


def main() -> int:
    p = argparse.ArgumentParser(description="Velox hostile HTTP test server")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=8080, help="0 for an ephemeral port")
    p.add_argument("--seed", type=int, default=1, help="deterministic content seed")
    p.add_argument("--loris-seconds", type=float, default=5.0)
    p.add_argument("--redirect-depth", type=int, default=5)
    p.add_argument("--throttle-bps", type=int, default=1024 * 1024)
    p.add_argument("--verbose", action="store_true")
    args = p.parse_args()

    httpd = make_server(args.host, args.port, args)
    actual_port = httpd.server_address[1]
    print(f"{actual_port}", flush=True)
    sys.stderr.write(
        f"velox testserver on http://{args.host}:{actual_port}  seed={args.seed}\n"
    )
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        httpd.shutdown()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())