mrok 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

mrok/controller/auth/manager.py

@@ -0,0 +1,31 @@
+from dynaconf.utils.boxing import DynaBox
+from fastapi import Request
+
+from mrok.controller.auth.base import UNAUTHORIZED_EXCEPTION, AuthIdentity, BaseHTTPAuthBackend
+from mrok.controller.auth.registry import get_authentication_backend
+
+
+class HTTPAuthManager:
+    def __init__(self, auth_settings: DynaBox):
+        self.auth_settings = auth_settings
+        self.active_backends: list[BaseHTTPAuthBackend] = []
+        self._setup_backends()
+
+    def _setup_backends(self):
+        enabled_keys = self.auth_settings.get("backends", [])
+
+        for key in enabled_keys:
+            backend_cls = get_authentication_backend(key)
+            if not backend_cls:
+                raise ValueError(f"Backend '{key}' is not registered.")
+
+            specific_config = self.auth_settings.get(key, {})
+            self.active_backends.append(backend_cls(specific_config))
+
+    async def __call__(self, request: Request) -> AuthIdentity:
+        for backend in self.active_backends:
+            identity = await backend(request)
+            if identity:
+                return identity
+
+        raise UNAUTHORIZED_EXCEPTION

mrok/controller/auth/registry.py

@@ -0,0 +1,17 @@
+from mrok.controller.auth.base import BaseHTTPAuthBackend
+
+BACKEND_REGISTRY: dict[str, type[BaseHTTPAuthBackend]] = {}
+
+
+def register_authentication_backend(name: str):
+    """Decorator to register a backend class with a unique key."""
+
+    def decorator(cls: type[BaseHTTPAuthBackend]):
+        BACKEND_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def get_authentication_backend(name: str) -> type[BaseHTTPAuthBackend] | None:
+    return BACKEND_REGISTRY.get(name)

mrok/frontend/app.py

@@ -1,14 +1,21 @@
-import re
+from http import HTTPStatus
+from pathlib import Path
+from typing import Any
 
 from httpcore import AsyncConnectionPool
+from jinja2 import Environment, FileSystemLoader, select_autoescape
 
 from mrok.conf import get_settings
+from mrok.frontend.utils import get_target_name, parse_accept_header
 from mrok.proxy.app import ProxyAppBase
 from mrok.proxy.backend import AIOZitiNetworkBackend
 from mrok.proxy.exceptions import InvalidTargetError
-from mrok.types.proxy import Scope
+from mrok.types.proxy import ASGISend, Scope
 
-RE_SUBDOMAIN = re.compile(r"(?i)^(?:EXT-\d{4}-\d{4}|INS-\d{4}-\d{4}-\d{4})$")
+ERROR_TEMPLATE_FORMATS = {
+    "application/json": "json",
+    "text/html": "html",
+}
 
 
 class FrontendProxyApp(ProxyAppBase):
@@ -19,10 +26,11 @@ class FrontendProxyApp(ProxyAppBase):
         max_connections: int | None = 10,
         max_keepalive_connections: int | None = None,
         keepalive_expiry: float | None = None,
-        retries=0,
+        retries: int = 0,
     ):
         self._identity_file = identity_file
-        self._proxy_domain = self._get_proxy_domain()
+        self._jinja_env_cache: dict[Path, Environment] = {}
+        self._templates_by_error = get_settings().frontend.get("errors", {})
         super().__init__(
             max_connections=max_connections,
             max_keepalive_connections=max_keepalive_connections,
@@ -46,30 +54,90 @@ class FrontendProxyApp(ProxyAppBase):
         )
 
     def get_upstream_base_url(self, scope: Scope) -> str:
-        target = self._get_target_name(
+        target = get_target_name(
             {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})}
         )
+        if not target:
+            raise InvalidTargetError()
+
         return f"http://{target.lower()}"
 
-    def _get_proxy_domain(self):
-        settings = get_settings()
-        return (
-            settings.proxy.domain
-            if settings.proxy.domain[0] == "."
-            else f".{settings.proxy.domain}"
-        )
+    async def send_error_response(
+        self,
+        scope: Scope,
+        send: ASGISend,
+        http_status: int,
+        body: str,
+        headers: list[tuple[bytes, bytes]] | None = None,
+    ):
+        request_headers = {
+            k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})
+        }
+        accept_header = request_headers.get("accept")
+        if not (accept_header and str(http_status) in self._templates_by_error):
+            return await super().send_error_response(scope, send, http_status, body)
 
