dais-shell 0.1.1__py3-none-any.whl

dais_shell/__init__.py

@@ -0,0 +1,62 @@
+import platform
+from typing import TypeAlias
+from .env_builder import EnvBuilder
+from .iostream_reader import IOStreamReaderResult, IOStreamReaderStatus
+from .runtimes import BaseShellRuntime, BashRuntime, PowerShellRuntime
+from .types import CommandStep
+from .types.exceptions import ShellError, CommandNotFoundError, ShellRuntimeNotFoundError, ForbiddenShellTargetError
+from .constants import DEFAULT_COMMAND_BLACKLIST
+
+ShellResult: TypeAlias = IOStreamReaderResult
+ShellResultStatus: TypeAlias = IOStreamReaderStatus
+
+class AgentShell:
+    def __init__(self,
+                 command_blacklist: set[str] | None = None,
+                 env_extra: dict[str, str] | None = None,
+                 max_lines: int = 10000,
+                 ):
+        self._runtime = self._create_runtime(max_lines)
+        self._command_blacklist = command_blacklist or DEFAULT_COMMAND_BLACKLIST
+        self._env_builder = EnvBuilder(blacklist=None, extra=env_extra)
+
+    @staticmethod
+    def _create_runtime(max_lines: int) -> BaseShellRuntime:
+        if platform.system() == "Windows":
+            return PowerShellRuntime(max_lines)
+        else:
+            return BashRuntime(max_lines)
+
+    def run_sync(self,
+                 step: CommandStep,
+                 on_stdout=None,
+                 on_stderr=None
+                 ) -> ShellResult:
+        step.validate_command(self._command_blacklist)
+        step.env = (self._env_builder
+                        .with_extra(step.env or {})
+                        .build())
+        return self._runtime.run_sync(step, on_stdout, on_stderr)
+
+    async def run(self,
+                  step: CommandStep,
+                  on_stdout=None,
+                  on_stderr=None
+                  ) -> ShellResult:
+        step.validate_command(self._command_blacklist)
+        step.env = (self._env_builder
+                        .with_extra(step.env or {})
+                        .build())
+        return await self._runtime.run(step, on_stdout, on_stderr)
+
+__all__ = [
+    "AgentShell",
+    "CommandStep",
+    "ShellResult",
+    "ShellResultStatus",
+
+    "ShellError",
+    "CommandNotFoundError",
+    "ShellRuntimeNotFoundError",
+    "ForbiddenShellTargetError",
+]

dais_shell/constants.py

@@ -0,0 +1,6 @@
+DEFAULT_COMMAND_BLACKLIST = {
+    "cmd", "cmd.exe",
+    "powershell", "powershell.exe",
+    "pwsh", "pwsh.exe",
+    "bash", "sh", "zsh", "fish"
+}

dais_shell/env_builder.py

