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

mrok/__init__.py

@@ -0,0 +1,6 @@
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("mrok")
+except PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.0.0.dev0"

mrok/agent/sidecar/__init__.py

@@ -0,0 +1,3 @@
+from mrok.agent.sidecar.main import run
+
+__all__ = ["run"]

mrok/agent/sidecar/app.py

@@ -0,0 +1,30 @@
+import asyncio
+import logging
+from collections.abc import Awaitable, Callable
+from pathlib import Path
+from typing import Any
+
+from mrok.http.forwarder import ForwardAppBase
+
+logger = logging.getLogger("mrok.proxy")
+
+Scope = dict[str, Any]
+ASGIReceive = Callable[[], Awaitable[dict[str, Any]]]
+ASGISend = Callable[[dict[str, Any]], Awaitable[None]]
+
+
+class ForwardApp(ForwardAppBase):
+    def __init__(
+        self, target_address: str | Path | tuple[str, int], read_chunk_size: int = 65536
+    ) -> None:
+        super().__init__(read_chunk_size=read_chunk_size)
+        self._target_address = target_address
+
+    async def select_backend(
+        self,
+        scope: Scope,
+        headers: dict[str, str],
+    ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter] | tuple[None, None]:
+        if isinstance(self._target_address, tuple):
+            return await asyncio.open_connection(*self._target_address)
+        return await asyncio.open_unix_connection(str(self._target_address))

mrok/agent/sidecar/main.py

@@ -0,0 +1,27 @@
+import asyncio
+import contextlib
+from functools import partial
+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
+
+
+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()
+
+
+def run(
+    identity_file: str,
+    target_addr: str | Path | tuple[str, int],
+    workers=4,
+    reload=False,
+):
+    start_fn = partial(run_sidecar, identity_file, target_addr)
+    master = Master(start_fn, workers=workers, reload=reload)
+    master.run()

mrok/agent/ziticorn.py

@@ -0,0 +1,29 @@
+import asyncio
+import contextlib
+import os
+from functools import partial
+
+from mrok.http.config import ASGIApplication, MrokBackendConfig
+from mrok.http.master import Master
+from mrok.http.server import MrokServer
+
+
+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 run(
+    app: ASGIApplication | str,
+    identity_file: str,
+    workers: int = 4,
+    reload: bool = False,
+):
+    start_fn = partial(run_ziticorn, app, identity_file)
+    master = Master(start_fn, workers=workers, reload=reload)
+    master.run()

mrok/cli/__init__.py

@@ -0,0 +1,3 @@
+from mrok.cli.main import app, run
+
+__all__ = ["app", "run"]

mrok/cli/commands/__init__.py

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

mrok/cli/commands/admin/__init__.py

@@ -0,0 +1,10 @@
+import typer
+
+from mrok.cli.commands.admin import bootstrap
+from mrok.cli.commands.admin.register import app as register_app
+from mrok.cli.commands.admin.unregister import app as unregister_app
+
+app = typer.Typer(help="mrok administrative commands.")
+app.add_typer(register_app)
+app.add_typer(unregister_app)
+bootstrap.register(app)

mrok/cli/commands/admin/bootstrap.py

@@ -0,0 +1,58 @@
+import asyncio
+import json
+import logging
+from pathlib import Path
+from typing import Annotated, Any
+
+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.ziti.bootstrap import bootstrap_identity
+
+logger = logging.getLogger(__name__)
+
+
+async def bootstrap(
+    settings: Settings, forced: bool, tags: TagsType | None
+) -> tuple[str, dict[str, Any] | None]:
+    async with ZitiManagementAPI(settings) as mgmt_api, ZitiClientAPI(settings) as client_api:
+        return await bootstrap_identity(
+            mgmt_api,
+            client_api,
+            settings.proxy.identity,
+            settings.proxy.mode,
+            forced,
+            tags,
+        )
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("bootstrap")
+    def run_bootstrap(
+        ctx: typer.Context,
+        identity_file: Path = typer.Argument(
+            Path("proxy_identity.json"),
+            help="Path to identity output file",
+            writable=True,
+        ),
+        forced: bool = typer.Option(
+            False,
+            "--force",
+            help="Regenerate identity even if it already exists",
+        ),
+        tags: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--tag",
+                "-t",
+                help="Add tag",
+                show_default=True,
+            ),
+        ] = None,
+    ):
+        """Run the mrok bootstrap."""
+        _, identity_json = asyncio.run(bootstrap(ctx.obj, forced, parse_tags(tags)))
+        if identity_json:
+            json.dump(identity_json, identity_file.open("w"))

mrok/cli/commands/admin/register/__init__.py

@@ -0,0 +1,8 @@
+import typer
+
+from mrok.cli.commands.admin.register import extensions, instances
+
+app = typer.Typer(name="register", help="Register resources into OpenZiti.")
+
+extensions.register(app)
+instances.register(app)

mrok/cli/commands/admin/register/extensions.py

@@ -0,0 +1,46 @@
+import asyncio
+import re
+from typing import Annotated
+
+import typer
+from rich import print
+
+from mrok.cli.commands.admin.utils import parse_tags
+from mrok.conf import Settings
+from mrok.ziti.api import ZitiManagementAPI
+from mrok.ziti.services import register_extension
+
+RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
+
+
+async def do_register(settings: Settings, extension_id: str, tags: list[str] | None):
+    async with ZitiManagementAPI(settings) as api:
+        await register_extension(settings, api, extension_id, tags=parse_tags(tags))
+
+
+def validate_extension_id(extension_id: str) -> str:
+    if not RE_EXTENSION_ID.fullmatch(extension_id):
+        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+    return extension_id
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("extension")
+    def register_extension(
+        ctx: typer.Context,
+        extension_id: str = typer.Argument(
+            ..., callback=validate_extension_id, help="Extension ID in format EXT-xxxx-yyyy"
+        ),
+        tags: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--tag",
+                "-t",
+                help="Add tag",
+                show_default=True,
+            ),
+        ] = None,
+    ):
+        """Register a new Extension in OpenZiti (service)."""
+        asyncio.run(do_register(ctx.obj, extension_id, tags))
+        print(f"🍻 [green]Extension [bold]{extension_id}[/bold] registered.[/green]")

mrok/cli/commands/admin/register/instances.py

@@ -0,0 +1,60 @@
+import asyncio
+import json
+import re
+from pathlib import Path
+from typing import Annotated
+
+import typer
+
+from mrok.cli.commands.admin.utils import parse_tags
+from mrok.conf import Settings
+from mrok.ziti.api import ZitiClientAPI, ZitiManagementAPI
+from mrok.ziti.identities import register_instance
+
+RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
+
+
+async def do_register(
+    settings: Settings, extension_id: str, instance_id: str, tags: list[str] | None
+):
+    async with ZitiManagementAPI(settings) as mgmt_api, ZitiClientAPI(settings) as client_api:
+        return await register_instance(
+            mgmt_api, client_api, extension_id, instance_id, tags=parse_tags(tags)
+        )
+
+
+def validate_extension_id(extension_id: str):
+    if not RE_EXTENSION_ID.fullmatch(extension_id):
+        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+    return extension_id
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("instance")
+    def register_instance(
+        ctx: typer.Context,
+        extension_id: str = typer.Argument(
+            ..., callback=validate_extension_id, help="Extension ID in format EXT-xxxx-yyyy"
+        ),
+        instance_id: str = typer.Argument(..., help="Instance ID"),
+        output: Path = typer.Argument(
+            ...,
+            file_okay=True,
+            dir_okay=False,
+            writable=True,
+            resolve_path=True,
+            help="Output file (default: stdout)",
+        ),
+        tags: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--tag",
+                "-t",
+                help="Add tag",
+                show_default=True,
+            ),
+        ] = None,
+    ):
+        """Register a new Extension Instance in OpenZiti (identity)."""
+        _, identity_file = asyncio.run(do_register(ctx.obj, extension_id, instance_id, tags))
+        json.dump(identity_file, output.open("w"))

mrok/cli/commands/admin/unregister/__init__.py

@@ -0,0 +1,8 @@
+import typer
+
+from mrok.cli.commands.admin.unregister import extensions, instances
+
+app = typer.Typer(name="unregister", help="Unregister resources from OpenZiti.")
+
+extensions.register(app)
+instances.register(app)

mrok/cli/commands/admin/unregister/extensions.py

@@ -0,0 +1,33 @@
+import asyncio
+import re
+
+import typer
+
+from mrok.conf import Settings
+from mrok.ziti.api import ZitiManagementAPI
+from mrok.ziti.services import unregister_extension
+
+RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
+
+
+async def do_unregister(settings: Settings, extension_id: str):
+    async with ZitiManagementAPI(settings) as api:
+        await unregister_extension(settings, api, extension_id)
+
+
+def validate_extension_id(extension_id: str):
+    if not RE_EXTENSION_ID.fullmatch(extension_id):
+        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+    return extension_id
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("extension")
+    def unregister_extension(
+        ctx: typer.Context,
+        extension_id: str = typer.Argument(
+            ..., callback=validate_extension_id, help="Extension ID in format EXT-xxxx-yyyy"
+        ),
+    ):
+        """Unregister a new Extension in OpenZiti (service)."""
+        asyncio.run(do_unregister(ctx.obj, extension_id))

mrok/cli/commands/admin/unregister/instances.py

@@ -0,0 +1,34 @@
+import asyncio
+import re
+
+import typer
+
+from mrok.conf import Settings
+from mrok.ziti.api import ZitiManagementAPI
+from mrok.ziti.identities import unregister_instance
+
+RE_EXTENSION_ID = re.compile(r"(?i)EXT-\d{4}-\d{4}")
+
+
+async def do_unregister(settings: Settings, extension_id: str, instance_id: str):
+    async with ZitiManagementAPI(settings) as api:
+        await unregister_instance(api, extension_id, instance_id)
+
+
+def validate_extension_id(extension_id: str):
+    if not RE_EXTENSION_ID.fullmatch(extension_id):
+        raise typer.BadParameter("ext_id must match EXT-xxxx-yyyy (case-insensitive)")
+    return extension_id
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("instance")
+    def unregister_instance(
+        ctx: typer.Context,
+        extension_id: str = typer.Argument(
+            ..., callback=validate_extension_id, help="Extension ID in format EXT-xxxx-yyyy"
+        ),
+        instance_id: str = typer.Argument(..., help="Instance ID"),
+    ):
+        """Register a new Extension Instance in OpenZiti (identity)."""
+        asyncio.run(do_unregister(ctx.obj, extension_id, instance_id))

mrok/cli/commands/admin/utils.py

@@ -0,0 +1,23 @@
+import typer
+
+from mrok.ziti.api import TagsType
+
+
+def parse_tags(pairs: list[str] | None) -> TagsType | None:
+    if not pairs:
+        return None
+
+    result: dict[str, str | bool | None] = {}
+    for item in pairs:
+        if "=" not in item:
+            raise typer.BadParameter(f"Invalid format {item!r}, expected key=value")
+        key, raw = item.split("=", 1)
+        raw_lower = raw.strip().lower()
+        if raw_lower in ("true", "false"):
+            val: str | bool | None = raw_lower == "true"
+        elif raw == "":
+            val = None
+        else:
+            val = raw
+        result[key.strip()] = val
+    return result

mrok/cli/commands/agent/__init__.py

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

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

@@ -0,0 +1,7 @@
+import typer
+
+from mrok.cli.commands.agent.run import asgi, sidecar
+
+app = typer.Typer(name="run", help="Run mrok agent.")
+asgi.register(app)
+sidecar.register(app)

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

@@ -0,0 +1,49 @@
+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
+
+
+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",
+        ),
+        workers: Annotated[
+            int,
+            typer.Option(
+                "--workers",
+                "-w",
+                help=f"Number of workers. Default: {default_workers}",
+                show_default=True,
+            ),
+        ] = default_workers,
+        reload: Annotated[
+            bool,
+            typer.Option(
+                "--reload",
+                "-r",
+                help="Enable auto-reload. Default: False",
+                show_default=True,
+            ),
+        ] = False,
+    ):
+        """Run an ASGI application exposing it through OpenZiti network."""
+        ziticorn.run(app, str(identity_file), workers=workers, reload=reload)

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

