moru 0.1.0__py3-none-any.whl

moru/sandbox_sync/commands/command.py

@@ -0,0 +1,328 @@
+from typing import Callable, 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_sync.commands.command_handle import CommandHandle
+
+
+class Commands:
+    """
+    Module for executing commands in the sandbox.
+    """
+
+    def __init__(
+        self,
+        envd_api_url: str,
+        connection_config: ConnectionConfig,
+        pool: httpcore.ConnectionPool,
+        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,
+            pool=pool,
+            json=True,
+            headers=connection_config.sandbox_headers,
+        )
+
+    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 = self._rpc.list(
+                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)
+
+    def kill(
+        self,
+        pid: int,
+        request_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Kills 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:
+            self._rpc.send_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)
+
+    def send_stdin(
+        self,
+        pid: int,
+        data: str,
+        request_timeout: Optional[float] = 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:
+            self._rpc.send_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
+    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[Callable[[str], None]] = None,
+        on_stderr: Optional[Callable[[str], None]] = 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
+    def run(
+        self,
+        cmd: str,
+        background: Literal[True],
+        envs: Optional[Dict[str, str]] = None,
+        user: Optional[Username] = None,
+        cwd: Optional[str] = None,
+        on_stdout: None = None,
+        on_stderr: None = None,
+        stdin: Optional[bool] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ) -> CommandHandle:
+        """
+        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 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: `CommandHandle` handle to interact with the running command
+        """
+        ...
+
+    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[Callable[[str], None]] = None,
+        on_stderr: Optional[Callable[[str], None]] = 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 = self._start(
+            cmd,
+            envs,
+            user,
+            cwd,
+            stdin,
+            timeout,
+            request_timeout,
+        )
+
+        return (
+            proc
+            if background
+            else proc.wait(
+                on_stdout=on_stdout,
+                on_stderr=on_stderr,
+            )
+        )
+
+    def _start(
+        self,
+        cmd: str,
+        envs: Optional[Dict[str, str]],
+        user: Username,
+        cwd: Optional[str],
+        stdin: bool,
+        timeout: Optional[float],
+        request_timeout: Optional[float],
+    ):
+        events = self._rpc.start(
+            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 = events.__next__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to start process: expected start event, got {start_event}"
+                )
+
+            return CommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def connect(
+        self,
+        pid: int,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ):
+        """
+        Connects to a running command.
+        You can use `CommandHandle.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 timeout: Timeout for the connection in **seconds**. Using `0` will not limit the connection time
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `CommandHandle` handle to interact with the running command
+        """
+        events = self._rpc.connect(
+            process_pb2.ConnectRequest(
+                process=process_pb2.ProcessSelector(pid=pid),
+            ),
+            headers={
+                KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
+            },
+            timeout=timeout,
+            request_timeout=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+        )
+
+        try:
+            start_event = events.__next__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to connect to process: expected start event, got {start_event}"
+                )
+
+            return CommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)

moru/sandbox_sync/commands/command_handle.py

@@ -0,0 +1,150 @@
+from typing import Optional, Callable, Any, Generator, Union, Tuple
+
+from moru.envd.rpc import handle_rpc_exception
+from moru.envd.process import process_pb2
+from moru.sandbox.commands.command_handle import (
+    CommandExitException,
+    CommandResult,
+    Stderr,
+    Stdout,
+    PtyOutput,
+)
+
+
+class CommandHandle:
+    """
+    Command execution handle.
+
+    It provides methods for waiting for the command to finish, retrieving stdout/stderr, and killing the command.
+    """
+
+    @property
+    def pid(self):
+        """
+        Command process ID.
+        """
+        return self._pid
+
+    def __init__(
+        self,
+        pid: int,
+        handle_kill: Callable[[], bool],
+        events: Generator[
+            Union[process_pb2.StartResponse, process_pb2.ConnectResponse], Any, None
+        ],
+    ):
+        self._pid = pid
+        self._handle_kill = handle_kill
+        self._events = events
+
+        self._stdout: str = ""
+        self._stderr: str = ""
+
+        self._result: Optional[CommandResult] = None
+        self._iteration_exception: Optional[Exception] = None
+
+    def __iter__(self):
+        """
+        Iterate over the command output.
+
+        :return: Generator of command outputs
+        """
+        return self._handle_events()
+
+    def _handle_events(
+        self,
+    ) -> Generator[
+        Union[
+            Tuple[Stdout, None, None],
+            Tuple[None, Stderr, None],
+            Tuple[None, None, PtyOutput],
+        ],
+        None,
+        None,
+    ]:
+        try:
+            for event in self._events:
+                if event.event.HasField("data"):
+                    if event.event.data.stdout:
+                        out = event.event.data.stdout.decode("utf-8", "replace")
+                        self._stdout += out
+                        yield out, None, None
+                    if event.event.data.stderr:
+                        out = event.event.data.stderr.decode("utf-8", "replace")
+                        self._stderr += out
+                        yield None, out, None
+                    if event.event.data.pty:
+                        yield None, None, event.event.data.pty
+                if event.event.HasField("end"):
+                    self._result = CommandResult(
+                        stdout=self._stdout,
+                        stderr=self._stderr,
+                        exit_code=event.event.end.exit_code,
+                        error=event.event.end.error,
+                    )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def disconnect(self) -> None:
+        """
+        Disconnect from the command.
+
+        The command is not killed, but SDK stops receiving events from the command.
+        You can reconnect to the command using `sandbox.commands.connect` method.
+        """
+        self._events.close()
+
+    def wait(
+        self,
+        on_pty: Optional[Callable[[PtyOutput], None]] = None,
+        on_stdout: Optional[Callable[[str], None]] = None,
+        on_stderr: Optional[Callable[[str], None]] = None,
+    ) -> CommandResult:
+        """
+        Wait for the command to finish and returns the result.
+        If the command exits with a non-zero exit code, it throws a `CommandExitException`.
+
+        :param on_pty: Callback for pty output
+        :param on_stdout: Callback for stdout output
+        :param on_stderr: Callback for stderr output
+
+        :return: `CommandResult` result of command execution
+        """
+        try:
+            for stdout, stderr, pty in self:
+                if stdout is not None and on_stdout:
+                    on_stdout(stdout)
+                elif stderr is not None and on_stderr:
+                    on_stderr(stderr)
+                elif pty is not None and on_pty:
+                    on_pty(pty)
+        except StopIteration:
+            pass
+        except Exception as e:
+            self._iteration_exception = handle_rpc_exception(e)
+
+        if self._iteration_exception:
+            raise self._iteration_exception
+
+        if self._result is None:
+            raise Exception("Command ended without an end event")
+
+        if self._result.exit_code != 0:
+            raise CommandExitException(
+                stdout=self._stdout,
+                stderr=self._stderr,
+                exit_code=self._result.exit_code,
+                error=self._result.error,
+            )
+
+        return self._result
+
+    def kill(self) -> bool:
+        """
+        Kills the command.
+
+        It uses `SIGKILL` signal to kill the command.
+
+        :return: Whether the command was killed successfully
+        """
+        return self._handle_kill()

