dais-shell 0.2.0__tar.gz → 0.3.1__tar.gz

{dais_shell-0.2.0 → dais_shell-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: dais-shell
-Version: 0.2.0
+Version: 0.3.1
 Summary: The shell tool for Dais.
 Author: BHznJNs
 Author-email: BHznJNs <441768875@qq.com>

{dais_shell-0.2.0 → dais_shell-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dais-shell"
-version = "0.2.0"
+version = "0.3.1"
 description = "The shell tool for Dais."
 readme = "README.md"
 authors = [{ name = "BHznJNs", email = "441768875@qq.com" }]

{dais_shell-0.2.0 → dais_shell-0.3.1}/src/dais_shell/__init__.py

@@ -5,7 +5,9 @@ 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, ShellError, ShellRuntimeNotFoundError, ForbiddenShellTargetError
+from .types.command_step import CommandStep
+from .types.shell_script import ShellScript
+from .types.exceptions import ShellError, ShellRuntimeNotFoundError, ForbiddenShellTargetError
 from .constants import DEFAULT_COMMAND_BLACKLIST
 
 ShellResult: TypeAlias = IOStreamReaderResult
@@ -30,19 +32,20 @@ class AgentShell:
             return BashRuntime(max_lines)
 
     def run_sync(self,
-                 step: CommandStep,
+                 step: CommandStep | ShellScript,
                  on_stdout=None,
                  on_stderr=None
                  ) -> ShellResult:
         return asyncio.run(self.run(step, on_stdout, on_stderr))
 
     async def run(self,
-                  step: CommandStep,
+                  step: CommandStep | ShellScript,
                   on_stdout=None,
                   on_stderr=None
                   ) -> ShellResult:
         step = replace(step)
-        step.validate_forbidden(self._command_blacklist)
+        if isinstance(step, CommandStep):
+            step.validate_forbidden(self._command_blacklist)
         step.env = (self._env_builder
                         .with_extra(step.env or {})
                         .build())
@@ -51,6 +54,7 @@ class AgentShell:
 __all__ = [
     "AgentShell",
     "CommandStep",
+    "ShellScript",
     "ShellResult",
     "ShellResultStatus",
 

{dais_shell-0.2.0 → dais_shell-0.3.1}/src/dais_shell/runtimes/BaseShellRuntime.py

@@ -1,19 +1,20 @@
 from abc import ABC, abstractmethod
-from ..types import CommandStep
+from ..types.command_step import CommandStep
+from ..types.shell_script import ShellScript
 from ..iostream_reader import IOStreamReaderResult
 
 
 class BaseShellRuntime(ABC):
     @abstractmethod
     def run_sync(self,
-                 step: CommandStep,
+                 step: CommandStep | ShellScript,
                  on_stdout=None,
                  on_stderr=None,
                  ) -> IOStreamReaderResult: ...
 
     @abstractmethod
     async def run(self,
-                  step: CommandStep,
+                  step: CommandStep | ShellScript,
                   on_stdout=None,
                   on_stderr=None
                   ) -> IOStreamReaderResult: ...

{dais_shell-0.2.0 → dais_shell-0.3.1}/src/dais_shell/runtimes/BashRuntime.py

@@ -5,7 +5,9 @@ from dataclasses import asdict, dataclass
 
 from dais_shell.utils.env_expander import EnvExpander
 from .BaseShellRuntime import BaseShellRuntime
-from ..types import CommandStep, ShellRuntimeNotFoundError
+from ..types.command_step import CommandStep
+from ..types.shell_script import ShellScript
+from ..types.exceptions import ShellRuntimeNotFoundError
 from ..iostream_reader import IOStreamReader, IOStreamReaderResult
 
 
@@ -39,21 +41,24 @@ class BashRuntime(BaseShellRuntime):
             step.to_wrapper_script(),
         ]
 
-    def _prepare_cmd(self, step: CommandStep) -> list[str]:
+    def _prepare_cmd(self, step: CommandStep | ShellScript) -> list[str]:
+        if isinstance(step, ShellScript):
+            return [self._shell, "-c", step.script]
+
         env_expander = EnvExpander(step.env or {})
         step.args = env_expander.expand(step.args)
         step = BashCommandStep.from_command_step(step)
         return self._make_bash_commands(step)
 
     def run_sync(self,
-                 step: CommandStep,
+                 step: CommandStep | ShellScript,
                  on_stdout=None,
                  on_stderr=None,
                 ) -> IOStreamReaderResult:
         return asyncio.run(self.run(step, on_stdout, on_stderr))
 
     async def run(self,
-                        step: CommandStep,
+                        step: CommandStep | ShellScript,
                         on_stdout=None,
                         on_stderr=None
                         ) -> IOStreamReaderResult:

{dais_shell-0.2.0 → dais_shell-0.3.1}/src/dais_shell/runtimes/PowershellRuntime.py

@@ -1,16 +1,62 @@
 import asyncio
 import base64
-import json
 import shutil
 import re
 import xml.etree.ElementTree as ET
 import subprocess
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, fields
 
 from dais_shell.utils.env_expander import EnvExpander
 from .BaseShellRuntime import BaseShellRuntime
 from ..iostream_reader import IOStreamReader, IOStreamReaderResult
-from ..types import CommandStep, ShellRuntimeNotFoundError
+from ..types.command_step import CommandStep
+from ..types.shell_script import ShellScript
+from ..types.exceptions import ShellRuntimeNotFoundError
+
+
+POWERSHELL_PREFERENCES = """
+$ErrorActionPreference = "Stop"
+$WarningPreference = "SilentlyContinue"
+$VerbosePreference = "SilentlyContinue"
+$DebugPreference = "SilentlyContinue"
+$InformationPreference = "SilentlyContinue"
+$ProgressPreference = "SilentlyContinue"
+
+$PSNativeCommandArgumentPassing = "Standard"
+""".strip()
+
+@dataclass
+class PowerShellScript(ShellScript):
+    @classmethod
+    def from_shell_script(cls, script: ShellScript):
+        data = asdict(script)
+        return cls(**{
+            field.name: data[field.name]
+            for field in fields(ShellScript)
+        })
+
+    def to_wrapper_script(self):
+        return f"""
+{POWERSHELL_PREFERENCES}
+
+chcp 65001 | Out-Null
+$OutputEncoding           = [System.Text.Encoding]::UTF8
+[Console]::OutputEncoding = [System.Text.Encoding]::UTF8
+[Console]::InputEncoding  = [System.Text.Encoding]::UTF8
+
+$LASTEXITCODE = $null
+
+& {{
+{self.script}
+
+$daisSuccess = $?
+$daisExitCode = $LASTEXITCODE
+
+if ($daisSuccess) {{ exit 0 }}
+if ($daisExitCode -ne $null) {{ exit $daisExitCode }}
+exit 1
+}}
+""".strip()
 
 
 @dataclass
@@ -24,8 +70,7 @@ class PowerShellCommandStep(CommandStep):
             return "'" + s.replace("'", "''") + "'"
 
         script = f"""
-$ErrorActionPreference = "Stop"
-$PSNativeCommandArgumentPassing = "Standard"
+{POWERSHELL_PREFERENCES}
 
 chcp 65001 | Out-Null
 $OutputEncoding           = [System.Text.Encoding]::UTF8
@@ -102,7 +147,13 @@ class PowerShellRuntime(BaseShellRuntime):
             "-EncodedCommand", encoded
         ]
 
-    def _prepare_cmd(self, step: CommandStep) -> list[str]:
+    def _prepare_cmd(self, step: CommandStep | ShellScript) -> list[str]:
+        if isinstance(step, ShellScript):
+            script = PowerShellScript.from_shell_script(step)
+            return self._make_powershell_commands(
+                self._encode(script.to_wrapper_script())
+            )
+
         env_expander = EnvExpander(step.env or {})
         step.args = env_expander.expand(step.args)
         step = PowerShellCommandStep.from_command_step(step)
@@ -112,7 +163,7 @@ class PowerShellRuntime(BaseShellRuntime):
 
     def run_sync(
         self,
-        step: CommandStep,
+        step: CommandStep | ShellScript,
         on_stdout=None,
         on_stderr=None,
     ) -> IOStreamReaderResult:
@@ -120,7 +171,7 @@ class PowerShellRuntime(BaseShellRuntime):
 
     async def run(
         self,
-        step: CommandStep,
+        step: CommandStep | ShellScript,
         on_stdout=None,
         on_stderr=None
     ) -> IOStreamReaderResult:

{dais_shell-0.2.0 → dais_shell-0.3.1}/src/dais_shell/types/__init__.py

@@ -1,2 +1,3 @@
 from .command_step import *
+from .shell_script import *
 from .exceptions import *

{dais_shell-0.2.0 → dais_shell-0.3.1}/src/dais_shell/types/command_step.py

@@ -22,6 +22,7 @@ class CommandStep:
             if name in filter:
                 raise ForbiddenShellTargetError(name)
 
+
 __all__ = [
     "CommandStep",
 ]

dais_shell-0.3.1/src/dais_shell/types/shell_script.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class ShellScript:
+    """Trusted shell source executed directly without command blacklist validation."""
+
+    script: str
+    cwd: str | Path
+    env: dict[str, str] | None = None
+    timeout: int | None = None
+
+
+__all__ = [
+    "ShellScript",
+]