-    def _get_target_from_header(self, headers: dict[str, str], name: str) -> str | None:
-        header_value = headers.get(name, "")
-        if self._proxy_domain in header_value:
-            if ":" in header_value:
-                header_value, _ = header_value.split(":", 1)
-            return header_value[: -len(self._proxy_domain)]
+        available_templates = self._templates_by_error[str(http_status)]
 
-    def _get_target_name(self, headers: dict[str, str]) -> str:
-        target = self._get_target_from_header(headers, "x-forwarded-host")
-        if not target:
-            target = self._get_target_from_header(headers, "host")
-        if not target or not RE_SUBDOMAIN.fullmatch(target):
-            raise InvalidTargetError()
-        return target
+        media_types = parse_accept_header(accept_header)
+        for media_type in media_types:
+            template_format = ERROR_TEMPLATE_FORMATS.get(media_type)
+            if template_format and template_format in available_templates:
+                template_path = available_templates[template_format]
+                rendered = await self._render_error_template(
+                    Path(template_path), scope, http_status, body
+                )
+                return await super().send_error_response(
+                    scope,
+                    send,
+                    http_status,
+                    rendered,
+                    headers=[(b"content-type", media_type.encode("latin-1"))],
+                )
+
+        return await super().send_error_response(scope, send, http_status, body)
+
+    async def _render_error_template(
+        self, template_path: Path, scope: Scope, http_status: int, body: str
+    ) -> str:
+        env = self._get_jinja_env(template_path)
+        template = env.get_template(template_path.name)
+        status_title = HTTPStatus(http_status).name.replace("_", " ").title()
+        context = {
+            "status": http_status,
+            "status_title": status_title,
+            "body": body,
+            "request": self._extract_request_context(scope),
+        }
+
+        return await template.render_async(context)
+
+    def _extract_request_context(self, scope: Scope) -> dict[str, Any]:
+        headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])}
+
+        return {
+            "method": scope.get("method"),
+            "path": scope.get("path"),
+            "raw_path": scope.get("raw_path", b"").decode("latin-1"),
+            "query_string": scope.get("query_string", b"").decode("latin-1"),
+            "scheme": scope.get("scheme"),
+            "headers": headers,
+            "client": scope.get("client"),
+            "server": scope.get("server"),
+            "http_version": scope.get("http_version"),
+        }
+
+    def _get_jinja_env(self, template_path: Path) -> Environment:
+        template_dir = template_path.parent
+
+        if template_dir not in self._jinja_env_cache:
+            self._jinja_env_cache[template_dir] = Environment(
+                loader=FileSystemLoader(str(template_dir)),
+                autoescape=select_autoescape(
+                    enabled_extensions=("html", "xml"),
+                    default_for_string=False,
+                ),
+                enable_async=True,
+            )
+
+        return self._jinja_env_cache[template_dir]

mrok/frontend/main.py

@@ -7,6 +7,7 @@ from uvicorn_worker import UvicornWorker
 
 from mrok.conf import get_settings
 from mrok.frontend.app import FrontendProxyApp
+from mrok.frontend.middleware import HealthCheckMiddleware
 from mrok.logging import get_logging_config
 
 
@@ -42,11 +43,13 @@ def run(
     max_keepalive_connections: int | None,
     keepalive_expiry: float | None,
 ):
-    app = FrontendProxyApp(
-        str(identity_file),
-        max_connections=max_connections,
-        max_keepalive_connections=max_keepalive_connections,
-        keepalive_expiry=keepalive_expiry,
+    app = HealthCheckMiddleware(
+        FrontendProxyApp(
+            str(identity_file),
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+        )
     )
 
     options = {

mrok/frontend/middleware.py

@@ -0,0 +1,35 @@
+import json
+
+from mrok.frontend.utils import get_target_name
+from mrok.types.proxy import ASGIApp, ASGIReceive, ASGISend, Scope
+
+
+class HealthCheckMiddleware:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: ASGIReceive, send: ASGISend):
+        if scope["type"] == "http" and scope["path"] == "/healthcheck":
+            target = get_target_name(
+                {k.decode("latin1"): v.decode("latin1") for k, v in scope.get("headers", {})}
+            )
+
+            if not target:
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": 200,
+                        "headers": [
+                            [b"content-type", b"application/json"],
+                        ],
+                    }
+                )
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": json.dumps({"status": "healthy"}).encode("utf-8"),
+                    }
+                )
+                return
+
+        await self.app(scope, receive, send)

mrok/frontend/utils.py

