Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
11 kB · 274 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275"""Runtime guard to prevent user cells from using subprocess/os.system to bypass bash() or native skills."""
from __future__ import annotations
import inspectimport osimport shleximport subprocessfrom typing import Anyfrom pathlib import Path
from .bash import _check_command_guards
_EXEMPT_MODULE_NAMES = { "git", "jj", "search", "nix", "ssh", "ssh_control", "repo", "rlm", "subprocess", "unittest", "unittest.mock", "importlib", "shutil", "tempfile", "multiprocessing", "asyncio",}
_ORIG_SUBPROCESS_RUN = getattr(subprocess, "_prime_orig_run", subprocess.run)_ORIG_SUBPROCESS_POPEN = getattr(subprocess, "_prime_orig_popen", subprocess.Popen)_ORIG_SUBPROCESS_CALL = getattr(subprocess, "_prime_orig_call", subprocess.call)_ORIG_SUBPROCESS_CHECK_CALL = getattr(subprocess, "_prime_orig_check_call", subprocess.check_call)_ORIG_SUBPROCESS_CHECK_OUTPUT = getattr(subprocess, "_prime_orig_check_output", subprocess.check_output)_ORIG_OS_SYSTEM = getattr(os, "_prime_orig_system", os.system)_ORIG_OS_KILL = getattr(os, "_prime_orig_kill", os.kill)_ORIG_OS_KILLPG = getattr(os, "_prime_orig_killpg", os.killpg)
subprocess._prime_orig_run = _ORIG_SUBPROCESS_RUNsubprocess._prime_orig_popen = _ORIG_SUBPROCESS_POPENsubprocess._prime_orig_call = _ORIG_SUBPROCESS_CALLsubprocess._prime_orig_check_call = _ORIG_SUBPROCESS_CHECK_CALLsubprocess._prime_orig_check_output = _ORIG_SUBPROCESS_CHECK_OUTPUTos._prime_orig_system = _ORIG_OS_SYSTEMos._prime_orig_kill = _ORIG_OS_KILLos._prime_orig_killpg = _ORIG_OS_KILLPG
_RLM_PKG_DIR = str(Path(__file__).resolve().parent)
def _is_runtime_file(filename: str) -> bool: if not filename: return False if "/prime-agent-runtime/" in filename: return True try: resolved = str(Path(filename).resolve()) return resolved.startswith(_RLM_PKG_DIR) except Exception: return _RLM_PKG_DIR in filename
def _is_exempt_caller(frame: Any) -> bool: f = frame while f: mod = f.f_globals.get("__name__", "") filename = f.f_code.co_filename or ""
# If we reached a user cell, __main__ (outside runtime), stdin, or string evaluation: if ( (mod == "__main__" and not _is_runtime_file(filename)) or filename.startswith("<cell-") or filename in ("<stdin>", "<string>") ): return False
# If a private helper (like _run) was invoked directly by an interactive cell, it is not exempt if f.f_code.co_name.startswith("_"): caller = f.f_back if caller: c_fn = caller.f_code.co_filename or "" c_mod = caller.f_globals.get("__name__", "") if ( (c_mod == "__main__" and not _is_runtime_file(c_fn)) or c_fn.startswith("<cell-") or c_fn in ("<stdin>", "<string>") ): return False
if mod in _EXEMPT_MODULE_NAMES or (mod and any(mod.startswith(m + ".") for m in _EXEMPT_MODULE_NAMES)): return True if ( "/skills/" in filename or _is_runtime_file(filename) or "/site-packages/" in filename or "/lib/python" in filename ): return True f = f.f_back return True
def _inspect_user_command(cmd_args: Any) -> None: if isinstance(cmd_args, str): # Check command guards (cd, vcs, search, shell wrappers) try: _check_command_guards(cmd_args) except ValueError as e: raise RuntimeError(str(e)) from None raise RuntimeError( f"Do not use subprocess or os.system: run shell commands with await bash({cmd_args!r}) instead." )
if isinstance(cmd_args, (list, tuple)) and cmd_args: raw_first = str(cmd_args[0]).strip().lower() first = Path(raw_first).name.lower() if first in ("rg", "ripgrep", "grep", "egrep", "fgrep", "ag", "ack"): raise RuntimeError( f"Do not use subprocess to run `{first}`: use search('pattern', path='...') or search.files('*.ext') instead." ) if first in ("git", "jj"): from .bash import _format_vcs_error msg = _format_vcs_error(" ".join(str(c) for c in cmd_args), first) msg = msg.replace("bash() rejects", "Do not use subprocess to run") raise RuntimeError(msg) if first in ("nix-shell", "nix"): if first == "nix-shell": raise RuntimeError("Do not use subprocess to run `nix-shell`: use `await nix.develop(...)` or `await nix.dev_env(...)` instead.") subcmd = None for c in cmd_args[1:]: if not str(c).startswith("-"): subcmd = str(c).lower() break if subcmd == "develop": raise RuntimeError("Do not use subprocess to run `nix develop`: use `shell = await nix.develop(...)` and `shell.run(...)` instead.") if subcmd == "fmt": raise RuntimeError("Do not use subprocess to run `nix fmt`: use `nix.fmt(...)` instead.") if subcmd == "build": raise RuntimeError("Do not use subprocess to run `nix build`: use `nix.build(...)` instead.") if subcmd == "run": raise RuntimeError("Do not use subprocess to run `nix run`: use `nix.run(...)` instead.") raise RuntimeError("Do not use subprocess to run `nix`: use the Python `nix` skill instead.") if first in ("find", "fd"): raise RuntimeError( f"Do not use subprocess to run `{first}`: use search.files('pattern', path='...') or Path.glob() instead." ) if first in ("ls", "dir"): raise RuntimeError( f"Do not use subprocess to run `{first}`: use search.ls('path', pattern='...') or os.listdir(...) instead." ) if first in ("bash", "sh", "zsh", "dash", "ksh", "nu"): # Check if -c or --command was passed inner_cmd = None for i, tok in enumerate(cmd_args[1:], 1): tok_str = str(tok) if tok_str in ("-c", "--command", "--commands") or (tok_str.startswith("-") and "c" in tok_str and not tok_str.startswith("--")): if i + 1 < len(cmd_args): inner_cmd = str(cmd_args[i + 1]) break if inner_cmd: try: _check_command_guards(inner_cmd) except ValueError as e: raise RuntimeError(str(e)) from None raise RuntimeError( f"Do not use subprocess to run shell scripts ({first}): use await bash(...) instead." ) raise RuntimeError( f"Do not use subprocess to run '{first}': use await bash(...) instead." )
def _guarded_run(*args: Any, **kwargs: Any) -> Any: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): cmd = args[0] if args else kwargs.get("args") _inspect_user_command(cmd) return _ORIG_SUBPROCESS_RUN(*args, **kwargs)
class _GuardedPopen(_ORIG_SUBPROCESS_POPEN): # type: ignore[type-arg,misc] """Popen wrapper that keeps Popen a class.
Libraries (e.g. mcp.os.win32.utilities) evaluate `subprocess.Popen[bytes]` in signatures at import time, which requires Popen to remain a class. """
def __init__(self, *args: Any, **kwargs: Any) -> None: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): cmd = args[0] if args else kwargs.get("args") _inspect_user_command(cmd) super().__init__(*args, **kwargs)
def _guarded_call(*args: Any, **kwargs: Any) -> Any: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): cmd = args[0] if args else kwargs.get("args") _inspect_user_command(cmd) return _ORIG_SUBPROCESS_CALL(*args, **kwargs)
def _guarded_check_call(*args: Any, **kwargs: Any) -> Any: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): cmd = args[0] if args else kwargs.get("args") _inspect_user_command(cmd) return _ORIG_SUBPROCESS_CHECK_CALL(*args, **kwargs)
def _guarded_check_output(*args: Any, **kwargs: Any) -> Any: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): cmd = args[0] if args else kwargs.get("args") _inspect_user_command(cmd) return _ORIG_SUBPROCESS_CHECK_OUTPUT(*args, **kwargs)
def _guarded_os_system(command: str) -> int: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): _inspect_user_command(command) return _ORIG_OS_SYSTEM(command)
def _guarded_os_kill(pid: int, sig: int) -> None: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): from .host_safety import is_pid_owned if pid != os.getpid() and not is_pid_owned(pid): raise PermissionError( f"os.kill rejects pid {pid}: process is not owned by this session. " "Scoped process control only permits signalling processes started by bash() in this session " "or explicitly adopted via `rlm.adopt_process(pid, reason=...)`. " "Use `kill_process(pid)` or `h.kill()` for an owned process." ) return _ORIG_OS_KILL(pid, sig)
def _guarded_os_killpg(pgid: int, sig: int) -> None: frame = inspect.currentframe() caller = frame.f_back if frame else None if caller and not _is_exempt_caller(caller): from .host_safety import is_pid_owned if pgid != os.getpgrp() and not is_pid_owned(pgid): raise PermissionError( f"os.killpg rejects pgid {pgid}: process group is not owned by this session. " "Scoped process control only permits signalling processes started by bash() in this session " "or explicitly adopted via `rlm.adopt_process(pid, reason=...)`. " "Use `kill_process(pid)` or `h.kill()` for an owned process." ) return _ORIG_OS_KILLPG(pgid, sig)
def install_subprocess_guards() -> None: subprocess.run = _guarded_run subprocess.Popen = _GuardedPopen subprocess.call = _guarded_call subprocess.check_call = _guarded_check_call subprocess.check_output = _guarded_check_output os.system = _guarded_os_system os.kill = _guarded_os_kill os.killpg = _guarded_os_killpg