mrok 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

mrok/agent/sidecar/main.py

@@ -1,27 +1,46 @@
-import asyncio
-import contextlib
-from functools import partial
+import logging
 from pathlib import Path
 
 from mrok.agent.sidecar.app import ForwardApp
-from mrok.http.config import MrokBackendConfig
-from mrok.http.master import Master
-from mrok.http.server import MrokServer
+from mrok.master import MasterBase
 
+logger = logging.getLogger("mrok.proxy")
 
-def run_sidecar(identity_file: str, target_addr: str | Path | tuple[str, int]):
-    config = MrokBackendConfig(ForwardApp(target_addr), identity_file)
-    server = MrokServer(config)
-    with contextlib.suppress(KeyboardInterrupt, asyncio.CancelledError):
-        server.run()
+
+class SidecarAgent(MasterBase):
+    def __init__(
+        self,
+        identity_file: str,
+        target_addr: str | Path | tuple[str, int],
+        workers: int = 4,
+        publishers_port: int = 50000,
+        subscribers_port: int = 50001,
+    ):
+        super().__init__(
+            identity_file,
+            workers,
+            False,
+            publishers_port,
+            subscribers_port,
+        )
+        self.target_address = target_addr
+
+    def get_asgi_app(self):
+        return ForwardApp(self.target_address)
 
 
 def run(
     identity_file: str,
     target_addr: str | Path | tuple[str, int],
-    workers=4,
-    reload=False,
+    workers: int = 4,
+    publishers_port: int = 50000,
+    subscribers_port: int = 50001,
 ):
-    start_fn = partial(run_sidecar, identity_file, target_addr)
-    master = Master(start_fn, workers=workers, reload=reload)
-    master.run()
+    agent = SidecarAgent(
+        identity_file,
+        target_addr,
+        workers=workers,
+        publishers_port=publishers_port,
+        subscribers_port=subscribers_port,
+    )
+    agent.run()

mrok/agent/ziticorn.py

@@ -1,29 +1,38 @@
-import asyncio
-import contextlib
-import os
-from functools import partial
+from mrok.http.types import ASGIApp
+from mrok.master import MasterBase
 
-from mrok.http.config import ASGIApplication, MrokBackendConfig
-from mrok.http.master import Master
-from mrok.http.server import MrokServer
 
+class ZiticornAgent(MasterBase):
+    def __init__(
+        self,
+        app: ASGIApp | str,
+        identity_file: str,
+        workers: int = 4,
+        reload: bool = False,
+        publishers_port: int = 50000,
+        subscribers_port: int = 50001,
+    ):
+        super().__init__(identity_file, workers, reload, publishers_port, subscribers_port)
+        self.app = app
 
-def run_ziticorn(app: ASGIApplication | str, identity_file: str):
-    import sys
-
-    sys.path.insert(0, os.getcwd())
-    config = MrokBackendConfig(app, identity_file)
-    server = MrokServer(config)
-    with contextlib.suppress(KeyboardInterrupt, asyncio.CancelledError):
-        server.run()
+    def get_asgi_app(self):
+        return self.app
 
 
 def run(
-    app: ASGIApplication | str,
+    app: ASGIApp | str,
     identity_file: str,
     workers: int = 4,
     reload: bool = False,
+    publishers_port: int = 50000,
+    subscribers_port: int = 50001,
 ):
-    start_fn = partial(run_ziticorn, app, identity_file)
-    master = Master(start_fn, workers=workers, reload=reload)
+    master = ZiticornAgent(
+        app,
+        identity_file,
+        workers=workers,
+        reload=reload,
+        publishers_port=publishers_port,
+        subscribers_port=subscribers_port,
+    )
     master.run()

mrok/cli/commands/__init__.py

@@ -1,7 +1,8 @@
-from mrok.cli.commands import admin, agent, controller
+from mrok.cli.commands import admin, agent, controller, proxy
 
 __all__ = [
     "admin",
     "agent",
     "controller",
+    "proxy",
 ]

mrok/cli/commands/agent/__init__.py

@@ -1,6 +1,8 @@
 import typer
 
+from mrok.cli.commands.agent.dev import app as dev_app
 from mrok.cli.commands.agent.run import app as run_app
 
 app = typer.Typer(help="mrok agent commands.")
 app.add_typer(run_app)
+app.add_typer(dev_app)

mrok/cli/commands/agent/dev/__init__.py

@@ -0,0 +1,7 @@
+import typer
+
+from mrok.cli.commands.agent.dev import console, web
+
+app = typer.Typer(name="dev", help="Dev tools for mrok agent.")
+console.register(app)
+web.register(app)

mrok/cli/commands/agent/dev/console.py

