May / tools /shell_tools.py
Mayank2027's picture
Create tools/shell_tools.py
337801b verified
Raw
History Blame Contribute Delete
14.6 kB
# File 7: tools/shell_tools.py
with open(f"{base_dir}/tools/shell_tools.py", "w") as f:
f.write('''"""Shell and Python Execution Tools"""
import os
import sys
import subprocess
import tempfile
import json
from typing import Dict, List, Optional
from config import Config
from utils import UI, Logger, truncate
class ShellTools:
"""Shell command and Python code execution with safety controls"""
def __init__(self):
self.logger = Logger()
self.blocked_patterns = [
"rm -rf /", "rm -rf /*", ":(){ :|:& };:", "> /dev/sda",
"mkfs", "dd if=/dev/zero", "chmod -R 777 /",
"wget -O- | sh", "curl | sh", "curl | bash"
]
def _is_safe(self, command: str) -> tuple:
"""Check if command is safe to run"""
cmd_lower = command.lower().strip()
for pattern in self.blocked_patterns:
if pattern.lower() in cmd_lower:
return False, f"Blocked dangerous pattern: {pattern}"
return True, "OK"
def bash_command(self, command: str, cwd: str = "",
timeout: int = 30, env_vars: Dict[str, str] = None,
capture_output: bool = True, confirm: bool = True) -> Dict:
"""Execute shell command with safety controls"""
try:
safe, reason = self._is_safe(command)
if not safe:
return {"success": False, "error": reason}
if cwd:
work_dir = os.path.join(Config.SANDBOX_PATH, cwd)
else:
work_dir = Config.SANDBOX_PATH
os.makedirs(work_dir, exist_ok=True)
destructive_keywords = ["rm", "del", "remove", "delete", "drop", "truncate",
"format", "mkfs", "dd", ">/dev"]
needs_confirm = any(kw in command.lower() for kw in destructive_keywords)
if needs_confirm and confirm:
if not UI.confirm(f"Run destructive command: {command}?"):
return {"success": False, "error": "User cancelled"}
env = os.environ.copy()
if env_vars:
env.update(env_vars)
result = subprocess.run(
command,
shell=True,
cwd=work_dir,
capture_output=capture_output,
text=True,
timeout=timeout,
env=env
)
self.logger.log("bash_command", {
"command": command,
"returncode": result.returncode,
"cwd": cwd
})
return {
"success": result.returncode == 0,
"command": command,
"returncode": result.returncode,
"stdout": truncate(result.stdout) if capture_output else "",
"stderr": truncate(result.stderr) if capture_output else "",
"cwd": cwd
}
except subprocess.TimeoutExpired:
return {"success": False, "error": f"Command timed out after {timeout}s"}
except Exception as e:
return {"success": False, "error": str(e)}
def python_execute(self, code: str, timeout: int = 30,
packages: List[str] = None, save_output: bool = False) -> Dict:
"""Execute Python code in isolated environment"""
try:
if packages:
for pkg in packages:
subprocess.run(
[sys.executable, "-m", "pip", "install", "-q", pkg],
capture_output=True, timeout=60
)
with tempfile.NamedTemporaryFile(mode="w", suffix=".py",
dir=Config.SANDBOX_PATH,
delete=False) as f:
f.write(code)
temp_file = f.name
env = os.environ.copy()
env["PYTHONPATH"] = Config.SANDBOX_PATH
result = subprocess.run(
[sys.executable, temp_file],
capture_output=True,
text=True,
timeout=timeout,
cwd=Config.SANDBOX_PATH,
env=env
)
try:
os.remove(temp_file)
except:
pass
output = result.stdout
if result.stderr:
output += "\n[STDERR]\n" + result.stderr
output_file = ""
if save_output and output:
output_file = f"python_output_{os.path.basename(temp_file)}.txt"
output_path = os.path.join(Config.SANDBOX_PATH, output_file)
with open(output_path, "w") as f:
f.write(output)
self.logger.log("python_execute", {
"code_length": len(code),
"returncode": result.returncode
})
return {
"success": result.returncode == 0,
"output": truncate(output, 2000),
"returncode": result.returncode,
"output_file": output_file if save_output else "",
"code_executed": code[:200] + "..." if len(code) > 200 else code
}
except subprocess.TimeoutExpired:
return {"success": False, "error": f"Code execution timed out after {timeout}s"}
except Exception as e:
return {"success": False, "error": str(e)}
def process_manager(self, action: str = "list", pid: int = None,
command: str = "", name_filter: str = "") -> Dict:
"""Manage system processes"""
try:
if action == "list":
if os.name == "nt":
cmd = "tasklist /fo csv"
else:
cmd = "ps aux"
result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
lines = result.stdout.strip().split("\n")
processes = []
for line in lines[1:]:
parts = line.split()
if len(parts) >= 2:
proc_name = " ".join(parts[10:]) if len(parts) > 10 else parts[-1]
if name_filter and name_filter.lower() not in proc_name.lower():
continue
processes.append({
"pid": parts[1] if len(parts) > 1 else "",
"cpu": parts[2] if len(parts) > 2 else "",
"mem": parts[3] if len(parts) > 3 else "",
"command": proc_name[:50]
})
return {
"success": True,
"action": "list",
"processes": processes[:50],
"total": len(processes)
}
elif action == "kill":
if not pid:
return {"success": False, "error": "PID required for kill action"}
if not UI.confirm(f"Kill process {pid}?"):
return {"success": False, "error": "User cancelled"}
if os.name == "nt":
result = subprocess.run(f"taskkill /PID {pid} /F", shell=True, capture_output=True, text=True)
else:
result = subprocess.run(f"kill -9 {pid}", shell=True, capture_output=True, text=True)
return {
"success": result.returncode == 0,
"action": "kill",
"pid": pid,
"output": result.stdout or result.stderr
}
elif action == "start":
if not command:
return {"success": False, "error": "Command required for start action"}
process = subprocess.Popen(
command,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
cwd=Config.SANDBOX_PATH
)
return {
"success": True,
"action": "start",
"pid": process.pid,
"command": command
}
else:
return {"success": False, "error": f"Unknown action: {action}"}
except Exception as e:
return {"success": False, "error": str(e)}
def system_info(self) -> Dict:
"""Get comprehensive system information"""
try:
info = {
"success": True,
"platform": sys.platform,
"python_version": sys.version,
"cpu_count": os.cpu_count(),
"cwd": os.getcwd(),
"sandbox": Config.SANDBOX_PATH
}
if hasattr(os, "statvfs"):
stat = os.statvfs(Config.SANDBOX_PATH)
info["disk_total"] = f"{stat.f_frsize * stat.f_blocks / (1024**3):.2f} GB"
info["disk_free"] = f"{stat.f_frsize * stat.f_bavail / (1024**3):.2f} GB"
try:
with open("/proc/meminfo", "r") as f:
meminfo = f.read()
for line in meminfo.split("\n"):
if "MemTotal" in line:
info["memory_total"] = line.split(":")[1].strip()
elif "MemAvailable" in line:
info["memory_available"] = line.split(":")[1].strip()
except:
pass
self.logger.log("system_info", info)
return info
except Exception as e:
return {"success": False, "error": str(e)}
def environment_vars(self, action: str = "list", key: str = "", value: str = "") -> Dict:
"""Manage environment variables"""
try:
if action == "list":
vars_dict = dict(os.environ)
sensitive = ["key", "token", "secret", "password", "api"]
filtered = {k: "***" if any(s in k.lower() for s in sensitive) else v
for k, v in vars_dict.items()}
return {
"success": True,
"action": "list",
"variables": filtered,
"count": len(filtered)
}
elif action == "get":
if not key:
return {"success": False, "error": "Key required"}
return {
"success": True,
"action": "get",
"key": key,
"value": os.environ.get(key, "(not set)")
}
elif action == "set":
if not key:
return {"success": False, "error": "Key required"}
os.environ[key] = value
return {
"success": True,
"action": "set",
"key": key,
"value": value
}
elif action == "delete":
if not key:
return {"success": False, "error": "Key required"}
if key in os.environ:
del os.environ[key]
return {
"success": True,
"action": "delete",
"key": key
}
else:
return {"success": False, "error": f"Unknown action: {action}"}
except Exception as e:
return {"success": False, "error": str(e)}
def package_manager(self, action: str = "list", package: str = "",
manager: str = "pip") -> Dict:
"""Manage packages"""
try:
commands = {
"pip": {
"list": "pip list --format=json",
"install": f"pip install {package}",
"uninstall": f"pip uninstall -y {package}",
"update": f"pip install --upgrade {package}",
"search": f"pip search {package} 2>/dev/null || pip index versions {package} 2>/dev/null || echo \'Search not available\'"
},
"npm": {
"list": "npm list --depth=0",
"install": f"npm install {package}",
"uninstall": f"npm uninstall {package}",
"update": f"npm update {package}",
"search": f"npm search {package}"
}
}
if manager not in commands:
return {"success": False, "error": f"Unsupported package manager: {manager}"}
cmd = commands[manager].get(action)
if not cmd:
return {"success": False, "error": f"Unknown action: {action}"}
if action in ["install", "uninstall"] and not UI.confirm(f"{action} {package} via {manager}?"):
return {"success": False, "error": "User cancelled"}
result = subprocess.run(
cmd, shell=True, capture_output=True, text=True, timeout=120
)
output = result.stdout
if action == "list" and manager == "pip":
try:
packages = json.loads(output)
output = packages
except:
pass
return {
"success": result.returncode == 0,
"action": action,
"manager": manager,
"package": package,
"output": output[:2000] if isinstance(output, str) else output,
"stderr": result.stderr[:500] if result.stderr else ""
}
except Exception as e:
return {"success": False, "error": str(e)}
''')
print("tools/shell_tools.py done")