moru 0.1.0__py3-none-any.whl

moru/sandbox/network.py

@@ -0,0 +1,8 @@
+"""
+Network configuration helpers for Moru sandboxes.
+"""
+
+"""
+CIDR range that represents all traffic.
+"""
+ALL_TRAFFIC = "0.0.0.0/0"

moru/sandbox/sandbox_api.py

@@ -0,0 +1,210 @@
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional, TypedDict, Union
+
+from typing_extensions import NotRequired, Unpack
+
+from moru import ConnectionConfig
+from moru.api.client.models import ListedSandbox, SandboxDetail, SandboxState
+from moru.connection_config import ApiParams
+from moru.sandbox.mcp import McpServer as BaseMcpServer
+
+
+class GitHubMcpServerConfig(TypedDict):
+    """
+    Configuration for a GitHub-based MCP server.
+    """
+
+    run_cmd: str
+    """
+    Command to run the MCP server. Must start a stdio-compatible server.
+    """
+    install_cmd: NotRequired[str]
+    """
+    Command to install dependencies for the MCP server. Working directory is the root of the github repository.
+    """
+    envs: NotRequired[Dict[str, str]]
+    """
+    Environment variables to set in the MCP process.
+    """
+
+
+# Extended MCP server configuration that includes base servers
+# and allows dynamic GitHub-based MCP servers with custom run and install commands.
+# For GitHub servers, use keys in the format "github/owner/repo"
+GitHubMcpServer = Dict[str, Union[GitHubMcpServerConfig, Any]]
+
+# Union type that combines base MCP servers with GitHub-based servers
+McpServer = Union[BaseMcpServer, GitHubMcpServer]
+
+
+class SandboxNetworkOpts(TypedDict):
+    """
+    Sandbox network configuration options.
+    """
+
+    allow_out: NotRequired[List[str]]
+    """
+    Allow outbound traffic from the sandbox to the specified addresses.
+    If `allow_out` is not specified, all outbound traffic is allowed.
+
+    Examples:
+    - To allow traffic to specific addresses: `["1.1.1.1", "8.8.8.0/24"]`
+    """
+
+    deny_out: NotRequired[List[str]]
+    """
+    Deny outbound traffic from the sandbox to the specified addresses.
+
+    Examples:
+    - To deny traffic to specific addresses: `["1.1.1.1", "8.8.8.0/24"]`
+    """
+
+    allow_public_traffic: NotRequired[bool]
+    """
+    Controls whether sandbox URLs should be publicly accessible or require authentication.
+    Defaults to True.
+    """
+
+    mask_request_host: NotRequired[str]
+    """
+    Allows specifying a custom host mask for all sandbox requests.
+    Supports ${PORT} variable. Defaults to "${PORT}-sandboxid.moru.io".
+
+    Examples:
+    - Custom subdomain: `"${PORT}-myapp.example.com"`
+    """
+
+
+@dataclass
+class SandboxInfo:
+    """Information about a sandbox."""
+
+    sandbox_id: str
+    """Sandbox ID."""
+    sandbox_domain: Optional[str]
+    """Domain where the sandbox is hosted."""
+    template_id: str
+    """Template ID."""
+    name: Optional[str]
+    """Template name."""
+    metadata: Dict[str, str]
+    """Saved sandbox metadata."""
+    started_at: datetime
+    """Sandbox start time."""
+    end_at: datetime
+    """Sandbox expiration date."""
+    state: SandboxState
+    """Sandbox state."""
+    cpu_count: int
+    """Sandbox CPU count."""
+    memory_mb: int
+    """Sandbox Memory size in MiB."""
+    envd_version: str
+    """Envd version."""
+    _envd_access_token: Optional[str]
+    """Envd access token."""
+
+    @classmethod
+    def _from_sandbox_data(
+        cls,
+        sandbox: Union[ListedSandbox, SandboxDetail],
+        envd_access_token: Optional[str] = None,
+        sandbox_domain: Optional[str] = None,
+    ):
+        return cls(
+            sandbox_domain=sandbox_domain,
+            sandbox_id=sandbox.sandbox_id,
+            template_id=sandbox.template_id,
+            name=(sandbox.alias if isinstance(sandbox.alias, str) else None),
+            metadata=(sandbox.metadata if isinstance(sandbox.metadata, dict) else {}),
+            started_at=sandbox.started_at,
+            end_at=sandbox.end_at,
+            state=sandbox.state,
+            cpu_count=sandbox.cpu_count,
+            memory_mb=sandbox.memory_mb,
+            envd_version=sandbox.envd_version,
+            _envd_access_token=envd_access_token,
+        )
+
+    @classmethod
+    def _from_listed_sandbox(cls, listed_sandbox: ListedSandbox):
+        return cls._from_sandbox_data(listed_sandbox)
+
+    @classmethod
+    def _from_sandbox_detail(cls, sandbox_detail: SandboxDetail):
+        return cls._from_sandbox_data(
+            sandbox_detail,
+            (
+                sandbox_detail.envd_access_token
+                if isinstance(sandbox_detail.envd_access_token, str)
+                else None
+            ),
+            sandbox_domain=(
+                sandbox_detail.domain
+                if isinstance(sandbox_detail.domain, str)
+                else None
+            ),
+        )
+
+
+@dataclass
+class SandboxQuery:
+    """Query parameters for listing sandboxes."""
+
+    metadata: Optional[dict[str, str]] = None
+    """Filter sandboxes by metadata."""
+
+    state: Optional[list[SandboxState]] = None
+    """Filter sandboxes by state."""
+
+
+@dataclass
+class SandboxMetrics:
+    """Sandbox metrics."""
+
+    cpu_count: int
+    """Number of CPUs."""
+    cpu_used_pct: float
+    """CPU usage percentage."""
+    disk_total: int
+    """Total disk space in bytes."""
+    disk_used: int
+    """Disk used in bytes."""
+    mem_total: int
+    """Total memory in bytes."""
+    mem_used: int
+    """Memory used in bytes."""
+    timestamp: datetime
+    """Timestamp of the metric entry."""
+
+
+class SandboxPaginatorBase:
+    def __init__(
+        self,
+        query: Optional[SandboxQuery] = None,
+        limit: Optional[int] = None,
+        next_token: Optional[str] = None,
+        **opts: Unpack[ApiParams],
+    ):
+        self._config = ConnectionConfig(**opts)
+
+        self.query = query
+        self.limit = limit
+
+        self._has_next = True
+        self._next_token = next_token
+
+    @property
+    def has_next(self) -> bool:
+        """
+        Returns True if there are more items to fetch.
+        """
+        return self._has_next
+
+    @property
+    def next_token(self) -> Optional[str]:
+        """
+        Returns the next token to use for pagination.
+        """
+        return self._next_token