@@ -0,0 +1,25 @@
+from typing import Annotated
+
+import typer
+
+from mrok.agent.devtools.inspector.app import InspectorApp
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("console")
+    def run_dev_console(
+        subscriber_port: Annotated[
+            int,
+            typer.Option(
+                "--subscriber-port",
+                "-s",
+                help=(
+                    "TCP port where the mrok inspector console application "
+                    "should connect to subscribe to request/response messages."
+                ),
+                show_default=True,
+            ),
+        ] = 50001,
+    ):
+        app = InspectorApp(subscriber_port)
+        app.run()

mrok/cli/commands/agent/dev/web.py

@@ -0,0 +1,37 @@
+from typing import Annotated
+
+import typer
+
+from mrok.agent.devtools.inspector.server import InspectorServer
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("web")
+    def run_web_console(
+        port: Annotated[
+            int,
+            typer.Option(
+                "--port",
+                "-p",
+                help="TCP port where the mrok inspector web application will listen for requests.",
+                show_default=True,
+            ),
+        ] = 7777,
+        subscriber_port: Annotated[
+            int,
+            typer.Option(
+                "--subscriber-port",
+                "-s",
+                help=(
+                    "TCP port where the mrok inspector web application "
+                    "should connect to subscribe to request/response messages."
+                ),
+                show_default=True,
+            ),
+        ] = 50001,
+    ):
+        server = InspectorServer(
+            port=port,
+            subscriber_port=subscriber_port,
+        )
+        server.serve()

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

@@ -1,16 +1,10 @@
-import multiprocessing
 from pathlib import Path
 from typing import Annotated
 
 import typer
 
-# from app.logging import get_logging_config
 from mrok.agent import ziticorn
-
-
-def number_of_workers():
-    return (multiprocessing.cpu_count() * 2) + 1
-
+from mrok.cli.utils import number_of_workers
 
 default_workers = number_of_workers()
 
@@ -18,14 +12,8 @@ default_workers = number_of_workers()
 def register(app: typer.Typer) -> None:
     @app.command("asgi")
     def run_asgi(
-        app: str = typer.Argument(
-            ...,
-            help="ASGI application",
-        ),
-        identity_file: Path = typer.Argument(
-            ...,
-            help="Identity json file",
-        ),
+        app: Annotated[str, typer.Argument(..., help="ASGI application")],
+        identity_file: Annotated[Path, typer.Argument(..., help="Identity json file")],
         workers: Annotated[
             int,
             typer.Option(
@@ -44,6 +32,37 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = False,
+        publishers_port: Annotated[
+            int,
+            typer.Option(
+                "--publishers-port",
+                "-p",
+                help=(
+                    "TCP port where the mrok agent "
+                    "should connect to publish to request/response messages."
+                ),
+                show_default=True,
+            ),
+        ] = 50000,
+        subscribers_port: Annotated[
+            int,
+            typer.Option(
+                "--subscribers-port",
+                "-s",
+                help=(
+                    "TCP port where the mrok agent should listen for incoming subscribers "
+                    "connections for request/response messages."
+                ),
+                show_default=True,
+            ),
+        ] = 50001,
     ):
         """Run an ASGI application exposing it through OpenZiti network."""
-        ziticorn.run(app, str(identity_file), workers=workers, reload=reload)
+        ziticorn.run(
+            app,
+            str(identity_file),
+            workers=workers,
+            reload=reload,
+            publishers_port=publishers_port,
+            subscribers_port=subscribers_port,
+        )

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

@@ -1,15 +1,10 @@
-import multiprocessing
 from pathlib import Path
 from typing import Annotated
 
 import typer
 
 from mrok.agent import sidecar
-
-
-def number_of_workers():
-    return (multiprocessing.cpu_count() * 2) + 1
-
+from mrok.cli.utils import number_of_workers
 
 default_workers = number_of_workers()
 
@@ -34,15 +29,30 @@ def register(app: typer.Typer) -> None:
                 show_default=True,
             ),
         ] = default_workers,
-        reload: Annotated[
-            bool,
+        publishers_port: Annotated[
+            int,
+            typer.Option(
+                "--publishers-port",
+                "-p",
+                help=(
+                    "TCP port where the mrok agent "
+                    "should connect to publish to request/response messages."
+                ),
+                show_default=True,
+            ),
+        ] = 50000,
+        subscribers_port: Annotated[
+            int,
             typer.Option(
-                "--reload",
-                "-r",
-                help="Enable auto-reload. Default: False",
+                "--subscribers-port",
+                "-s",
+                help=(
+                    "TCP port where the mrok agent should listen for incoming subscribers "
+                    "connections for request/response messages."
+                ),
                 show_default=True,
             ),
-        ] = False,
+        ] = 50001,
     ):
         """Run a Sidecar Proxy to expose a web application through OpenZiti."""
         if ":" in str(target):