@@ -0,0 +1,83 @@
+import re
+
+from mrok.conf import get_settings
+
+
+def parse_accept_header(accept: str | None) -> list[str]:
+    if not accept:
+        return ["*/*"]
+
+    result: list[tuple[str, float, int]] = []
+
+    for index, item in enumerate(accept.split(",")):
+        item = item.strip()
+        if not item:
+            continue
+
+        parts = [p.strip() for p in item.split(";")]
+        media_type = parts[0].lower()
+
+        q = 1.0
+        for param in parts[1:]:
+            if param.startswith("q="):
+                try:
+                    q = float(param[2:])
+                except ValueError:
+                    q = 0.0
+
+        result.append((media_type, q, index))
+
+    # Sort by:
+    # 1) q value (desc)
+    # 2) specificity (more specific first)
+    # 3) original order (stable)
+    result.sort(
+        key=lambda x: (
+            -x[1],
+            -_media_type_specificity(x[0]),
+            x[2],
+        )
+    )
+
+    return [media_type for media_type, _, _ in result]
+
+
+def _media_type_specificity(media_type: str) -> int:
+    if media_type == "*/*":
+        return 0
+    if media_type.endswith("/*"):
+        return 1
+    return 2
+
+
+def get_frontend_domain():
+    settings = get_settings()
+    return (
+        settings.frontend.domain
+        if settings.frontend.domain[0] == "."
+        else f".{settings.frontend.domain}"
+    )
+
+
+def _get_target_from_header(headers: dict[str, str], name: str) -> str | None:
+    domain_name = get_frontend_domain()
+    header_value = headers.get(name, "")
+    if domain_name in header_value:
+        if ":" in header_value:
+            header_value, _ = header_value.split(":", 1)
+        return header_value[: -len(domain_name)]
+
+
+def get_target_name(headers: dict[str, str]) -> str | None:
+    settings = get_settings()
+
+    target = _get_target_from_header(headers, "x-forwarded-host")
+    if not target:
+        target = _get_target_from_header(headers, "host")
+
+    if target and (
+        re.fullmatch(settings.identifiers.extension.regex, target)
+        or re.fullmatch(settings.identifiers.instance.regex, target)
+    ):
+        return target
+    return None

mrok/logging.py

@@ -1,12 +1,14 @@
+import logging
 import logging.config
 
-from rich.console import Console
-from rich.logging import RichHandler
-from textual_serve.server import LogHighlighter
-
 from mrok.conf import Settings
 
 
+class HealthCheckFilter(logging.Filter):
+    def filter(self, record):
+        return "/healthcheck" not in record.getMessage()
+
+
 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"
@@ -30,6 +32,11 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict:
             },
             "plain": {"format": "%(message)s"},
         },
+        "filters": {
+            "healthcheck_filter": {
+                "()": HealthCheckFilter,
+            }
+        },
         "handlers": {
             "console": {
                 "class": "logging.StreamHandler",
@@ -58,17 +65,30 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict:
                 "handlers": [handler],
                 "level": log_level,
                 "propagate": False,
+                "filters": ["healthcheck_filter"],
             },
             "gunicorn.error": {
                 "handlers": [handler],
                 "level": log_level,
                 "propagate": False,
             },
+            "uvicorn.access": {
+                "handlers": [handler],
+                "level": log_level,
+                "propagate": False,
+                "filters": ["healthcheck_filter"],
+            },
             "mrok": {
                 "handlers": [mrok_handler],
                 "level": log_level,
                 "propagate": False,
             },
+            "mrok.access": {
+                "handlers": [mrok_handler],
+                "level": log_level,
+                "propagate": False,
+                "filters": ["healthcheck_filter"],
+            },
         },
     }
 
@@ -78,21 +98,3 @@ def get_logging_config(settings: Settings, cli_mode: bool = False) -> dict:
 def setup_logging(settings: Settings, cli_mode: bool = False) -> None:
     logging_config = get_logging_config(settings, cli_mode)
     logging.config.dictConfig(logging_config)
-
-
-def setup_inspector_logging(console: Console) -> None:
-    logging.basicConfig(
-        level="WARNING",
-        format="%(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-        handlers=[
-            RichHandler(
-                show_path=False,
-                show_time=True,
-                rich_tracebacks=True,
-                tracebacks_show_locals=True,
-                highlighter=LogHighlighter(),
-                console=console,
-            )
-        ],
-    )

mrok/proxy/app.py

@@ -57,7 +57,7 @@ class ProxyAppBase(abc.ABC):
             return
 
         if scope.get("type") != "http":
-            await self._send_error(send, 500, "Unsupported")
+            await self.send_error_response(scope, send, 500, "Unsupported")
             return
 
         try:
@@ -105,15 +105,23 @@ class ProxyAppBase(abc.ABC):
             await response.aclose()
 
         except ProxyError as pe:
-            await self._send_error(send, pe.http_status, pe.message)
+            await self.send_error_response(scope, send, pe.http_status, pe.message)
 
         except Exception:
             logger.exception("Unexpected error in forwarder")
-            await self._send_error(send, 502, "Bad Gateway")
+            await self.send_error_response(scope, send, 502, "Bad Gateway")
 
-    async def _send_error(self, send: ASGISend, http_status: int, body: str):
+    async def send_error_response(
+        self,
+        scope: Scope,
+        send: ASGISend,
+        http_status: int,
+        body: str,
+        headers: list[tuple[bytes, bytes]] | None = None,
+    ):
+        headers = headers or [(b"content-type", b"text/plain")]
         try:
-            await send({"type": "http.response.start", "status": http_status, "headers": []})
+            await send({"type": "http.response.start", "status": http_status, "headers": headers})
             await send({"type": "http.response.body", "body": body.encode()})
         except Exception as e:  # pragma: no cover
             logger.error(f"Cannot send error response: {e}")

mrok/proxy/middleware.py

@@ -2,10 +2,8 @@ import asyncio
 import logging
 import time
 
-from mrok.proxy.constants import MAX_REQUEST_BODY_BYTES, MAX_RESPONSE_BODY_BYTES
 from mrok.proxy.metrics import MetricsCollector
 from mrok.proxy.models import FixedSizeByteBuffer, HTTPHeaders, HTTPRequest, HTTPResponse
-from mrok.proxy.utils import must_capture_request, must_capture_response
 from mrok.types.proxy import (
     ASGIApp,
     ASGIReceive,
@@ -15,6 +13,9 @@ from mrok.types.proxy import (
     Scope,
 )
 
+MAX_REQUEST_BODY_BYTES = 2 * 1024 * 1024
+MAX_RESPONSE_BODY_BYTES = 5 * 1024 * 1024
+
 logger = logging.getLogger("mrok.proxy")
 
 
@@ -42,7 +43,7 @@ class CaptureMiddleware:
         state = {}
 
         req_buf = FixedSizeByteBuffer(MAX_REQUEST_BODY_BYTES)
-        capture_req_body = must_capture_request(method, req_headers)
+        capture_req_body = method.upper() not in ("GET", "HEAD", "OPTIONS", "TRACE")
 
         request = HTTPRequest(
             method=method,
@@ -68,9 +69,7 @@ class CaptureMiddleware:
                 resp_headers = HTTPHeaders.from_asgi(msg.get("headers", []))
                 state["resp_headers_raw"] = resp_headers
 
-                state["capture_resp_body"] = must_capture_response(resp_headers)
-
-            if state["capture_resp_body"] and msg["type"] == "http.response.body":
+            if msg["type"] == "http.response.body":
                 body = msg.get("body", b"")
                 resp_buf.write(body)
 
@@ -91,8 +90,8 @@ class CaptureMiddleware:
             status=state["status"] or 0,
             headers=state["resp_headers_raw"],
             duration=duration,
-            body=resp_buf.getvalue() if state["capture_resp_body"] else None,
-            body_truncated=resp_buf.overflow if state["capture_resp_body"] else None,
+            body=resp_buf.getvalue(),
+            body_truncated=resp_buf.overflow,
         )
         asyncio.create_task(self._on_response_complete(response))
 

mrok/proxy/models.py

@@ -1,13 +1,40 @@
 from __future__ import annotations
 
+import base64
+import binascii
 import json
+from collections.abc import Mapping
 from pathlib import Path
-from typing import Literal
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic.functional_serializers import PlainSerializer
+from pydantic.functional_validators import PlainValidator
 from pydantic_core import core_schema
 
 
+def serialize_b64(v: bytes) -> str:
+    return base64.b64encode(v).decode("ascii")
+
+
+def deserialize_b64(v):
+    if isinstance(v, bytes):
+        return v
+    if isinstance(v, str):
+        try:
+            return base64.b64decode(v, validate=True)
+        except binascii.Error as e:  # pragma: no branch
+            raise ValueError("Invalid base64 data") from e
+    raise TypeError("Expected bytes or base64 string")  # pragma: no cover
+
+
+Base64Bytes = Annotated[
+    bytes,
+    PlainValidator(deserialize_b64),
+    PlainSerializer(serialize_b64, return_type=str),
+]
+
+
 class X509Credentials(BaseModel):
     key: str
     cert: str
@@ -16,14 +43,13 @@ class X509Credentials(BaseModel):
     @field_validator("key", "cert", "ca", mode="before")
     @classmethod
     def strip_pem_prefix(cls, value: str) -> str:
-        if isinstance(value, str) and value.startswith("pem:"):
+        if isinstance(value, str) and value.startswith("pem:"):  # pragma: no branch
             return value[4:]
-        return value
+        return value  # pragma: no cover
 
 
 class ServiceMetadata(BaseModel):
     model_config = ConfigDict(extra="ignore")
-    identity: str
     extension: str
     instance: str
     domain: str | None = None
@@ -34,10 +60,10 @@ class Identity(BaseModel):
     model_config = ConfigDict(extra="ignore")
     zt_api: str = Field(validation_alias="ztAPI")
     id: X509Credentials
+    mrok: ServiceMetadata
+    enable_ha: bool = Field(default=False, validation_alias="enableHa")
     zt_apis: str | None = Field(default=None, validation_alias="ztAPIs")
     config_types: str | None = Field(default=None, validation_alias="configTypes")
-    enable_ha: bool = Field(default=False, validation_alias="enableHa")
-    mrok: ServiceMetadata | None = None
 
     @staticmethod
     def load_from_file(path: str | Path) -> Identity:
@@ -77,7 +103,7 @@ class FixedSizeByteBuffer:
 
 
 class HTTPHeaders(dict):
-    def __init__(self, initial=None):
+    def __init__(self, initial: Mapping[Any, Any] | None = None):
         super().__init__()
         if initial:
             for k, v in initial.items():
@@ -131,9 +157,9 @@ class HTTPRequest(BaseModel):
     method: str
     url: str
     headers: HTTPHeaders
-    query_string: bytes
+    query_string: Base64Bytes | None = None
     start_time: float
-    body: bytes | None = None
+    body: Base64Bytes | None = None
     body_truncated: bool | None = None
 
 
@@ -143,7 +169,7 @@ class HTTPResponse(BaseModel):
     status: int
     headers: HTTPHeaders
     duration: float
-    body: bytes | None = None
+    body: Base64Bytes | None = None
     body_truncated: bool | None = None
 
 

mrok/proxy/ziticorn.py

@@ -1,4 +1,3 @@
-import json
 import logging
 import socket
 from collections.abc import Callable
@@ -10,6 +9,7 @@ from uvicorn import config, server
 from uvicorn.lifespan.on import LifespanOn
 from uvicorn.protocols.http.httptools_impl import HttpToolsProtocol as UvHttpToolsProtocol
 
+from mrok.proxy.models import Identity
 from mrok.types.proxy import ASGIApp
 
 logger = logging.getLogger("mrok.proxy")
@@ -48,10 +48,8 @@ class BackendConfig(config.Config):
         backlog: int = 2048,
     ):
         self.identity_file = identity_file