@@ -0,0 +1,54 @@
+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
+
+
+default_workers = number_of_workers()
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("sidecar")
+    def run_sidecar(
+        identity_file: Path = typer.Argument(
+            ...,
+            help="Identity json file",
+        ),
+        target: Path = typer.Argument(
+            ...,
+            help="Target service (host:port or path to unix domain socket)",
+        ),
+        workers: Annotated[
+            int,
+            typer.Option(
+                "--workers",
+                "-w",
+                help=f"Number of workers. Default: {default_workers}",
+                show_default=True,
+            ),
+        ] = default_workers,
+        reload: Annotated[
+            bool,
+            typer.Option(
+                "--reload",
+                "-r",
+                help="Enable auto-reload. Default: False",
+                show_default=True,
+            ),
+        ] = False,
+    ):
+        """Run a Sidecar Proxy to expose a web application through OpenZiti."""
+        if ":" in str(target):
+            host, port = str(target).split(":", 1)
+            target_addr = (host or "127.0.0.1", int(port))
+        else:
+            target_addr = str(target)  # type: ignore
+
+        sidecar.run(str(identity_file), target_addr, workers=workers, reload=reload)

mrok/cli/commands/controller/__init__.py

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

mrok/cli/commands/controller/openapi.py

@@ -0,0 +1,47 @@
+import json
+from enum import Enum
+from pathlib import Path
+from typing import Annotated
+
+import typer
+import yaml
+
+from mrok.controller.openapi import generate_openapi_spec
+
+
+class OutputFormat(str, Enum):
+    json = "json"
+    yaml = "yaml"
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("openapi")
+    def generate_spec(
+        ctx: typer.Context,
+        output: Annotated[
+            Path | None,
+            typer.Option(
+                "--output",
+                "-o",
+                help="Output file",
+            ),
+        ] = Path("mrok_openapi_spec.yml"),
+        output_format: Annotated[
+            OutputFormat,
+            typer.Option(
+                "--output-format",
+                "-f",
+                help="Output file format",
+            ),
+        ] = OutputFormat.yaml,
+    ):
+        """
+        Generates the mrok controller OpenAPI spec file.
+        """
+        from mrok.controller.app import app
+
+        dump_fn = json.dump if output_format == OutputFormat.json else yaml.dump
+        spec = generate_openapi_spec(app, ctx.obj)
+
+        with open(output, "w") as f:  # type: ignore
+            dump_fn(spec, f, indent=2)

mrok/cli/commands/controller/run.py

@@ -0,0 +1,87 @@
+import multiprocessing
+from collections.abc import Callable
+from typing import Annotated, Any
+
+import typer
+from gunicorn.app.base import BaseApplication
+
+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 {}
+        self.application = application
+        super().__init__()
+
+    def load_config(self):
+        config = {
+            key: value
+            for key, value in self.options.items()
+            if key in self.cfg.settings and value is not None
+        }
+        for key, value in config.items():
+            self.cfg.set(key.lower(), value)
+
+    def load(self):
+        return self.application
+
+
+default_workers = number_of_workers()
+
+
+def register(app: typer.Typer) -> None:
+    @app.command("run")
+    def run_controller(
+        ctx: typer.Context,
+        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,
+        dev: Annotated[
+            bool,
+            typer.Option(
+                "--reload",
+                "-r",
+                help="Enable auto-reload. Default: False",
+                show_default=True,
+            ),
+        ] = False,
+    ):
+        """Run the mrok controller with Gunicorn and Uvicorn workers."""
+        options = {
+            "bind": f"{host}:{port}",
+            "workers": workers,
+            "worker_class": "uvicorn_worker.UvicornWorker",
+            "logconfig_dict": get_logging_config(ctx.obj),
+            "reload": dev,
+        }
+        StandaloneApplication(asgi_app, options).run()