Anton
052591281d
- Complete agent system with planning and execution - Bridge API with FastAPI server - WhatsApp-style mobile dashboard - Enhanced web search with query improvement - Security guards and policy checks - Comprehensive documentation
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
"""Safe shell execution tool with allowlist enforcement."""
|
|
from __future__ import annotations
|
|
|
|
import subprocess
|
|
|
|
from core.state import PlanStep, StepResult, StepStatus
|
|
|
|
# Commands explicitly allowed
|
|
ALLOWED_COMMANDS: frozenset[str] = frozenset({"pwd", "ls", "echo", "cat", "python", "git"})
|
|
|
|
# Commands explicitly blocked (belt-and-suspenders)
|
|
BLOCKED_COMMANDS: frozenset[str] = frozenset({
|
|
"bash", "sh", "rm", "sudo", "curl", "wget",
|
|
"chmod", "chown", "docker", "systemctl",
|
|
})
|
|
|
|
# For commands with restricted sub-commands
|
|
RESTRICTED_SUBCOMMANDS: dict[str, set[str]] = {
|
|
"python": {"--version"},
|
|
"git": {"status"},
|
|
}
|
|
|
|
|
|
def run_shell(step: PlanStep) -> StepResult:
|
|
"""Execute a whitelisted shell command via subprocess (shell=False)."""
|
|
raw_command: str = step.payload.get("command", "").strip()
|
|
parts = raw_command.split()
|
|
|
|
if not parts:
|
|
return StepResult(step_id=step.step_id, status=StepStatus.FAILED, error="Empty command")
|
|
|
|
binary = parts[0]
|
|
|
|
if binary in BLOCKED_COMMANDS:
|
|
return StepResult(
|
|
step_id=step.step_id,
|
|
status=StepStatus.FAILED,
|
|
error=f"Command '{binary}' is blocked by security policy",
|
|
)
|
|
|
|
if binary not in ALLOWED_COMMANDS:
|
|
return StepResult(
|
|
step_id=step.step_id,
|
|
status=StepStatus.FAILED,
|
|
error=f"Command '{binary}' is not in allowlist. Allowed: {sorted(ALLOWED_COMMANDS)}",
|
|
)
|
|
|
|
# Validate restricted sub-commands
|
|
if binary in RESTRICTED_SUBCOMMANDS:
|
|
allowed_sub = RESTRICTED_SUBCOMMANDS[binary]
|
|
sub = " ".join(parts[1:])
|
|
if sub not in allowed_sub:
|
|
return StepResult(
|
|
step_id=step.step_id,
|
|
status=StepStatus.FAILED,
|
|
error=f"Only '{binary} {sorted(allowed_sub)}' is permitted, got: '{binary} {sub}'",
|
|
)
|
|
|
|
try:
|
|
proc = subprocess.run(
|
|
parts,
|
|
shell=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=10,
|
|
)
|
|
output = proc.stdout + proc.stderr
|
|
if proc.returncode != 0:
|
|
return StepResult(
|
|
step_id=step.step_id,
|
|
status=StepStatus.FAILED,
|
|
output_text=output,
|
|
error=f"Non-zero exit code: {proc.returncode}",
|
|
)
|
|
return StepResult(step_id=step.step_id, status=StepStatus.DONE, output_text=output)
|
|
|
|
except subprocess.TimeoutExpired:
|
|
return StepResult(
|
|
step_id=step.step_id,
|
|
status=StepStatus.FAILED,
|
|
error="Command timed out after 10s",
|
|
)
|
|
except FileNotFoundError:
|
|
return StepResult(
|
|
step_id=step.step_id,
|
|
status=StepStatus.FAILED,
|
|
error=f"Binary '{binary}' not found on PATH",
|
|
)
|