+        self.identity = Identity.load_from_file(self.identity_file)
         self.ziti_load_timeout_ms = ziti_load_timeout_ms
-        self.service_name, self.identity_name, self.instance_id = self.get_identity_info(
-            identity_file
-        )
         super().__init__(
             app,
             loop="asyncio",
@@ -59,26 +57,19 @@ class BackendConfig(config.Config):
             backlog=backlog,
         )
 
-    def get_identity_info(self, identity_file: str | Path):
-        with open(identity_file) as f:
-            identity_data = json.load(f)
-            try:
-                identity_name = identity_data["mrok"]["identity"]
-                instance_id, service_name = identity_name.split(".", 1)
-                return service_name, identity_name, instance_id
-            except KeyError:
-                raise ValueError("Invalid identity file: identity file is not mrok compatible.")
-
     def bind_socket(self) -> socket.socket:
-        logger.info(f"Connect to Ziti service '{self.service_name} ({self.instance_id})'")
+        logger.info(
+            "Connect to Ziti service "
+            f"'{self.identity.mrok.extension} ({self.identity.mrok.instance})'"
+        )
 
         ctx, err = openziti.load(str(self.identity_file), timeout=self.ziti_load_timeout_ms)
         if err != 0:
             raise RuntimeError(f"Failed to load Ziti identity from {self.identity_file}: {err}")
 
-        sock = ctx.bind(self.service_name)
+        sock = ctx.bind(self.identity.mrok.extension)
         sock.listen(self.backlog)
-        logger.info(f"listening on ziti service {self.service_name} for connections")
+        logger.info(f"listening on ziti service {self.identity.mrok.extension} for connections")
         return sock
 
     def configure_logging(self) -> None: