mrok 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

mrok/agent/devtools/inspector/utils.py

@@ -0,0 +1,149 @@
+from collections.abc import Generator
+from dataclasses import dataclass
+from io import BytesIO
+
+from multipart import MultipartParser
+
+TEXTUAL_CONTENT_TYPES = {
+    "application/json",
+    "application/xml",
+    "application/javascript",
+    "application/x-www-form-urlencoded",
+}
+
+TEXTUAL_PREFIXES = ("text/",)
+
+CONTENT_TYPE_TO_LANGUAGE = {
+    "application/json": "json",
+    "application/ld+json": "json",
+    "application/problem+json": "json",
+    "application/schema+json": "json",
+    "application/xml": "xml",
+    "text/xml": "xml",
+    "application/xhtml+xml": "html",
+    "text/html": "html",
+    "text/css": "css",
+    "application/javascript": "javascript",
+    "application/x-javascript": "javascript",
+    "text/javascript": "javascript",
+    "application/ecmascript": "javascript",
+    "text/markdown": "markdown",
+    "text/x-markdown": "markdown",
+    "application/yaml": "yaml",
+    "application/x-yaml": "yaml",
+    "text/yaml": "yaml",
+    "application/toml": "toml",
+    "application/x-toml": "toml",
+    "application/sql": "sql",
+    "text/x-sql": "sql",
+    "application/java": "java",
+    "text/x-java-source": "java",
+    "application/python": "python",
+    "text/x-python": "python",
+    "application/x-python-code": "python",
+    "application/rust": "rust",
+    "text/x-rust": "rust",
+    "application/go": "go",
+    "text/x-go": "go",
+    "application/bash": "bash",
+    "application/x-sh": "bash",
+    "text/x-shellscript": "bash",
+    "application/regex": "regex",
+    "text/x-regex": "regex",
+}
+
+
+@dataclass
+class ContentTypeInfo:
+    content_type: str
+    binary: bool
+    charset: str | None = None
+    boundary: str | None = None
+
+
+def parse_content_type(content_type_header: str) -> ContentTypeInfo:
+    parts = content_type_header.split(";")
+    content_type = parts[0].strip().lower()
+
+    charset = None
+    boundary = None
+
+    for part in parts[1:]:
+        part = part.strip()
+        if "=" in part:
+            key, value = part.split("=", 1)
+            key = key.strip().lower()
+            value = value.strip().strip('"')
+            if key == "charset":
+                charset = value
+            elif key == "boundary":
+                boundary = value
+
+    binary = not is_textual(content_type)
+
+    if charset is None and not binary:
+        charset = "utf-8"
+
+    return ContentTypeInfo(
+        content_type=content_type, binary=binary, charset=charset, boundary=boundary
+    )
+
+
+def parse_form_data(data: bytes, boundary: str) -> Generator[tuple[str, str]]:
+    parser = MultipartParser(BytesIO(data), boundary)
+    for part in parser:
+        if is_textual(part.content_type):
+            yield part.name, part.value
+            continue
+        yield part.name, "<binary>"
+
+
+def is_textual(content_type: str) -> bool:
+    ct = content_type.lower()
+    if ct in TEXTUAL_CONTENT_TYPES:
+        return True
+    if any(ct.startswith(p) for p in TEXTUAL_PREFIXES):
+        return True
+    return False
+
+
+def build_tree(node, data):
+    if isinstance(data, dict):
+        for key, value in data.items():
+            child = node.add(str(key))
+            build_tree(child, value)
+    elif isinstance(data, list):
+        for index, value in enumerate(data):
+            child = node.add(f"[{index}]")
+            build_tree(child, value)
+    else:
+        node.add(repr(data))
+
+
+def hexdump(data, width=16):
+    lines = []
+    for i in range(0, len(data), width):
+        chunk = data[i : i + width]
+        hex_part = " ".join(f"{b:02x}" for b in chunk)
+        ascii_part = "".join(chr(b) if 32 <= b <= 126 else "." for b in chunk)
+        lines.append(f"{hex_part:<{width * 3}} {ascii_part}")
+    return "\n".join(lines)
+
+
+def humanize_bytes(num_bytes: int) -> tuple[float, str]:  # type: ignore[return-value]
+    if num_bytes < 0:
+        raise ValueError("num_bytes must be non-negative")
+
+    units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
+    value = float(num_bytes)
+
+    for unit in units:
+        if value < 1024 or unit == units[-1]:
+            return round(value, 2), unit
+        value /= 1024
+
+
+def get_highlighter_language_by_content_type(content_type: str) -> str | None:
+    if content_type in CONTENT_TYPE_TO_LANGUAGE:
+        return CONTENT_TYPE_TO_LANGUAGE[content_type]
+    return None

mrok/agent/sidecar/app.py