moru/sandbox_sync/commands/pty.py

@@ -0,0 +1,230 @@
+import moru_connect
+import httpcore
+
+from typing import Dict, Optional
+
+from packaging.version import Version
+from moru.envd.process import process_connect, process_pb2
+from moru.connection_config import (
+    Username,
+    ConnectionConfig,
+    KEEPALIVE_PING_HEADER,
+    KEEPALIVE_PING_INTERVAL_SEC,
+)
+from moru.exceptions import SandboxException
+from moru.envd.rpc import authentication_header, handle_rpc_exception
+from moru.sandbox.commands.command_handle import PtySize
+from moru.sandbox_sync.commands.command_handle import CommandHandle
+
+
+class Pty:
+    """
+    Module for interacting with PTYs (pseudo-terminals) in the sandbox.
+    """
+
+    def __init__(
+        self,
+        envd_api_url: str,
+        connection_config: ConnectionConfig,
+        pool: httpcore.ConnectionPool,
+        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,
+            pool=pool,
+            json=True,
+            headers=connection_config.sandbox_headers,
+        )
+
+    def kill(
+        self,
+        pid: int,
+        request_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Kill PTY.
+
+        :param pid: Process ID of the PTY
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `true` if the PTY was killed, `false` if the PTY was not found
+        """
+        try:
+            self._rpc.send_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)
+
+    def send_stdin(
+        self,
+        pid: int,
+        data: bytes,
+        request_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Send input to a PTY.
+
+        :param pid: Process ID of the PTY
+        :param data: Input data to send
+        :param request_timeout: Timeout for the request in **seconds**
+        """
+        try:
+            self._rpc.send_input(
+                process_pb2.SendInputRequest(
+                    process=process_pb2.ProcessSelector(pid=pid),
+                    input=process_pb2.ProcessInput(
+                        pty=data,
+                    ),
+                ),
+                request_timeout=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def create(
+        self,
+        size: PtySize,
+        user: Optional[Username] = None,
+        cwd: Optional[str] = None,
+        envs: Optional[Dict[str, str]] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ) -> CommandHandle:
+        """
+        Start a new PTY (pseudo-terminal).
+
+        :param size: Size of the PTY
+        :param user: User to use for the PTY
+        :param cwd: Working directory for the PTY
+        :param envs: Environment variables for the PTY
+        :param timeout: Timeout for the PTY in **seconds**
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: Handle to interact with the PTY
+        """
+        envs = envs or {}
+        envs["TERM"] = "xterm-256color"
+        events = self._rpc.start(
+            process_pb2.StartRequest(
+                process=process_pb2.ProcessConfig(
+                    cmd="/bin/bash",
+                    envs=envs,
+                    args=["-i", "-l"],
+                    cwd=cwd,
+                ),
+                pty=process_pb2.PTY(
+                    size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols)
+                ),
+            ),
+            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 = events.__next__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to start process: expected start event, got {start_event}"
+                )
+
+            return CommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def connect(
+        self,
+        pid: int,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ) -> CommandHandle:
+        """
+        Connect to a running PTY.
+
+        :param pid: Process ID of the PTY to connect to. You can get the list of running PTYs using `sandbox.pty.list()`.
+        :param timeout: Timeout for the PTY connection in **seconds**. Using `0` will not limit the connection time
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: Handle to interact with the PTY
+        """
+        events = self._rpc.connect(
+            process_pb2.ConnectRequest(
+                process=process_pb2.ProcessSelector(pid=pid),
+            ),
+            headers={
+                KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
+            },
+            timeout=timeout,
+            request_timeout=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+        )
+
+        try:
+            start_event = events.__next__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to connect to process: expected start event, got {start_event}"
+                )
+
+            return CommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def resize(
+        self,
+        pid: int,
+        size: PtySize,
+        request_timeout: Optional[float] = None,
+    ) -> None:
+        """
+        Resize PTY.
+        Call this when the terminal window is resized and the number of columns and rows has changed.
+
+        :param pid: Process ID of the PTY
+        :param size: New size of the PTY
+        :param request_timeout: Timeout for the request in **seconds**s
+        """
+        self._rpc.update(
+            process_pb2.UpdateRequest(
+                process=process_pb2.ProcessSelector(pid=pid),
+                pty=process_pb2.PTY(
+                    size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols),
+                ),
+            ),
+            request_timeout=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+        )