#!/usr/bin/env python3
"""HTTP download/upload test server for TR-143 diagnostics.

Serves both directions of Device.IP.Diagnostics.{Download,Upload}Diagnostics as
implemented by tr143d (curl based, see tr143d/scripts/{download,upload}).

    python3 tr143_server.py [--port 8000] [--save-dir DIR] [--default-size 1G]

Download routes (DownloadURL) -- all plain GET, response always ends with an
accurate Content-Length so size based tests terminate:

    /anything.bin?size=90MB     explicit size (suffixes: k/K, M/MB, G/GB)
    /testfile_90000000.bin      size taken from digits in the basename
    /100MB.bin                  same, with unit suffix
    /real/file/on/disk          served from --dir, when --dir is given
    /anything.bin               --default-size (1G) synthetic bytes

    /stream                     endless body, for time based tests
    /anything?size=inf          same as /stream

Time based tests (TimeBasedTestDuration > 0) run curl with --max-time=duration
and remap exit 28 to success, so the body must outlast the duration: either use
/stream or leave the default 1G size. The server never ends a transfer on a
timer of its own -- that would give curl exit 0 where the client expects 28.

Upload routes (UploadURL) -- PUT (curl -T -, chunked) or POST, any path:

    /upload.bin                 body drained and discarded, 200 with no body

The empty response body is mandatory: tr143d/scripts/upload runs curl without
-o and feeds the captured output straight to json_load, so anything the server
writes corrupts curl's -w JSON and the test reports Error_Other.

Pass --save-dir to write uploads to disk instead of discarding them.
"""

import argparse
import os
import socket
import sys
import time
from http import HTTPStatus
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import unquote, urlsplit, parse_qs

BLOCK = 256 * 1024
FILLER = bytes(BLOCK)
DEFAULT_UPLOAD_NAME = "speedtest_raw.bin"

# Set from the command line in main().
OPTS = None


def parse_size(text):
    """Parses '1000', '90MB', '1G', 'inf' into a byte count (None for endless)."""
    text = text.strip().lower()
    if text in ("inf", "infinite", "endless"):
        return None
    units = {"": 1, "b": 1, "k": 1000, "kb": 1000, "m": 1000 * 1000,
             "mb": 1000 * 1000, "g": 1000 * 1000 * 1000, "gb": 1000 * 1000 * 1000}
    digits = text
    suffix = ""
    while digits and not digits[-1].isdigit():
        suffix = digits[-1] + suffix
        digits = digits[:-1]
    if not digits or suffix not in units:
        raise ValueError("bad size %r" % text)
    return int(digits) * units[suffix]


def size_from_basename(path):
    """Extracts a size from the last digit run in the basename, e.g. 100MB.bin."""
    name = os.path.basename(path)
    stem = name.rsplit(".", 1)[0] if "." in name else name
    # Walk back over an optional unit suffix and then the digits.
    end = len(stem)
    while end > 0 and stem[end - 1].isalpha():
        end -= 1
    start = end
    while start > 0 and stem[start - 1].isdigit():
        start -= 1
    if start == end:
        return None
    try:
        return parse_size(stem[start:])
    except ValueError:
        return None


class TR143Handler(SimpleHTTPRequestHandler):
    # HTTP/1.1 is what makes handle_expect_100() fire; without it every upload
    # worker stalls on its Expect: 100-continue and the delay lands in BOMTime.
    protocol_version = "HTTP/1.1"
    server_version = "tr143_server/1.0"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=OPTS.dir, **kwargs)

    # -- logging ----------------------------------------------------------

    def log_message(self, fmt, *args):
        pass

    def log_error(self, fmt, *args):
        pass

    def report(self, nbytes, elapsed, status, note=""):
        if OPTS.quiet:
            return
        mbits = (nbytes * 8 / elapsed / 1e6) if elapsed > 0 else 0.0
        print("%s %-4s %-40s %-21s %13d B %7.3f s %9.2f Mbit/s %s %s"
              % (time.strftime("%H:%M:%S"), self.command, self.path[:40],
                 "%s:%s" % self.client_address[:2], nbytes, elapsed, mbits,
                 status, note),
              file=sys.stderr, flush=True)

    # -- helpers ----------------------------------------------------------

    def split_target(self):
        parts = urlsplit(self.path)
        return unquote(parts.path), parse_qs(parts.query)

    def safe_local_path(self, url_path):
        """Converts the URL request path into a safe local relative path."""
        normalized = os.path.normpath(url_path.lstrip("/"))
        if normalized.startswith("..") or os.path.isabs(normalized):
            raise ValueError("Directory traversal attempt blocked.")
        return normalized

    def reply(self, status, extra=None):
        """Sends a bodiless response. Never write a body: see module docstring."""
        self.send_response(status)
        self.send_header("Content-Length", "0")
        for key, value in (extra or {}).items():
            self.send_header(key, value)
        if self.close_connection:
            self.send_header("Connection", "close")
        self.end_headers()

    # -- download ---------------------------------------------------------

    def resolve_download(self, url_path, query):
        """Returns the byte count to send, or None for an endless body."""
        if "size" in query:
            return parse_size(query["size"][0])

        if url_path.rstrip("/").endswith("/stream"):
            return None

        size = size_from_basename(url_path)
        if size is not None:
            return size

        return OPTS.default_size

    def on_disk(self, url_path):
        if not OPTS.dir:
            return False
        try:
            local = self.safe_local_path(url_path)
        except ValueError:
            return False
        if not local:
            return False
        return os.path.isfile(os.path.join(OPTS.dir, local))

    def send_synthetic_headers(self, size):
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Accept-Ranges", "none")
        self.send_header("Cache-Control", "no-store")
        if size is None:
            # Endless body: no length is knowable, so delimit by close.
            self.close_connection = True
            self.send_header("Connection", "close")
        else:
            self.send_header("Content-Length", str(size))
        self.end_headers()

    def serve_download(self, body=True):
        url_path, query = self.split_target()

        if "size" not in query and self.on_disk(url_path):
            started = time.monotonic()
            local = os.path.join(OPTS.dir, self.safe_local_path(url_path))
            handler = super().do_GET if body else super().do_HEAD
            handler()
            self.report(os.path.getsize(local), time.monotonic() - started,
                        "200", "from disk: %s" % local)
            return

        try:
            size = self.resolve_download(url_path, query)
        except ValueError as exc:
            self.close_connection = True
            self.reply(HTTPStatus.BAD_REQUEST)
            self.report(0, 0, "400", str(exc))
            return

        if size is not None and size > OPTS.max_size:
            self.close_connection = True
            self.reply(HTTPStatus.BAD_REQUEST)
            self.report(0, 0, "400", "size above --max-size")
            return

        self.send_synthetic_headers(size)
        if not body:
            self.report(0, 0, "200", "HEAD")
            return

        started = time.monotonic()
        sent = 0
        note = ""
        try:
            if size is None:
                # Runs until the client aborts; never stops on its own timer.
                note = "endless"
                while True:
                    self.wfile.write(FILLER)
                    sent += BLOCK
            else:
                remaining = size
                while remaining > 0:
                    chunk = FILLER if remaining >= BLOCK else FILLER[:remaining]
                    self.wfile.write(chunk)
                    remaining -= len(chunk)
                    sent += len(chunk)
        except (BrokenPipeError, ConnectionResetError):
            # Expected on every time based test: curl aborts at --max-time.
            self.close_connection = True
            note = "client aborted"
        self.report(sent, time.monotonic() - started, "200", note)

    def do_GET(self):
        self.serve_download(body=True)

    def do_HEAD(self):
        self.serve_download(body=False)

    # -- upload -----------------------------------------------------------

    def drain(self, sink):
        """Consumes the whole request body, feeding it to sink. Returns bytes read."""
        total = 0
        encoding = self.headers.get("Transfer-Encoding", "").lower()

        if "chunked" in encoding:
            while True:
                line = self.rfile.readline(1024)
                if not line:
                    raise ValueError("truncated chunked body")
                try:
                    size = int(line.split(b";", 1)[0].strip(), 16)
                except ValueError:
                    raise ValueError("bad chunk size %r" % line)
                if size == 0:
                    # Trailers, terminated by an empty line.
                    while True:
                        trailer = self.rfile.readline(1024)
                        if not trailer or trailer in (b"\r\n", b"\n"):
                            break
                    return total
                remaining = size
                while remaining > 0:
                    data = self.rfile.read(min(remaining, BLOCK))
                    if not data:
                        raise ValueError("truncated chunk")
                    sink(data)
                    remaining -= len(data)
                    total += len(data)
                self.rfile.readline(1024)  # CRLF after chunk data

        if "Content-Length" in self.headers:
            try:
                remaining = int(self.headers["Content-Length"])
            except ValueError:
                raise ValueError("bad Content-Length")
            while remaining > 0:
                data = self.rfile.read(min(remaining, BLOCK))
                if not data:
                    raise ValueError("truncated body")
                sink(data)
                remaining -= len(data)
                total += len(data)
            return total

        # Neither framing: reading to EOF would hang under keep-alive.
        raise LookupError("no Content-Length and not chunked")

    def upload_target(self, url_path):
        local = self.safe_local_path(url_path)
        if not local or not os.path.basename(local):
            local = os.path.join(local, DEFAULT_UPLOAD_NAME)
        target = os.path.join(OPTS.save_dir, local)
        parent = os.path.dirname(target)
        if parent:
            os.makedirs(parent, exist_ok=True)
        return target

    def handle_upload(self):
        url_path, _ = self.split_target()
        started = time.monotonic()
        state = {"file": None}
        note = "discarded"
        try:
            if OPTS.save_dir:
                target = self.upload_target(url_path)
                note = "saved to %s" % target

                # Opened on the first byte so a rejected request leaves no file.
                def sink(data):
                    if state["file"] is None:
                        state["file"] = open(target, "wb")
                    state["file"].write(data)
            else:
                def sink(data):
                    pass

            nbytes = self.drain(sink)
        except LookupError as exc:
            self.close_connection = True
            self.reply(HTTPStatus.LENGTH_REQUIRED)
            self.report(0, time.monotonic() - started, "411", str(exc))
            return
        except (BrokenPipeError, ConnectionResetError):
            self.close_connection = True
            self.report(0, time.monotonic() - started, "-", "client aborted")
            return
        except Exception as exc:
            # The body framing is no longer trustworthy, so do not keep alive.
            self.close_connection = True
            self.reply(HTTPStatus.BAD_REQUEST)
            self.report(0, time.monotonic() - started, "400", str(exc))
            return
        finally:
            if state["file"]:
                state["file"].close()

        self.reply(HTTPStatus.OK)
        self.report(nbytes, time.monotonic() - started, "200", note)

    do_PUT = handle_upload
    do_POST = handle_upload


class TR143Server(ThreadingHTTPServer):
    daemon_threads = True
    allow_reuse_address = True
    request_queue_size = 64

    def __init__(self, addr, handler, ipv6=False):
        if ipv6:
            self.address_family = socket.AF_INET6
        super().__init__(addr, handler)

    def server_bind(self):
        if self.address_family == socket.AF_INET6:
            try:
                self.socket.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
            except OSError:
                pass
        super().server_bind()

    def handle_error(self, request, client_address):
        exc = sys.exc_info()[1]
        if isinstance(exc, (BrokenPipeError, ConnectionResetError, TimeoutError)):
            return
        super().handle_error(request, client_address)


def main():
    parser = argparse.ArgumentParser(
        description="HTTP download/upload test server for TR-143 diagnostics.")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--bind", default="::",
                        help="bind address; '::' (default) serves IPv4 and IPv6")
    parser.add_argument("--dir", default=None,
                        help="serve real files from this directory when the request "
                             "path matches one; off by default so that no-hint URLs "
                             "always get synthetic bytes of --default-size")
    parser.add_argument("--save-dir", default=None,
                        help="write uploads here instead of discarding them")
    parser.add_argument("--default-size", default="1G", type=parse_size,
                        help="synthetic download size when the URL gives no hint")
    parser.add_argument("--max-size", default="8G", type=parse_size,
                        help="refuse synthetic downloads larger than this")
    parser.add_argument("--quiet", action="store_true",
                        help="do not log one line per transfer")
    global OPTS
    OPTS = parser.parse_args()

    for name in ("default_size", "max_size"):
        if getattr(OPTS, name) is None:
            parser.error("--%s must be a finite byte count" % name.replace("_", "-"))
    if OPTS.save_dir:
        os.makedirs(OPTS.save_dir, exist_ok=True)

    ipv6 = ":" in OPTS.bind
    server = TR143Server((OPTS.bind, OPTS.port), TR143Handler, ipv6=ipv6)
    if not OPTS.quiet:
        print("tr143_server listening on [%s]:%d  (default download size: %d B, "
              "uploads: %s, real files: %s)"
              % (OPTS.bind, OPTS.port, OPTS.default_size,
                 OPTS.save_dir or "discarded", OPTS.dir or "off"),
              file=sys.stderr, flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()