@@ -0,0 +1,103 @@
+import os
+import platform
+
+WINDOWS_ESSENTIAL_VARS = {
+    "SYSTEMROOT",
+    "SYSTEMDRIVE",
+    "PATHEXT",
+    "COMSPEC",
+    "WINDIR",
+
+    "USERPROFILE",
+    "USERNAME",
+    "COMPUTERNAME",
+    "APPDATA",
+    "LOCALAPPDATA",
+    "PUBLIC",
+
+    "PROCESSOR_ARCHITECTURE"
+}
+
+UNIX_ESSENTIAL_VARS = {
+    "HOME",
+    "USER",
+    "LOGNAME",
+    "SHELL",
+    "TERM",
+    "LINES", "COLUMNS",
+    "PWD",
+    "XDG_RUNTIME_DIR", "XDG_CONFIG_HOME"
+}
+
+MACOS_ESSENTIAL_VARS = UNIX_ESSENTIAL_VARS | {
+    "DYLD_LIBRARY_PATH",
+    "DYLD_FRAMEWORK_PATH",
+    "DYLD_FALLBACK_LIBRARY_PATH",
+
+    "SSH_AUTH_SOCK",
+    "SECURITYSESSIONID",
+
+    "__CF_USER_TEXT_ENCODING",
+    "DEVELOPER_DIR",
+
+    "TERM_PROGRAM",
+    "TERM_PROGRAM_VERSION",
+}
+
+UNIVERSAL_ESSENTIAL_VARS = {
+    "PATH",
+    "TEMP", "TMP", "TMPDIR",
+    "LANG", "LC_ALL", "LC_CTYPE",
+}
+
+PROGRAM_ESSENTIAL_VARS = {
+    "PYTHONIOENCODING",
+    "GOPATH",
+    "LUA_PATH",
+    "JAVA_HOME",
+    "RUSTUP_HOME",
+    "CARGO_HOME",
+    "PNPM_HOME",
+    "ANDROID_HOME"
+}
+
+SECURITY_ADDITIONS = {
+    "SSL_CERT_FILE",
+    "SSL_CERT_DIR",
+    "REQUESTS_CA_BUNDLE"
+}
+
+ESSENTIAL_VARS = UNIVERSAL_ESSENTIAL_VARS | PROGRAM_ESSENTIAL_VARS | SECURITY_ADDITIONS
+if platform.system() == "Windows":
+    ESSENTIAL_VARS |= WINDOWS_ESSENTIAL_VARS
+elif platform.system() == "Darwin":
+    ESSENTIAL_VARS |= MACOS_ESSENTIAL_VARS
+else:
+    ESSENTIAL_VARS |= UNIX_ESSENTIAL_VARS
+
+class EnvBuilder:
+    def __init__(self,
+                 blacklist: set[str] | None = None,
+                 extra: dict[str, str] | None = None,
+                 ):
+        self._blacklist = blacklist
+        self._extra = extra
+
+    def with_extra(self, extra: dict[str, str]) -> "EnvBuilder":
+        final_extra = (self._extra or {}).copy()
+        final_extra.update(extra)
+        return EnvBuilder(self._blacklist, final_extra)
+
+    def build(self) -> dict[str, str]:
+        base_env = os.environ.copy()
+        final_env = {}
+
+        # insert essential vars
+        for key, var in base_env.items():
+            if key.upper() in ESSENTIAL_VARS:
+                final_env[key] = var
+
+        # insert extra vars
+        if self._extra:
+            final_env.update(self._extra)
+        return final_env

dais_shell/iostream_reader.py

@@ -0,0 +1,171 @@
+import asyncio
+import subprocess
+import threading
+import time
+from enum import Enum
+from dataclasses import dataclass
+from collections import deque
+from typing import Callable, TextIO
+
+IOStreamCallback = Callable[[str], None]
+IOStreamBuffer = deque[str]
+
+class IOStreamReaderStatus(str, Enum):
+    SUCCESS = "success"
+    TIMEOUT = "timeout"
+    CANCELED = "canceled"
+    ERROR = "error"
+
+@dataclass
+class IOStreamReaderResult:
+    returncode: int
+    status: IOStreamReaderStatus
+    error: Exception | None
+    stdout_buf: IOStreamBuffer
+    stderr_buf: IOStreamBuffer
+
+    @property
+    def stdout(self) -> str:
+        """Get the full text of stdout"""
+        return "\n".join(self.stdout_buf)
+
+    @property
+    def stderr(self) -> str:
+        """Get the full text of stderr"""
+        return "\n".join(self.stderr_buf)
+
+
+class IOStreamReaderSync:
+    def __init__(self,
+                 proc: subprocess.Popen,
+                 on_stdout: IOStreamCallback | None = None,
+                 on_stderr: IOStreamCallback | None = None,
+                 max_lines: int = 10000):
+        self._proc = proc
+        self._max_lines = max_lines
+        self._on_stdout = on_stdout
+        self._on_stderr = on_stderr
+        self._cancel_event = threading.Event()
+
+    def cancel(self):
+        self._cancel_event.set()
+
+    @staticmethod
+    def _consumer(pipe: TextIO, callback: IOStreamCallback | None, buf: IOStreamBuffer):
+        for line in iter(pipe.readline, ""):
+            buf.append(line)
+            if callback: callback(line)
+
+    def read(self, timeout_sec: int | None = None) -> IOStreamReaderResult:
+        if (returncode := self._proc.poll()) is not None:
+            return IOStreamReaderResult(
+                returncode,
+                status=IOStreamReaderStatus.SUCCESS,
+                error=None,
+                stdout_buf=IOStreamBuffer(),
+                stderr_buf=IOStreamBuffer())
+
+        stdout_buf = IOStreamBuffer(maxlen=self._max_lines)
+        stderr_buf = IOStreamBuffer(maxlen=self._max_lines)
+
+        stdout_consumer = threading.Thread(
+            target=IOStreamReaderSync._consumer,
+            args=(self._proc.stdout, self._on_stdout, stdout_buf))
+        stderr_consumer = threading.Thread(
+            target=IOStreamReaderSync._consumer,
+            args=(self._proc.stderr, self._on_stderr, stderr_buf))
+
+        stdout_consumer.start()
+        stderr_consumer.start()
+
+        start_time = time.monotonic()
+        status = IOStreamReaderStatus.SUCCESS
+        error: Exception | None = None
+        try:
+            while self._proc.poll() is None:
+                if self._cancel_event.is_set():
+                    self._proc.kill()
+                    status = IOStreamReaderStatus.CANCELED
+                    break
+
+                if (timeout_sec is not None and
+                   (time.monotonic() - start_time) > timeout_sec):
+                    self._proc.kill()
+                    status = IOStreamReaderStatus.TIMEOUT
+                    break
+                time.sleep(0.1)
+        except Exception as exc:
+            self._proc.kill()
+            status = IOStreamReaderStatus.ERROR
+            error = exc
+
+        stdout_consumer.join(timeout=3)
+        stderr_consumer.join(timeout=3)
+
+        returncode = self._proc.returncode
+        if returncode is None:
+            returncode = self._proc.wait()
+        return IOStreamReaderResult(returncode, status, error, stdout_buf, stderr_buf)
+
+class IOStreamReader:
+    def __init__(self,
+                 proc: asyncio.subprocess.Process,
+                 max_lines: int,
+                 on_stdout: IOStreamCallback | None = None,
+                 on_stderr: IOStreamCallback | None = None,
+                 ):
+        self._proc = proc
+        self._max_lines = max_lines
+        self._on_stdout = on_stdout
+        self._on_stderr = on_stderr
+
+    @staticmethod
+    async def _consumer(stream: asyncio.StreamReader, callback: IOStreamCallback | None, buf: IOStreamBuffer):
+        while not stream.at_eof():
+            line = await stream.readline()
+            if not line: break
+            text = line.decode("utf-8", errors="replace")
+            buf.append(text)
+            if callback: callback(text)
+
+    async def read(self, timeout_sec: int | None = None) -> IOStreamReaderResult:
+        stdout_buf = IOStreamBuffer(maxlen=self._max_lines)
+        stderr_buf = IOStreamBuffer(maxlen=self._max_lines)
+
+        assert self._proc.stdout is not None
+        assert self._proc.stderr is not None
+        consumer_task = [
+            asyncio.create_task(IOStreamReader._consumer(self._proc.stdout, self._on_stdout, stdout_buf)),
+            asyncio.create_task(IOStreamReader._consumer(self._proc.stderr, self._on_stderr, stderr_buf))
+        ]
+
+        status = IOStreamReaderStatus.SUCCESS
+        error: Exception | None = None
+        try:
+            if timeout_sec is not None:
+                returncode = await asyncio.wait_for(self._proc.wait(), timeout=timeout_sec)
+            else:
+                returncode = await self._proc.wait()
+        except asyncio.TimeoutError:
+            self._proc.kill()
+            returncode = await self._proc.wait()
+            status = IOStreamReaderStatus.TIMEOUT
+        except asyncio.CancelledError:
+            self._proc.kill()
+            returncode = await self._proc.wait()
+            status = IOStreamReaderStatus.CANCELED
+        except Exception as exc:
+            self._proc.kill()
+            returncode = await self._proc.wait()
+            status = IOStreamReaderStatus.ERROR
+            error = exc
+        finally:
+            try:
+                await asyncio.wait_for(
+                    asyncio.gather(*consumer_task, return_exceptions=True),
+                    timeout=2)
+            except asyncio.TimeoutError:
+                for task in consumer_task: task.cancel()
+                await asyncio.gather(*consumer_task, return_exceptions=True)
+
+        return IOStreamReaderResult(returncode, status, error, stdout_buf, stderr_buf)

dais_shell/runtimes/BaseShellRuntime.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+from ..types import CommandStep
+from ..iostream_reader import IOStreamReaderResult
+
+class BaseShellRuntime(ABC):
+    @abstractmethod
+    def run_sync(self,
+                 step: CommandStep,
+                 on_stdout=None,
+                 on_stderr=None,
+                 ) -> IOStreamReaderResult: ...
+
+    @abstractmethod
+    async def run(self,
+                  step: CommandStep,
+                  on_stdout=None,
+                  on_stderr=None
+                  ) -> IOStreamReaderResult: ...

dais_shell/runtimes/BashRuntime.py

@@ -0,0 +1,90 @@
+import asyncio
+import shutil
+import subprocess
+from dataclasses import dataclass
+from .BaseShellRuntime import BaseShellRuntime
+from ..types import CommandStep
+from ..types.exceptions import ShellRuntimeNotFoundError
+from ..iostream_reader import IOStreamReaderResult, IOStreamReader, IOStreamReaderSync
+
+@dataclass
+class BashCommandStep(CommandStep):
+    @classmethod
+    def from_command_step(cls, step: CommandStep):
+        return cls(
+            command=step.command,
+            args=step.args,
+            env=step.env,
+            cwd=step.cwd,
+            timeout=step.timeout
+        )
+
+    def to_wrapper_script(self):
+        return 'exec "$1" "${@:2}"'
+
+# --- --- --- --- --- ---
+
+class BashRuntime(BaseShellRuntime):
+    def __init__(self, max_lines: int):
+        self._shell = self._detect_shell()
+        self._max_lines = max_lines
+
+    def _detect_shell(self) -> str:
+        if bash := shutil.which("bash"):
+            return bash
+        if sh := shutil.which("sh"):
+            return sh
+        raise ShellRuntimeNotFoundError("Bash")
+
+    def _make_bash_commands(self, step: BashCommandStep):
+        return [
+            self._shell,
+            "-c",
+            step.to_wrapper_script(),
+            "--",
+            step.command,
+            *step.args
+        ]
+
+    def _prepare_cmd(self, step: CommandStep) -> list[str]:
+        step = BashCommandStep.from_command_step(step)
+        resolved = shutil.which(step.command)
+        is_shell_command = resolved is None
+        if is_shell_command:
+            return self._make_bash_commands(step)
+        else:
+            return [resolved, *step.args]
+
+    def run_sync(self,
+                 step: CommandStep,
+                 on_stdout=None,
+                 on_stderr=None,
+                ) -> IOStreamReaderResult:
+        proc = subprocess.Popen(
+            self._prepare_cmd(step),
+            cwd=step.cwd,
+            env=step.env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+        )
+
+        reader = IOStreamReaderSync(proc, on_stdout, on_stderr, self._max_lines)
+        return reader.read(step.timeout)
+
+    async def run(self,
+                        step: CommandStep,
+                        on_stdout=None,
+                        on_stderr=None
+                        ) -> IOStreamReaderResult:
+        proc = await asyncio.create_subprocess_exec(
+            *self._prepare_cmd(step),
+            cwd=step.cwd,
+            env=step.env,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+
+        reader = IOStreamReader(proc, self._max_lines, on_stdout, on_stderr)
+        return await reader.read(step.timeout)

dais_shell/runtimes/PowershellRuntime.py

@@ -0,0 +1,114 @@
+import asyncio
+import base64
+import os
+import json
+import shutil
+import subprocess
+from dataclasses import dataclass
+from .BaseShellRuntime import BaseShellRuntime
+from ..iostream_reader import IOStreamReaderResult, IOStreamReader, IOStreamReaderSync
+from ..types import CommandStep
+from ..types.exceptions import ShellRuntimeNotFoundError
+
+CREATE_NO_WINDOW = 0x08000000 if os.name == "nt" else 0
+
+@dataclass
+class PowerShellCommandStep(CommandStep):
+    @classmethod
+    def from_command_step(cls, step: CommandStep):
+        return cls(
+            command=step.command,
+            args=step.args,
+            env=step.env,
+            cwd=step.cwd,
+            timeout=step.timeout
+        )
+
+    def to_wrapper_script(self):
+        cmd_json  = json.dumps(self.command)
+        args_json = json.dumps(self.args)
+        script = f"""
+$ErrorActionPreference = "Stop"
+$PSNativeCommandArgumentPassing = "Standard"
+[Console]::OutputEncoding = [System.Text.Encoding]::UTF8
+
+$command  = ConvertFrom-Json '{cmd_json}'
+$arguments = ,(ConvertFrom-Json '{args_json}')
+
+& $command @arguments
+exit $LASTEXITCODE"""
+        return script.strip()
+
+# --- --- --- --- --- ---
+
+class PowerShellRuntime(BaseShellRuntime):
+    def __init__(self, max_lines: int):
+        self._shell = self._detect_shell()
+        self._max_lines = max_lines
+
+    @staticmethod
+    def _detect_shell() -> str:
+        if pwsh := shutil.which("pwsh"):
+            return pwsh
+        if powershell := shutil.which("powershell"):
+            return powershell
+        raise ShellRuntimeNotFoundError("PowerShell")
+
+    @staticmethod
+    def _encode(source: str) -> str:
+        return base64.b64encode(
+            source.encode("utf-16-le")
+        ).decode("ascii")
+
+    def _make_powershell_commands(self, encoded: str):
+        return [
+            self._shell,
+            "-NoProfile",
+            "-NonInteractive",
+            "-ExecutionPolicy", "Bypass",
+            "-EncodedCommand", encoded
+        ]
+
+    def _prepare_cmd(self, step: CommandStep) -> list[str]:
+        step = PowerShellCommandStep.from_command_step(step)
+        script = step.to_wrapper_script()
+        encoded = self._encode(script)
+        return self._make_powershell_commands(encoded)
+
+    def run_sync(
+        self,
+        step: CommandStep,
+        on_stdout=None,
+        on_stderr=None,
+    ) -> IOStreamReaderResult:
+        proc = subprocess.Popen(
+            self._prepare_cmd(step),
+            cwd=step.cwd,
+            env=step.env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+            creationflags=CREATE_NO_WINDOW
+        )
+
+        reader = IOStreamReaderSync(proc, on_stdout, on_stderr, self._max_lines)
+        return reader.read(step.timeout)
+
+    async def run(
+        self,
+        step: CommandStep,
+        on_stdout=None,
+        on_stderr=None
+    ) -> IOStreamReaderResult:
+        proc = await asyncio.create_subprocess_exec(
+            *self._prepare_cmd(step),
+            cwd=step.cwd,
+            env=step.env,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            creationflags=CREATE_NO_WINDOW
+        )
+
+        reader = IOStreamReader(proc, self._max_lines, on_stdout, on_stderr)
+        return await reader.read(step.timeout)

dais_shell/runtimes/__init__.py

@@ -0,0 +1,3 @@
+from .BaseShellRuntime import BaseShellRuntime
+from .BashRuntime import BashRuntime
+from .PowershellRuntime import PowerShellRuntime

dais_shell/types/__init__.py

@@ -0,0 +1,26 @@
+import os
+import shutil
+from abc import abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from .exceptions import CommandNotFoundError, ForbiddenShellTargetError
+
+@dataclass
+class CommandStep:
+    command: str
+    args: list[str]
+    cwd: str | Path
+    env: dict[str, str] | None = None
+    timeout: int | None = None
+
+    @abstractmethod
+    def to_wrapper_script(self) -> str: ...
+
+    def validate_command(self, filter: set[str] | None = None):
+        if filter is not None:
+            name = os.path.basename(self.command).lower()
+            if name in filter:
+                raise ForbiddenShellTargetError(name)
+        resolved = shutil.which(self.command)
+        if resolved is None:
+            raise CommandNotFoundError(self.command)

dais_shell/types/exceptions.py

@@ -0,0 +1,16 @@
+class ShellError(Exception): ...
+
+class CommandNotFoundError(ShellError):
+    def __init__(self, command: str):
+        self.command = command
+        super().__init__(f"Command not found: {command}")
+
+class ShellRuntimeNotFoundError(ShellError):
+    def __init__(self, runtime: str):
+        self.runtime = runtime
+        super().__init__(f"Shell runtime not found: {runtime}")
+
+class ForbiddenShellTargetError(ShellError):
+    def __init__(self, command: str):
+        self.command = command
+        super().__init__(f"Refusing to execute shell program as target: {command}")

dais_shell-0.1.1.dist-info/METADATA

@@ -0,0 +1,12 @@
+Metadata-Version: 2.3
+Name: dais-shell
+Version: 0.1.1
+Summary: The shell tool for Dais.
+Author: BHznJNs
+Author-email: BHznJNs <441768875@qq.com>
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# Dais Shell
+
+The shell tool for [Dais](https://github.com/Dais-Project/Dais).

dais_shell-0.1.1.dist-info/RECORD

@@ -0,0 +1,14 @@
+dais_shell/__init__.py,sha256=GIEykP4yW5bC31JR1km2_fOVJGyxU_kPWgZd02fDobY,2202
+dais_shell/constants.py,sha256=G5_Bo503U7CB9l5b7cjiQe2gNrGsn4HIQFnBokBrpQo,146
+dais_shell/env_builder.py,sha256=ZIJlxZeOhngU1_nD8xJcvfC2CTLYkv7G8tewI7zppJk,2208
+dais_shell/iostream_reader.py,sha256=IZEXgxHseDUyFNr8OGr4EAtUyohHuMdvkGSv_iaWcKg,6044
+dais_shell/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dais_shell/runtimes/BaseShellRuntime.py,sha256=997cIoDK866hbvOhDVLO2IAXUZ3TIUV552VnrQi91w0,540
+dais_shell/runtimes/BashRuntime.py,sha256=UGQa_yl_9I9IkG_lWA1h-ApImONKWuKyFJlaEOhNPyE,2741
+dais_shell/runtimes/PowershellRuntime.py,sha256=8G9OhGJYUu12P1H-Mibjs_14IDy_E8EsfFPxSgQXo-c,3342
+dais_shell/runtimes/__init__.py,sha256=VfiKA1QYRYFpdFJDRzzwYitxKuoK86OadjO8wh4Zd4U,133
+dais_shell/types/__init__.py,sha256=6GFk9q51LQfOnIWwd-lgI9w85SodNLzVXbLiwzVhB-k,775
+dais_shell/types/exceptions.py,sha256=F4x3Q9MTK9V_6XbVRHcfzAUahXvDv6YEO1mVa7uSGhs,579
+dais_shell-0.1.1.dist-info/WHEEL,sha256=iHtWm8nRfs0VRdCYVXocAWFW8ppjHL-uTJkAdZJKOBM,80
+dais_shell-0.1.1.dist-info/METADATA,sha256=bNu9cxvfuPEo6NBbZZ7ZFxbumpj4sbqQS4ncEHkT4Mc,289
+dais_shell-0.1.1.dist-info/RECORD,,

dais_shell-0.1.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.9.30
+Root-Is-Purelib: true
+Tag: py3-none-any