@@ -5,7 +5,7 @@ from typing import Literal
 from httpcore import AsyncConnectionPool
 
 from mrok.proxy.app import ProxyAppBase
-from mrok.proxy.types import Scope
+from mrok.types.proxy import Scope
 
 logger = logging.getLogger("mrok.agent")
 
@@ -18,10 +18,10 @@ class SidecarProxyApp(ProxyAppBase):
         self,
         target: str | Path | tuple[str, int],
         *,
-        max_connections=1000,
-        max_keepalive_connections=10,
-        keepalive_expiry=120,
-        retries=0,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
     ):
         self._target = target
         self._target_type, self._target_address = self._parse_target()
@@ -34,10 +34,10 @@ class SidecarProxyApp(ProxyAppBase):
 
     def setup_connection_pool(
         self,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 10,
-        keepalive_expiry: float | None = 120.0,
-        retries: int = 0,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
     ) -> AsyncConnectionPool:
         if self._target_type == "unix":
             return AsyncConnectionPool(

mrok/agent/sidecar/main.py

@@ -13,26 +13,47 @@ class SidecarAgent(MasterBase):
         identity_file: str,
         target: str | Path | tuple[str, int],
         workers: int = 4,
+        events_enabled: bool = True,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
+        retries: int = 0,
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
     ):
         super().__init__(
             identity_file,
-            workers,
-            False,
-            publishers_port,
-            subscribers_port,
+            workers=workers,
+            reload=False,
+            events_enabled=events_enabled,
+            events_pub_port=publishers_port,
+            events_sub_port=subscribers_port,
         )
         self._target = target
+        self._max_connections = max_connections
+        self._max_keepalive_connections = max_keepalive_connections
+        self._keepalive_expiry = keepalive_expiry
+        self._retries = retries
 
     def get_asgi_app(self):
-        return SidecarProxyApp(self._target)
+        return SidecarProxyApp(
+            self._target,
+            max_connections=self._max_connections,
+            max_keepalive_connections=self._max_keepalive_connections,
+            keepalive_expiry=self._keepalive_expiry,
+            retries=self._retries,
+        )
 
 
 def run(
     identity_file: str,
     target_addr: str | Path | tuple[str, int],
     workers: int = 4,
+    events_enabled: bool = True,
+    max_connections: int | None = 10,
+    max_keepalive_connections: int | None = None,
+    keepalive_expiry: float | None = None,
+    retries: int = 0,
     publishers_port: int = 50000,
     subscribers_port: int = 50001,
 ):
@@ -40,6 +61,11 @@ def run(
         identity_file,
         target_addr,
         workers=workers,
+        events_enabled=events_enabled,
+        max_connections=max_connections,
+        max_keepalive_connections=max_keepalive_connections,
+        keepalive_expiry=keepalive_expiry,
+        retries=retries,
         publishers_port=publishers_port,
         subscribers_port=subscribers_port,
     )

mrok/agent/ziticorn.py

@@ -1,5 +1,5 @@
 from mrok.proxy.master import MasterBase
-from mrok.proxy.types import ASGIApp
+from mrok.types.proxy import ASGIApp
 
 
 class ZiticornAgent(MasterBase):
@@ -12,7 +12,13 @@ class ZiticornAgent(MasterBase):
         publishers_port: int = 50000,
         subscribers_port: int = 50001,
     ):
-        super().__init__(identity_file, workers, reload, publishers_port, subscribers_port)
+        super().__init__(
+            identity_file,
+            workers=workers,
+            reload=reload,
+            events_pub_port=publishers_port,
+            events_sub_port=subscribers_port,
+        )
         self.app = app
 
     def get_asgi_app(self):

mrok/cli/commands/admin/bootstrap.py

@@ -8,14 +8,15 @@ import typer
 
 from mrok.cli.commands.admin.utils import parse_tags
 from mrok.conf import Settings
-from mrok.ziti.api import TagsType, ZitiClientAPI, ZitiManagementAPI
+from mrok.types.ziti import Tags
+from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
 from mrok.ziti.bootstrap import bootstrap_identity
 
 logger = logging.getLogger(__name__)
 
 
 async def bootstrap(
-    settings: Settings, forced: bool, tags: TagsType | None
+    settings: Settings, forced: bool, tags: Tags | None
 ) -> tuple[str, dict[str, Any] | None]:
     async with ZitiManagementAPI(settings) as mgmt_api, ZitiClientAPI(settings) as client_api:
         return await bootstrap_identity(

mrok/cli/commands/admin/utils.py

@@ -2,10 +2,10 @@ from datetime import datetime
 
 import typer
 
-from mrok.ziti.api import TagsType
+from mrok.types.ziti import Tags
 
 
-def parse_tags(pairs: list[str] | None) -> TagsType | None:
+def parse_tags(pairs: list[str] | None) -> Tags | None:
     if not pairs:
         return None
 

mrok/cli/commands/agent/run/sidecar.py

@@ -25,10 +25,55 @@ def register(app: typer.Typer) -> None:
             typer.Option(
                 "--workers",
                 "-w",
-                help=f"Number of workers. Default: {default_workers}",
+                help="Number of workers.",
                 show_default=True,
             ),
         ] = default_workers,
+        max_connections: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connections",
+                help=(
+                    "The maximum number of concurrent HTTP connections that "
+                    "the pool should allow. Any attempt to send a request on a pool that "
+                    "would exceed this amount will block until a connection is available."
+                ),
+                show_default=True,
+            ),
+        ] = 10,
+        max_keepalive_connections: Annotated[
+            int | None,
+            typer.Option(
+                "--max-pool-keepalive-connections",
+                help=(
+                    "The maximum number of idle HTTP connections "
+                    "that will be maintained in the pool."
+                ),
+                show_default=True,
+            ),
+        ] = None,
+        keepalive_expiry: Annotated[
+            float | None,
+            typer.Option(
+                "--max-pool-keepalive-expiry",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = None,
+        retries: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connect-retries",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 0,
         publishers_port: Annotated[
             int,
             typer.Option(
@@ -53,6 +98,14 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = 50001,
+        no_events: Annotated[
+            bool,
+            typer.Option(
+                "--no-events",
+                help="Disable events. Default: False",
+                show_default=True,
+            ),
+        ] = False,
     ):
         """Run a Sidecar Proxy to expose a web application through OpenZiti."""
         if ":" in str(target):
@@ -65,6 +118,11 @@ def register(app: typer.Typer) -> None:
             str(identity_file),
             target_addr,
             workers=workers,
+            events_enabled=not no_events,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+            retries=retries,
             publishers_port=publishers_port,
             subscribers_port=subscribers_port,
         )

mrok/cli/commands/frontend/run.py

@@ -44,6 +44,48 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = default_workers,
+        max_connections: Annotated[
+            int,
+            typer.Option(
+                "--max-pool-connections",
+                help=(
+                    "The maximum number of concurrent HTTP connections that "
+                    "the pool should allow. Any attempt to send a request on a pool that "
+                    "would exceed this amount will block until a connection is available."
+                ),
+                show_default=True,
+            ),
+        ] = 1000,
+        max_keepalive_connections: Annotated[
+            int | None,
+            typer.Option(
+                "--max-pool-keepalive-connections",
+                help=(
+                    "The maximum number of idle HTTP connections "
+                    "that will be maintained in the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 100,
+        keepalive_expiry: Annotated[
+            float | None,
+            typer.Option(
+                "--max-pool-keepalive-expiry",
+                help=(
+                    "The duration in seconds that an idle HTTP connection "
+                    "may be maintained for before being expired from the pool."
+                ),
+                show_default=True,
+            ),
+        ] = 300,
     ):
         """Run the mrok frontend with Gunicorn and Uvicorn workers."""
-        frontend.run(identity_file, host, port, workers)
+        frontend.run(
+            identity_file,
+            host,
+            port,
+            workers,
+            max_connections=max_connections,
+            max_keepalive_connections=max_keepalive_connections,
+            keepalive_expiry=keepalive_expiry,
+        )

mrok/cli/main.py

@@ -1,5 +1,6 @@
 import inspect
 import sys
+from typing import Annotated
 
 import typer
 from pyfiglet import Figlet
@@ -79,11 +80,23 @@ for name, module in inspect.getmembers(commands):
     elif hasattr(module, "app"):  # pragma: no branch
         app.add_typer(module.app, name=name.replace("_", "-"))
 
+_debug_mode = False
+
 
 @app.callback()
 def main(
     ctx: typer.Context,
+    debug: Annotated[
+        bool,
+        typer.Option(
+            "--debug",
+            help="Run the CLI in debug mode",
+            show_default=True,
+        ),
+    ] = False,
 ):
+    global _debug_mode
+    _debug_mode = debug
     settings = get_settings()
     setup_logging(settings, cli_mode=True)
     ctx.obj = settings
@@ -93,5 +106,8 @@ def run():
     try:
         app()
     except Exception as e:
-        err_console.print(f"[bold red]Error:[/bold red] {e}")
+        if _debug_mode:
+            raise
+        message = str(e) or "Unexpected error. Debug it with --debug"
+        err_console.print(f"[bold red]Error:[/bold red] {message}")
         sys.exit(-1)

mrok/constants.py

@@ -2,3 +2,24 @@ import re
 
 RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
 RE_INSTANCE_ID = re.compile(r"(?i)INS-\d{4}-\d{4}-\d{4}")
+
+
+BINARY_CONTENT_TYPES = {
+    "application/octet-stream",
+    "application/pdf",
+}
+
+BINARY_PREFIXES = (
+    "image/",
+    "video/",
+    "audio/",
+)
+
+TEXTUAL_CONTENT_TYPES = {
+    "application/json",
+    "application/xml",
+    "application/javascript",
+    "application/x-www-form-urlencoded",
+}
+
+TEXTUAL_PREFIXES = ("text/",)

mrok/controller/schemas.py

@@ -9,12 +9,12 @@ from pydantic import (
     computed_field,
 )
 
-from mrok.ziti.api import TagsType
+from mrok.types.ziti import Tags
 
 
 class BaseSchema(BaseModel):
     model_config = ConfigDict(from_attributes=True, extra="ignore")
-    tags: TagsType | None = None
+    tags: Tags | None = None
 
 
 class IdSchema(BaseModel):

mrok/frontend/app.py

@@ -6,7 +6,7 @@ from mrok.conf import get_settings
 from mrok.proxy.app import ProxyAppBase
 from mrok.proxy.backend import AIOZitiNetworkBackend
 from mrok.proxy.exceptions import InvalidTargetError
-from mrok.proxy.types import Scope
+from mrok.types.proxy import Scope
 
 RE_SUBDOMAIN = re.compile(r"(?i)^(?:EXT-\d{4}-\d{4}|INS-\d{4}-\d{4}-\d{4})$")
 
@@ -16,9 +16,9 @@ class FrontendProxyApp(ProxyAppBase):
         self,
         identity_file: str,
         *,
-        max_connections: int = 1000,
-        max_keepalive_connections: int = 10,
-        keepalive_expiry: float = 120.0,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
         retries=0,
     ):
         self._identity_file = identity_file
@@ -32,10 +32,10 @@ class FrontendProxyApp(ProxyAppBase):
 
     def setup_connection_pool(
         self,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 100,
-        keepalive_expiry: float | None = 120.0,
-        retries: int = 0,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
     ) -> AsyncConnectionPool:
         return AsyncConnectionPool(
             max_connections=max_connections,

mrok/frontend/main.py

@@ -38,8 +38,16 @@ def run(
     host: str,
     port: int,
     workers: int,
+    max_connections: int | None,
+    max_keepalive_connections: int | None,
+    keepalive_expiry: float | None,
 ):
-    app = FrontendProxyApp(str(identity_file))
+    app = FrontendProxyApp(
+        str(identity_file),
+        max_connections=max_connections,
+        max_keepalive_connections=max_keepalive_connections,
+        keepalive_expiry=keepalive_expiry,
+    )
 
     options = {
         "bind": f"{host}:{port}",

mrok/logging.py

@@ -1,9 +1,5 @@
 import logging.config
 
-from rich.console import Console
-from rich.logging import RichHandler
-from textual_serve.server import LogHighlighter
-
 from mrok.conf import Settings
 
 
@@ -78,21 +74,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

@@ -4,8 +4,8 @@ import logging
 from httpcore import AsyncConnectionPool, Request
 
 from mrok.proxy.exceptions import ProxyError
-from mrok.proxy.streams import ASGIRequestBodyStream
-from mrok.proxy.types import ASGIReceive, ASGISend, Scope
+from mrok.proxy.stream import ASGIRequestBodyStream
+from mrok.types.proxy import ASGIReceive, ASGISend, Scope
 
 logger = logging.getLogger("mrok.proxy")
 
@@ -26,9 +26,9 @@ class ProxyAppBase(abc.ABC):
     def __init__(
         self,
         *,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 10,
-        keepalive_expiry: float | None = 120.0,
+        max_connections: int | None = 10,
+        max_keepalive_connections: int | None = None,
+        keepalive_expiry: float | None = None,
         retries: int = 0,
     ) -> None:
         self._pool = self.setup_connection_pool(
@@ -41,10 +41,10 @@ class ProxyAppBase(abc.ABC):
     @abc.abstractmethod
     def setup_connection_pool(
         self,
-        max_connections: int | None = 1000,
-        max_keepalive_connections: int | None = 10,
-        keepalive_expiry: float | None = 120.0,
-        retries: int = 0,
+        max_connections: int | None,
+        max_keepalive_connections: int | None,
+        keepalive_expiry: float | None,
+        retries: int,
     ) -> AsyncConnectionPool:
         raise NotImplementedError()
 
@@ -78,6 +78,7 @@ class ProxyAppBase(abc.ABC):
                 content=body_stream,
             )
             response = await self._pool.handle_async_request(request)
+            logger.debug(f"connection pool status: {self._pool}")
             response_headers = []
             for k, v in response.headers:
                 if k.lower() not in HOP_BY_HOP_HEADERS: