mrok 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

mrok/http/forwarder.py

@@ -0,0 +1,299 @@
+import abc
+import asyncio
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+logger = logging.getLogger("mrok.proxy")
+
+Scope = dict[str, Any]
+ASGIReceive = Callable[[], Awaitable[dict[str, Any]]]
+ASGISend = Callable[[dict[str, Any]], Awaitable[None]]
+
+
+class BackendNotFoundError(Exception):
+    pass
+
+
+class ForwardAppBase(abc.ABC):
+    """Generic HTTP forwarder base class.
+
+    Subclasses must implement `select_backend(scope)` to return an
+    (asyncio.StreamReader, asyncio.StreamWriter) pair connected to the
+    desired backend. The base class implements the HTTP/1.1 framing
+    and streaming logic (requests and responses).
+    """
+
+    def __init__(self, read_chunk_size: int = 65536) -> None:
+        # number of bytes to read per iteration when streaming bodies
+        self._read_chunk_size: int = int(read_chunk_size)
+
+    @abc.abstractmethod
+    async def select_backend(
+        self,
+        scope: Scope,
+        headers: dict[str, str],
+    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter] | tuple[None, None]:
+        """Return (reader, writer) connected to the target backend."""
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend) -> None:
+        """ASGI callable entry point.
+
+        Delegates to smaller helper methods for readability. Subclasses only
+        need to implement backend selection.
+        """
+        if scope.get("type") != "http":
+            await send({"type": "http.response.start", "status": 500, "headers": []})
+            await send({"type": "http.response.body", "body": b"Unsupported"})
+            return
+
+        method = scope.get("method", "GET")
+        path_qs = self.format_path(scope)
+
+        headers = list(scope.get("headers", []))
+        headers = self.ensure_host_header(headers, scope)
+        reader, writer = await self.select_backend(
+            scope, {k[0].decode().lower(): k[1].decode() for k in headers}
+        )
+
+        if not (reader and writer):
+            await send({"type": "http.response.start", "status": 502, "headers": []})
+            await send({"type": "http.response.body", "body": b"Bad Gateway"})
+            return
+
+        use_chunked = self.ensure_chunked_if_needed(headers)
+
+        await self.write_request_line_and_headers(writer, method, path_qs, headers)
+
+        await self.stream_request_body(receive, writer, use_chunked)
+
+        status_line = await reader.readline()
+        if not status_line:
+            await send({"type": "http.response.start", "status": 502, "headers": []})
+            await send({"type": "http.response.body", "body": b"Bad Gateway"})
+            writer.close()
+            await writer.wait_closed()
+            return
+
+        status, headers_out, raw_headers = await self.read_status_and_headers(reader, status_line)
+
+        await send({"type": "http.response.start", "status": status, "headers": headers_out})
+
+        await self.stream_response_body(reader, send, raw_headers)
+
+        writer.close()
+        await writer.wait_closed()
+
+    def format_path(self, scope: Scope) -> str:
+        raw_path = scope.get("raw_path")
+        if raw_path:
+            return raw_path.decode()
+        q = scope.get("query_string", b"")
+        path = scope.get("path", "/")
+        path_qs = path
+        if q:
+            path_qs += "?" + q.decode()
+        return path_qs
+
+    def ensure_host_header(
+        self, headers: list[tuple[bytes, bytes]], scope: Scope
+    ) -> list[tuple[bytes, bytes]]:
+        if any(n.lower() == b"host" for n, _ in headers):
+            return headers
+        server = scope.get("server")
+        if server:
+            host = f"{server[0]}:{server[1]}" if server[1] else server[0]
+            headers.append((b"host", host.encode()))
+        return headers
+
+    def ensure_chunked_if_needed(self, headers: list[tuple[bytes, bytes]]) -> bool:
+        has_content_length = any(n.lower() == b"content-length" for n, _ in headers)
+        has_transfer_encoding = any(n.lower() == b"transfer-encoding" for n, _ in headers)
+        if not has_content_length and not has_transfer_encoding:
+            headers.append((b"transfer-encoding", b"chunked"))
+            return True
+        return False
+
+    async def write_request_line_and_headers(
+        self,
+        writer: asyncio.StreamWriter,
+        method: str,
+        path_qs: str,
+        headers: list[tuple[bytes, bytes]],
+    ) -> None:
+        writer.write(f"{method} {path_qs} HTTP/1.1\r\n".encode())
+        for name, value in headers:
+            if name.lower() == b"expect":
+                continue
+            writer.write(name + b": " + value + b"\r\n")
+        writer.write(b"\r\n")
+        await writer.drain()
+
+    async def stream_request_body(
+        self, receive: ASGIReceive, writer: asyncio.StreamWriter, use_chunked: bool
+    ) -> None:
+        if use_chunked:
+            await self.stream_request_chunked(receive, writer)
+            return
+
+        await self.stream_request_until_end(receive, writer)
+
+    async def stream_request_chunked(
+        self, receive: ASGIReceive, writer: asyncio.StreamWriter
+    ) -> None:
+        """Send request body to backend using HTTP/1.1 chunked encoding."""
+        while True:
+            event = await receive()
+            if event["type"] == "http.request":
+                body = event.get("body", b"") or b""
+                if body:
+                    writer.write(f"{len(body):X}\r\n".encode())
+                    writer.write(body)
+                    writer.write(b"\r\n")
+                    await writer.drain()
+                if not event.get("more_body", False):
+                    break
+            elif event["type"] == "http.disconnect":
+                writer.close()
+                return
+
+        writer.write(b"0\r\n\r\n")
+        await writer.drain()
+
+    async def stream_request_until_end(
+        self, receive: ASGIReceive, writer: asyncio.StreamWriter
+    ) -> None:
+        """Send request body to backend when content length/transfer-encoding
+        already provided (no chunking).
+        """
+        while True:
+            event = await receive()
+            if event["type"] == "http.request":
+                body = event.get("body", b"") or b""
+                if body:
+                    writer.write(body)
+                    await writer.drain()
+                if not event.get("more_body", False):
+                    break
+            elif event["type"] == "http.disconnect":
+                writer.close()
+                return
+
+    async def read_status_and_headers(
+        self, reader: asyncio.StreamReader, first_line: bytes
+    ) -> tuple[int, list[tuple[bytes, bytes]], dict[bytes, bytes]]:
+        parts = first_line.decode(errors="ignore").split(" ", 2)
+        status = int(parts[1]) if len(parts) >= 2 and parts[1].isdigit() else 502
+        headers: list[tuple[bytes, bytes]] = []
+        raw_headers: dict[bytes, bytes] = {}
+        while True:
+            line = await reader.readline()
+            if line in (b"\r\n", b"\n", b""):
+                break
+            i = line.find(b":")
+            if i == -1:
+                continue
+            name = line[:i].strip().lower()
+            value = line[i + 1 :].strip()
+            headers.append((name, value))
+            raw_headers[name] = value
+
+        return status, headers, raw_headers
+
+    def is_chunked(self, te_value: bytes) -> bool:
+        """Return True if transfer-encoding header tokens include 'chunked'."""
+        if not te_value:
+            return False
+        # split on commas, strip spaces and check tokens
+        tokens = [t.strip() for t in te_value.split(b",")]
+        return any(t.lower() == b"chunked" for t in tokens)
+
+    def parse_content_length(self, cl_value: bytes | None) -> int | None:
+        """Parse Content-Length header value to int, or return None if invalid."""
+        if cl_value is None:
+            return None
+        try:
+            return int(cl_value)
+        except Exception:
+            return None
+
+    async def drain_trailers(self, reader: asyncio.StreamReader) -> None:
+        """Consume trailer header lines until an empty line is encountered."""
+        while True:
+            trailer = await reader.readline()
+            if trailer in (b"\r\n", b"\n", b""):
+                break
+
+    async def stream_response_chunked(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
+        """Read chunked-encoded response from reader, decode and forward to ASGI send."""
+        while True:
+            size_line = await reader.readline()
+            if not size_line:
+                break
+            size_str = size_line.split(b";", 1)[0].strip()
+            try:
+                size = int(size_str, 16)
+            except Exception:
+                break
+            if size == 0:
+                # consume trailers
+                await self.drain_trailers(reader)
+                break
+            try:
+                chunk = await reader.readexactly(size)
+            except Exception:
+                break
+            # consume the CRLF after the chunk
+            try:
+                await reader.readexactly(2)
+            except Exception:
+                logger.warning("failed to read CRLF after chunk from backend")
+            await send({"type": "http.response.body", "body": chunk, "more_body": True})
+
+        await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+    async def stream_response_with_content_length(
+        self, reader: asyncio.StreamReader, send: ASGISend, content_length: int
+    ) -> None:
+        """Read exactly content_length bytes and forward to ASGI send events."""
+        remaining = content_length
+        sent_final = False
+        while remaining > 0:
+            to_read = min(self._read_chunk_size, remaining)
+            chunk = await reader.read(to_read)
+            if not chunk:
+                break
+            remaining -= len(chunk)
+            more = remaining > 0
+            await send({"type": "http.response.body", "body": chunk, "more_body": more})
+            if not more:
+                sent_final = True
+
+        if not sent_final:
+            await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+    async def stream_response_until_eof(self, reader: asyncio.StreamReader, send: ASGISend) -> None:
+        """Read from reader until EOF and forward chunks to ASGI send events."""
+        while True:
+            chunk = await reader.read(self._read_chunk_size)
+            if not chunk:
+                break
+            await send({"type": "http.response.body", "body": chunk, "more_body": True})
+        await send({"type": "http.response.body", "body": b"", "more_body": False})
+
+    async def stream_response_body(
+        self, reader: asyncio.StreamReader, send: ASGISend, raw_headers: dict[bytes, bytes]
+    ) -> None:
+        te = raw_headers.get(b"transfer-encoding", b"").lower()
+        cl = raw_headers.get(b"content-length")
+
+        if self.is_chunked(te):
+            await self.stream_response_chunked(reader, send)
+            return
+
+        content_length = self.parse_content_length(cl)
+        if content_length is not None:
+            await self.stream_response_with_content_length(reader, send, content_length)
+            return
+
+        await self.stream_response_until_eof(reader, send)

mrok/http/lifespan.py

@@ -0,0 +1,10 @@
+import logging
+
+from uvicorn.config import Config
+from uvicorn.lifespan.on import LifespanOn
+
+
+class MrokLifespan(LifespanOn):
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+        self.logger = logging.getLogger("mrok.proxy")

mrok/http/master.py

@@ -0,0 +1,90 @@
+import logging
+import os
+import signal
+import threading
+from collections.abc import Callable
+from pathlib import Path
+
+from watchfiles import watch
+from watchfiles.filters import PythonFilter
+from watchfiles.run import CombinedProcess, start_process
+
+logger = logging.getLogger("mrok.agent")
+
+
+def print_path(path):
+    try:
+        return f'"{path.relative_to(Path.cwd())}"'
+    except ValueError:
+        return f'"{path}"'
+
+
+class Master:
+    def __init__(
+        self,
+        start_fn: Callable,
+        workers: int,
+        reload: bool,
+    ):
+        self.start_fn = start_fn
+        self.workers = workers
+        self.reload = reload
+        self.worker_processes: list[CombinedProcess] = []
+        self.stop_event = threading.Event()
+        self.watch_filter = PythonFilter(ignore_paths=None)
+        self.watcher = watch(
+            Path.cwd(),
+            watch_filter=self.watch_filter,
+            stop_event=self.stop_event,
+            yield_on_timeout=True,
+        )
+        self.setup_signals_handler()
+
+    def setup_signals_handler(self):
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            signal.signal(sig, self.handle_signal)
+
+    def handle_signal(self, *args, **kwargs):
+        self.stop_event.set()
+
+    def start(self):
+        for _ in range(self.workers):
+            p = start_process(
+                self.start_fn,
+                "function",
+                (),
+                None,
+            )
+            logger.info(f"Worker [{p.pid}] started")
+            self.worker_processes.append(p)
+
+    def stop(self):
+        for process in self.worker_processes:
+            process.stop(sigint_timeout=5, sigkill_timeout=1)
+        self.worker_processes = []
+
+    def restart(self):
+        self.stop()
+        self.start()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        changes = next(self.watcher)
+        if changes:
+            return list({Path(change[1]) for change in changes})
+        return None
+
+    def run(self):
+        self.start()
+        logger.info(f"Master process started: {os.getpid()}")
+        if self.reload:
+            for files_changed in self:
+                if files_changed:
+                    logger.warning(
+                        f"{', '.join(map(print_path, files_changed))} changed, reloading...",
+                    )
+                    self.restart()
+        else:
+            self.stop_event.wait()

mrok/http/protocol.py

@@ -0,0 +1,11 @@
+import logging
+
+from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol
+
+
+class MrokHttpToolsProtocol(HttpToolsProtocol):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logger = logging.getLogger("mrok.proxy")
+        self.access_logger = logging.getLogger("mrok.access")
+        self.access_log = self.access_logger.hasHandlers()

mrok/http/server.py

@@ -0,0 +1,14 @@
+import logging
+import socket
+
+from uvicorn import server
+
+server.logger = logging.getLogger("mrok.proxy")
+
+
+class MrokServer(server.Server):
+    async def serve(self, sockets: list[socket.socket] | None = None) -> None:
+        if not sockets:
+            sockets = [self.config.bind_socket()]
+        with self.capture_signals():
+            await self._serve(sockets)

mrok/logging.py

@@ -0,0 +1,76 @@
+import logging.config
+
+from mrok.conf import Settings
+
+
+def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict:
+    log_level = "DEBUG" if settings.logging.debug else "INFO"
+    handler = "rich" if settings.logging.rich else "console"
+
+    if cli_mode:
+        mrok_handler = "cli"
+    else:
+        mrok_handler = handler
+
+    logging_config = {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "verbose": {
+                "format": "{asctime} {name} {levelname} (pid: {process}) {message}",
+                "style": "{",
+            },
+            "rich": {
+                "format": "{name} {message}",
+                "style": "{",
+            },
+            "plain": {"format": "%(message)s"},
+        },
+        "handlers": {
+            "console": {
+                "class": "logging.StreamHandler",
+                "formatter": "verbose",
+                "stream": "ext://sys.stderr",
+            },
+            "cli": {
+                "class": "logging.StreamHandler",
+                "formatter": "plain",
+                "stream": "ext://sys.stdout",
+            },
+            "rich": {
+                "class": "rich.logging.RichHandler",
+                "level": log_level,
+                "formatter": "rich",
+                "log_time_format": lambda x: x.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
+                "rich_tracebacks": True,
+            },
+        },
+        "root": {
+            "handlers": ["rich"],
+            "level": "WARNING",
+        },
+        "loggers": {
+            "gunicorn.access": {
+                "handlers": [handler],
+                "level": log_level,
+                "propagate": False,
+            },
+            "gunicorn.error": {
+                "handlers": [handler],
+                "level": log_level,
+                "propagate": False,
+            },
+            "mrok": {
+                "handlers": [mrok_handler],
+                "level": log_level,
+                "propagate": False,
+            },
+        },
+    }
+
+    return logging_config
+
+
+def setup_logging(settings: Settings, cli_mode: bool = False) -> None:
+    logging_config = get_logging_config(settings, cli_mode)
+    logging.config.dictConfig(logging_config)

mrok/ziti/__init__.py

@@ -0,0 +1,15 @@
+from mrok.ziti.api import (
+    ZitiAPIError,
+    ZitiAuthError,
+    ZitiBadRequestError,
+    ZitiClientAPI,
+    ZitiManagementAPI,
+)
+
+__all__ = [
+    "ZitiAPIError",
+    "ZitiAuthError",
+    "ZitiBadRequestError",
+    "ZitiClientAPI",
+    "ZitiManagementAPI",
+]