@@ -51,4 +61,10 @@ def register(app: typer.Typer) -> None:
         else:
             target_addr = str(target)  # type: ignore
 
-        sidecar.run(str(identity_file), target_addr, workers=workers, reload=reload)
+        sidecar.run(
+            str(identity_file),
+            target_addr,
+            workers=workers,
+            publishers_port=publishers_port,
+            subscribers_port=subscribers_port,
+        )

mrok/cli/commands/agent/utils.py

@@ -0,0 +1,5 @@
+import multiprocessing
+
+
+def number_of_workers() -> int:
+    return (multiprocessing.cpu_count() * 2) + 1

mrok/cli/commands/controller/run.py

@@ -1,18 +1,14 @@
-import multiprocessing
 from collections.abc import Callable
 from typing import Annotated, Any
 
 import typer
 from gunicorn.app.base import BaseApplication
 
+from mrok.cli.utils import number_of_workers
 from mrok.controller.app import app as asgi_app
 from mrok.logging import get_logging_config
 
 
-def number_of_workers():
-    return (multiprocessing.cpu_count() * 2) + 1
-
-
 class StandaloneApplication(BaseApplication):  # pragma: no cover
     def __init__(self, application: Callable, options: dict[str, Any] | None = None):
         self.options = options or {}

mrok/cli/commands/proxy/__init__.py

@@ -0,0 +1,6 @@
+import typer
+
+from mrok.cli.commands.proxy import run
+
+app = typer.Typer(help="mrok proxy commands.")
+run.register(app)

mrok/cli/commands/proxy/run.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from mrok import proxy
+from mrok.cli.utils import number_of_workers
+
+default_workers = number_of_workers()
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("run")
+    def run_proxy(
+        ctx: typer.Context,
+        identity_file: Path = typer.Argument(
+            ...,
+            help="Identity json file",
+        ),
+        host: Annotated[
+            str,
+            typer.Option(
+                "--host",
+                "-h",
+                help="Host to bind to. Default: 127.0.0.1",
+                show_default=True,
+            ),
+        ] = "127.0.0.1",
+        port: Annotated[
+            int,
+            typer.Option(
+                "--port",
+                "-P",
+                help="Port to bind to. Default: 8000",
+                show_default=True,
+            ),
+        ] = 8000,
+        workers: Annotated[
+            int,
+            typer.Option(
+                "--workers",
+                "-w",
+                help=f"Number of workers. Default: {default_workers}",
+                show_default=True,
+            ),
+        ] = default_workers,
+    ):
+        """Run the mrok proxy with Gunicorn and Uvicorn workers."""
+        proxy.run(identity_file, host, port, workers)

mrok/cli/utils.py

@@ -0,0 +1,5 @@
+import multiprocessing
+
+
+def number_of_workers() -> int:
+    return (multiprocessing.cpu_count() * 2) + 1

mrok/conf.py

@@ -15,6 +15,12 @@ DEFAULT_SETTINGS = {
         "ssl_verify": False,
     },
     "PAGINATION": {"limit": 50},
+    "SIDECAR": {
+        "textual_port": 4040,
+        "store_port": 5051,
+        "store_size": 1000,
+        "textual_command": "python mrok/agent/sidecar/inspector.py",
+    },
 }
 
 _settings = None

mrok/controller/auth.py

@@ -80,8 +80,8 @@ async def authenticate(
         )
         return payload
     except jwt.InvalidKeyError as e:
-        logger.error(f"Invalid jwt token: {e}")
+        logger.error(f"Invalid jwt token: {e} ({credentials.credentials})")
         raise UNAUTHORIZED_EXCEPTION
     except jwt.InvalidTokenError as e:
-        logger.error(f"Invalid jwt token: {e}")
+        logger.error(f"Invalid jwt token: {e} ({credentials.credentials})")
         raise UNAUTHORIZED_EXCEPTION

mrok/datastructures.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+from pydantic_core import core_schema
+
+
+class FixedSizeByteBuffer:
+    def __init__(self, max_size: int):
+        self._max_size = max_size
+        self._buf = bytearray()
+        self.overflow = False
+
+    def write(self, data: bytes) -> None:
+        if not data:
+            return
+
+        remaining = self._max_size - len(self._buf)
+        if remaining <= 0:
+            self.overflow = True
+            return
+
+        if len(data) > remaining:
+            self._buf.extend(data[:remaining])
+            self.overflow = True
+        else:
+            self._buf.extend(data)
+
+    def getvalue(self) -> bytes:
+        return bytes(self._buf)
+
+    def clear(self) -> None:
+        self._buf.clear()
+        self.overflow = False
+
+
+class HTTPHeaders(dict):
+    def __init__(self, initial=None):
+        super().__init__()
+        if initial:
+            for k, v in initial.items():
+                super().__setitem__(str(k).lower(), str(v))
+
+    def __getitem__(self, key: str) -> str:
+        return super().__getitem__(key.lower())
+
+    def __setitem__(self, key: str, value: str) -> None:
+        super().__setitem__(str(key).lower(), str(value))
+
+    def __delitem__(self, key: str) -> None:
+        super().__delitem__(key.lower())
+
+    def get(self, k, default=None):
+        return super().get(k.lower(), default)
+
+    @classmethod
+    def from_asgi(cls, items: list[tuple[bytes, bytes]]) -> HTTPHeaders:
+        d = {k.decode("latin-1"): v.decode("latin-1") for k, v in items}
+        return cls(d)
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source, handler):
+        """Provide a pydantic-core schema so Pydantic treats this as a mapping of str->str.
+
+        We generate the schema for `dict[str, str]` using the provided handler and wrap
+        it with a validator that converts the validated dict into `HTTPHeaders`.
+        """
+        # handler may be a callable or an object with `generate_schema`; handle both
+        try:
+            dict_schema = handler.generate_schema(dict[str, str])
+        except AttributeError:
+            dict_schema = handler(dict[str, str])
+
+        def _wrap(v, validator):
+            # `validator` will validate input according to `dict_schema` and return a dict
+            validated = validator(input_value=v)
+            if isinstance(validated, HTTPHeaders):
+                return validated
+            return cls(validated)
+
+        return core_schema.no_info_wrap_validator_function(
+            _wrap,
+            dict_schema,
+            serialization=core_schema.plain_serializer_function_ser_schema(lambda v: dict(v)),
+        )
+
+
+class HTTPRequest(BaseModel):
+    method: str
+    url: str
+    headers: HTTPHeaders
+    query_string: bytes
+    start_time: float
+    body: bytes | None = None
+    body_truncated: bool | None = None
+
+
+class HTTPResponse(BaseModel):
+    type: Literal["response"] = "response"
+    request: HTTPRequest
+    status: int
+    headers: HTTPHeaders
+    duration: float
+    body: bytes | None = None
+    body_truncated: bool | None = None
+
+
+class ProcessMetrics(BaseModel):
+    cpu: float
+    mem: float
+
+
+class DataTransferMetrics(BaseModel):
+    bytes_in: int
+    bytes_out: int
+
+
+class RequestsMetrics(BaseModel):
+    rps: int
+    total: int
+    successful: int
+    failed: int
+
+
+class ResponseTimeMetrics(BaseModel):
+    avg: float
+    min: int
+    max: int
+    p50: int
+    p90: int
+    p99: int
+
+
+class WorkerMetrics(BaseModel):
+    worker_id: str
+    data_transfer: DataTransferMetrics
+    requests: RequestsMetrics
+    response_time: ResponseTimeMetrics
+    process: ProcessMetrics
+
+
+class Meta(BaseModel):
+    identity: str
+    extension: str
+    instance: str
+    domain: str
+    tags: dict[str, str] | None = None
+
+
+class Status(BaseModel):
+    type: Literal["status"] = "status"
+    meta: Meta
+    metrics: WorkerMetrics
+
+
+class Event(BaseModel):
+    type: Literal["status", "response"]
+    data: Status | HTTPResponse = Field(discriminator="type")

mrok/http/config.py

@@ -8,21 +8,18 @@ from typing import Any
 import openziti
 from uvicorn import config
 
-from mrok.conf import get_settings
 from mrok.http.protocol import MrokHttpToolsProtocol
-from mrok.logging import setup_logging
+from mrok.http.types import ASGIApp
 
 logger = logging.getLogger("mrok.proxy")
 
 config.LIFESPAN["auto"] = "mrok.http.lifespan:MrokLifespan"
 
-ASGIApplication = config.ASGIApplication
-
 
 class MrokBackendConfig(config.Config):
     def __init__(
         self,
-        app: ASGIApplication | Callable[..., Any] | str,
+        app: ASGIApp | Callable[..., Any] | str,
         identity_file: str | Path,
         ziti_load_timeout_ms: int = 5000,
         backlog: int = 2048,
@@ -62,4 +59,4 @@ class MrokBackendConfig(config.Config):
         return sock
 
     def configure_logging(self) -> None:
-        setup_logging(get_settings())
+        return

mrok/http/constants.py

@@ -0,0 +1,22 @@
+MAX_REQUEST_BODY_BYTES = 2 * 1024 * 1024
+MAX_RESPONSE_BODY_BYTES = 5 * 1024 * 1024
+
+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/",)