moru/sandbox/signature.py

@@ -0,0 +1,45 @@
+import base64
+import hashlib
+import time
+
+from typing import Optional, TypedDict, Literal
+
+Operation = Literal["read", "write"]
+
+
+class Signature(TypedDict):
+    signature: str
+    expiration: Optional[int]  # Unix timestamp or None
+
+
+def get_signature(
+    path: str,
+    operation: Operation,
+    user: Optional[str],
+    envd_access_token: Optional[str],
+    expiration_in_seconds: Optional[int] = None,
+) -> Signature:
+    """
+    Generate a v1 signature for sandbox file URLs.
+    """
+    if not envd_access_token:
+        raise ValueError("Access token is not set and signature cannot be generated!")
+
+    expiration = (
+        int(time.time()) + expiration_in_seconds if expiration_in_seconds else None
+    )
+
+    # if user is None, set it to empty string to handle default user
+    if user is None:
+        user = ""
+
+    raw = (
+        f"{path}:{operation}:{user}:{envd_access_token}"
+        if expiration is None
+        else f"{path}:{operation}:{user}:{envd_access_token}:{expiration}"
+    )
+
+    digest = hashlib.sha256(raw.encode("utf-8")).digest()
+    encoded = base64.b64encode(digest).rstrip(b"=").decode("ascii")
+
+    return {"signature": f"v1_{encoded}", "expiration": expiration}

moru/sandbox/utils.py

@@ -0,0 +1,34 @@
+from typing import TypeVar, Any, cast, Optional, Type
+import functools
+
+T = TypeVar("T")
+
+
+class class_method_variant(object):
+    def __init__(self, class_method_name):
+        self.class_method_name = class_method_name
+
+    method: Any
+
+    def __call__(self, method: T) -> T:
+        self.method = method
+        return cast(T, self)
+
+    def __get__(self, obj, objtype: Optional[Type[Any]] = None):
+        @functools.wraps(self.method)
+        def _wrapper(*args, **kwargs):
+            if obj is not None:
+                # Method was called as an instance method, e.g.
+                # instance.method(...)
+                return self.method(obj, *args, **kwargs)
+            elif len(args) > 0 and objtype is not None and isinstance(args[0], objtype):
+                # Method was called as a class method with the instance as the
+                # first argument, e.g. Class.method(instance, ...) which in
+                # Python is the same thing as calling an instance method
+                return self.method(args[0], *args[1:], **kwargs)
+            else:
+                # Method was called as a class method, e.g. Class.method(...)
+                class_method = getattr(objtype, self.class_method_name)
+                return class_method(*args, **kwargs)
+
+        return _wrapper

