#!/usr/bin/env python3
# upstream_py: Python-based upstream HTTP server used by the
# paivana reverse-proxy tests.  Implements the same small set
# of canned endpoints as upstream_mhd.c.

import sys
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

UPSTREAM_NAME = "py"


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

    def log_message(self, fmt, *args):
        sys.stderr.write("upstream_py: " + (fmt % args) + "\n")

    def _send_text(self, code, body, content_type="text/plain"):
        if isinstance(body, str):
            body = body.encode("utf-8")
        self.send_response(code)
        self.send_header("X-Upstream", UPSTREAM_NAME)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _read_body(self):
        cl = int(self.headers.get("Content-Length", "0") or 0)
        if cl > 0:
            return self.rfile.read(cl)
        return b""

    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header("X-Upstream", UPSTREAM_NAME)
        self.send_header("Allow",
                         "GET, POST, PUT, PATCH, DELETE, HEAD, OPTIONS")
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_HEAD(self):
        if self.path == "/hello" or self.path.startswith("/large/"):
            # no body regardless
            body_len = 0
            if self.path.startswith("/large/"):
                try:
                    n = int(self.path[len("/large/"):])
                    body_len = max(0, min(n, 10 * 1024 * 1024))
                except ValueError:
                    body_len = 0
            self.send_response(200)
            self.send_header("X-Upstream", UPSTREAM_NAME)
            self.send_header("Content-Type", "application/octet-stream")
            self.send_header("Content-Length", str(body_len))
            self.end_headers()
            return
        self._send_text(404, "not found\n")

    def do_GET(self):
        if self.path == "/hello":
            self._send_text(200, f"Hello from {UPSTREAM_NAME}\n")
            return
        if self.path.startswith("/status/"):
            try:
                code = int(self.path[len("/status/"):])
            except ValueError:
                code = 500
            if code < 100 or code > 599:
                code = 500
            self._send_text(code, f"status {code}\n")
            return
        if self.path.startswith("/large/"):
            try:
                n = int(self.path[len("/large/"):])
            except ValueError:
                n = 0
            n = max(0, min(n, 10 * 1024 * 1024))
            buf = bytes(((ord('A') + i % 26) for i in range(n)))
            self._send_text(200, buf, "application/octet-stream")
            return
        if self.path.startswith("/slow/"):
            try:
                ms = int(self.path[len("/slow/"):])
            except ValueError:
                ms = 0
            ms = max(0, min(ms, 30000))
            time.sleep(ms / 1000.0)
            self._send_text(200, "slept\n")
            return
        if self.path == "/conn-response":
            # Name two of our own response headers in Connection: a
            # conforming proxy must strip both (RFC 9110 §7.6.1) but
            # keep X-Keep-Resp.  X-Hop-Before is deliberately emitted
            # *before* the Connection header that names it, X-Hop-After
            # *after* it, so the proxy cannot get away with a
            # single-pass filter.
            body = b"conn\n"
            self.send_response(200)
            self.send_header("X-Upstream", UPSTREAM_NAME)
            self.send_header("X-Hop-Before", "must-not-leak")
            self.send_header("Connection", "X-Hop-Before, X-Hop-After")
            self.send_header("X-Hop-After", "must-not-leak")
            self.send_header("X-Keep-Resp", "survivor")
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        if self.path == "/echo-headers":
            parts = []
            for k, v in self.headers.items():
                parts.append(f"{k}: {v}\n")
            self._send_text(200, "".join(parts))
            return
        self._send_text(404, "not found\n")

    def do_POST(self):
        body = self._read_body()
        if self.path == "/echo":
            self._send_text(200, body, "application/octet-stream")
            return
        if self.path == "/upload":
            self._send_text(200, f"Received {len(body)} bytes\n")
            return
        self._send_text(404, "not found\n")

    def do_PUT(self):
        body = self._read_body()
        if self.path == "/put":
            self._send_text(200, f"PUT received {len(body)}\n")
            return
        self._send_text(404, "not found\n")

    def do_PATCH(self):
        body = self._read_body()
        if self.path == "/patch":
            self._send_text(200, f"PATCH received {len(body)}\n")
            return
        self._send_text(404, "not found\n")

    def do_DELETE(self):
        # drain any body
        self._read_body()
        if self.path.startswith("/item"):
            self.send_response(204)
            self.send_header("X-Upstream", UPSTREAM_NAME)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        self._send_text(404, "not found\n")


def parse_port(arg):
    # int("0") would bind an ephemeral port and report itself
    # listening on 0, and int("garbage") would end in a traceback
    # rather than in a usable message.
    try:
        port = int(arg, 10)
    except ValueError:
        port = -1
    if port < 1 or port > 65535:
        sys.stderr.write(f"invalid port {arg!r}\n")
        sys.exit(1)
    return port


def main():
    port = 8403
    if len(sys.argv) > 1:
        port = parse_port(sys.argv[1])
    # Loopback, not "": this server echoes an arbitrary POST body and
    # hands out 10 MiB on request, and has no business being on the
    # network for the duration of `make check'.
    server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
    sys.stderr.write(f"upstream_py listening on port {port}\n")
    sys.stderr.flush()
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    server.server_close()


if __name__ == "__main__":
    main()