moru/sandbox_async/commands/command.py

@@ -0,0 +1,336 @@
+from typing import Dict, List, Literal, Optional, Union, overload
+
+import moru_connect
+import httpcore
+from packaging.version import Version
+from moru.connection_config import (
+    ConnectionConfig,
+    Username,
+    KEEPALIVE_PING_HEADER,
+    KEEPALIVE_PING_INTERVAL_SEC,
+)
+from moru.envd.process import process_connect, process_pb2
+from moru.envd.rpc import authentication_header, handle_rpc_exception
+from moru.envd.versions import ENVD_COMMANDS_STDIN
+from moru.exceptions import SandboxException
+from moru.sandbox.commands.main import ProcessInfo
+from moru.sandbox.commands.command_handle import CommandResult
+from moru.sandbox_async.commands.command_handle import AsyncCommandHandle, Stderr, Stdout
+from moru.sandbox_async.utils import OutputHandler
+
+
+class Commands:
+    """
+    Module for executing commands in the sandbox.
+    """
+
+    def __init__(
+        self,
+        envd_api_url: str,
+        connection_config: ConnectionConfig,
+        pool: httpcore.AsyncConnectionPool,
+        envd_version: Version,
+    ) -> None:
+        self._connection_config = connection_config
+        self._envd_version = envd_version
+        self._rpc = process_connect.ProcessClient(
+            envd_api_url,
+            # TODO: Fix and enable compression again — the headers compression is not solved for streaming.
+            # compressor=moru_connect.GzipCompressor,
+            async_pool=pool,
+            json=True,
+            headers=connection_config.sandbox_headers,
+        )
+
+    async def list(
+        self,
+        request_timeout: Optional[float] = None,
+    ) -> List[ProcessInfo]:
+        """
+        Lists all running commands and PTY sessions.
+
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: List of running commands and PTY sessions
+        """
+        try:
+            res = await self._rpc.alist(
+                process_pb2.ListRequest(),
+                request_timeout=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+            return [
+                ProcessInfo(
+                    pid=p.pid,
+                    tag=p.tag,
+                    cmd=p.config.cmd,
+                    args=list(p.config.args),
+                    envs=dict(p.config.envs),
+                    cwd=p.config.cwd,
+                )
+                for p in res.processes
+            ]
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    async def kill(
+        self,
+        pid: int,
+        request_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Kill a running command specified by its process ID.
+        It uses `SIGKILL` signal to kill the command.
+
+        :param pid: Process ID of the command. You can get the list of processes using `sandbox.commands.list()`
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `True` if the command was killed, `False` if the command was not found
+        """
+        try:
+            await self._rpc.asend_signal(
+                process_pb2.SendSignalRequest(
+                    process=process_pb2.ProcessSelector(pid=pid),
+                    signal=process_pb2.Signal.SIGNAL_SIGKILL,
+                ),
+                request_timeout=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+            return True
+        except Exception as e:
+            if isinstance(e, moru_connect.ConnectException):
+                if e.status == moru_connect.Code.not_found:
+                    return False
+            raise handle_rpc_exception(e)
+
+    async def send_stdin(
+        self,
+        pid: int,
+        data: str,
+        request_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Send data to command stdin.
+
+        :param pid Process ID of the command. You can get the list of processes using `sandbox.commands.list()`.
+        :param data: Data to send to the command
+        :param request_timeout: Timeout for the request in **seconds**
+        """
+        try:
+            await self._rpc.asend_input(
+                process_pb2.SendInputRequest(
+                    process=process_pb2.ProcessSelector(pid=pid),
+                    input=process_pb2.ProcessInput(
+                        stdin=data.encode(),
+                    ),
+                ),
+                request_timeout=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    @overload
+    async def run(
+        self,
+        cmd: str,
+        background: Union[Literal[False], None] = None,
+        envs: Optional[Dict[str, str]] = None,
+        user: Optional[Username] = None,
+        cwd: Optional[str] = None,
+        on_stdout: Optional[OutputHandler[Stdout]] = None,
+        on_stderr: Optional[OutputHandler[Stderr]] = None,
+        stdin: Optional[bool] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ) -> CommandResult:
+        """
+        Start a new command and wait until it finishes executing.
+
+        :param cmd: Command to execute
+        :param background: **`False` if the command should be executed in the foreground**, `True` if the command should be executed in the background
+        :param envs: Environment variables used for the command
+        :param user: User to run the command as
+        :param cwd: Working directory to run the command
+        :param on_stdout: Callback for command stdout output
+        :param on_stderr: Callback for command stderr output
+        :param stdin: If `True`, the command will have a stdin stream that you can send data to using `sandbox.commands.send_stdin()`
+        :param timeout: Timeout for the command connection in **seconds**. Using `0` will not limit the command connection time
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `CommandResult` result of the command execution
+        """
+        ...
+
+    @overload
+    async def run(
+        self,
+        cmd: str,
+        background: Literal[True],
+        envs: Optional[Dict[str, str]] = None,
+        user: Optional[Username] = None,
+        cwd: Optional[str] = None,
+        on_stdout: Optional[OutputHandler[Stdout]] = None,
+        on_stderr: Optional[OutputHandler[Stderr]] = None,
+        stdin: Optional[bool] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ) -> AsyncCommandHandle:
+        """
+        Start a new command and return a handle to interact with it.
+
+        :param cmd: Command to execute
+        :param background: `False` if the command should be executed in the foreground, **`True` if the command should be executed in the background**
+        :param envs: Environment variables used for the command
+        :param user: User to run the command as
+        :param cwd: Working directory to run the command
+        :param on_stdout: Callback for command stdout output
+        :param on_stderr: Callback for command stderr output
+        :param stdin: If `True`, the command will have a stdin stream that you can send data to using `sandbox.commands.send_stdin()`
+        :param timeout: Timeout for the command connection in **seconds**. Using `0` will not limit the command connection time
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `AsyncCommandHandle` handle to interact with the running command
+        """
+        ...
+
+    async def run(
+        self,
+        cmd: str,
+        background: Union[bool, None] = None,
+        envs: Optional[Dict[str, str]] = None,
+        user: Optional[Username] = None,
+        cwd: Optional[str] = None,
+        on_stdout: Optional[OutputHandler[Stdout]] = None,
+        on_stderr: Optional[OutputHandler[Stderr]] = None,
+        stdin: Optional[bool] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ):
+        # Check version for stdin support
+        if stdin is False and self._envd_version < ENVD_COMMANDS_STDIN:
+            raise SandboxException(
+                f"Sandbox envd version {self._envd_version} can't specify stdin, it's always turned on. "
+                f"Please rebuild your template if you need this feature."
+            )
+
+        # Default to `False`
+        stdin = stdin or False
+
+        proc = await self._start(
+            cmd,
+            envs,
+            user,
+            cwd,
+            timeout,
+            request_timeout,
+            stdin,
+            on_stdout=on_stdout,
+            on_stderr=on_stderr,
+        )
+
+        return proc if background else await proc.wait()
+
+    async def _start(
+        self,
+        cmd: str,
+        envs: Optional[Dict[str, str]],
+        user: Username,
+        cwd: Optional[str],
+        timeout: Optional[float],
+        request_timeout: Optional[float],
+        stdin: bool,
+        on_stdout: Optional[OutputHandler[Stdout]],
+        on_stderr: Optional[OutputHandler[Stderr]],
+    ) -> AsyncCommandHandle:
+        events = self._rpc.astart(
+            process_pb2.StartRequest(
+                process=process_pb2.ProcessConfig(
+                    cmd="/bin/bash",
+                    envs=envs,
+                    args=["-l", "-c", cmd],
+                    cwd=cwd,
+                ),
+                stdin=stdin,
+            ),
+            headers={
+                **authentication_header(self._envd_version, user),
+                KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
+            },
+            timeout=timeout,
+            request_timeout=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+        )
+
+        try:
+            start_event = await events.__anext__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to start process: expected start event, got {start_event}"
+                )
+
+            return AsyncCommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+                on_stdout=on_stdout,
+                on_stderr=on_stderr,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    async def connect(
+        self,
+        pid: int,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+        on_stdout: Optional[OutputHandler[Stdout]] = None,
+        on_stderr: Optional[OutputHandler[Stderr]] = None,
+    ) -> AsyncCommandHandle:
+        """
+        Connects to a running command.
+        You can use `AsyncCommandHandle.wait()` to wait for the command to finish and get execution results.
+
+        :param pid: Process ID of the command to connect to. You can get the list of processes using `sandbox.commands.list()`
+        :param request_timeout: Request timeout in **seconds**
+        :param timeout: Timeout for the command connection in **seconds**. Using `0` will not limit the command connection time
+        :param on_stdout: Callback for command stdout output
+        :param on_stderr: Callback for command stderr output
+
+        :return: `AsyncCommandHandle` handle to interact with the running command
+        """
+        events = self._rpc.aconnect(
+            process_pb2.ConnectRequest(
+                process=process_pb2.ProcessSelector(pid=pid),
+            ),
+            timeout=timeout,
+            request_timeout=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+            headers={
+                KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
+            },
+        )
+
+        try:
+            start_event = await events.__anext__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to connect to process: expected start event, got {start_event}"
+                )
+
+            return AsyncCommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+                on_stdout=on_stdout,
+                on_stderr=on_stderr,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)