"""
System tools for OllamaTools.
This module provides functions for executing system commands and working with files.
"""
import os
import glob
import subprocess
import logging
import shlex
from typing import List, Optional, Dict, Any, Union, Set
# Import the command validation module
from ollamatools.core.utils.validators.command import (
validate_and_sanitize_command,
CommandValidationError,
get_secure_env,
is_sudo_command
)
# Import the path validation module
from ollamatools.core.utils.validators.path import (
validate_path_for_reading,
validate_path_for_writing,
validate_path_for_searching,
PathValidationError,
configure_path_sandbox
)
# Set up logging
logger = logging.getLogger(__name__)
def execute_command(
command: str,
args: Optional[List[str]] = None,
timeout: int = 60,
permission_manager = None,
custom_allowlist: Optional[set] = None
) -> Dict[str, Any]:
"""
Execute a shell command with arguments.
Args:
command: Command to execute. Can be a simple command or a full command line.
args: Command arguments (used only if command is a simple command)
timeout: Command timeout in seconds
permission_manager: Permission manager instance
custom_allowlist: Optional custom allowlist for commands
Returns:
dict: Result with stdout, stderr, and returncode
"""
# Handle case where command contains both command and arguments as a string
if isinstance(command, str) and ' ' in command and args is None:
# Split the command into parts using shlex to handle quotes properly
try:
command_parts = shlex.split(command)
command = command_parts[0]
args = command_parts[1:] if len(command_parts) > 1 else []
logger.info(f"Split command '{command}' with args {args}")
except ValueError as e:
logger.error(f"Invalid command syntax: {str(e)}")
return {"error": f"Invalid command syntax: {str(e)}"}
# Default args to empty list if None
if args is None:
args = []
# Convert args to strings if they aren't already
args = [str(arg) for arg in args]
# Detect sudo before validation to ensure we prompt correctly
is_sudo = is_sudo_command(command, args)
if is_sudo:
logger.warning(f"Sudo command detected: {command} {' '.join(args)}")
# Check permission if permission manager is provided
if permission_manager:
function_args = {"command": command, "args": args, "timeout": timeout}
permission_granted = permission_manager.request_permission(
function_name="execute_command",
args=function_args,
is_sudo=is_sudo
)
if not permission_granted:
logger.warning(f"Permission denied for command: {command} {' '.join(args)}")
return {"error": "Permission denied"}
try:
# Validate and sanitize the command and arguments
validated_command, validated_args = validate_and_sanitize_command(
command, args, custom_allowlist
)
# Construct the full command
full_command = [validated_command] + validated_args
logger.info(f"Executing validated command: {' '.join(full_command)}")
# Get secure environment variables
secure_env = get_secure_env()
# Execute the command
result = subprocess.run(
full_command,
capture_output=True,
check=False,
text=True,
timeout=timeout,
env=secure_env
)
# Return the result
return {
"stdout": result.stdout,
"stderr": result.stderr,
"returncode": result.returncode,
"command": ' '.join(full_command) # Include the executed command for reference
}
except CommandValidationError as e:
logger.error(f"Command validation error: {str(e)}")
return {
"error": f"Command validation error: {str(e)}",
"command": f"{command} {' '.join(args) if args else ''}"
}
except subprocess.TimeoutExpired:
logger.error(f"Command timed out after {timeout} seconds: {command} {' '.join(args)}")
return {
"error": f"Command timed out after {timeout} seconds",
"command": f"{command} {' '.join(args) if args else ''}"
}
except Exception as e:
logger.error(f"Error executing command: {str(e)}")
return {
"error": str(e),
"command": f"{command} {' '.join(args) if args else ''}"
}
def read_file(
path: str,
max_size: Optional[int] = None,
permission_manager = None,
sandbox_path: Optional[str] = None,
additional_allowed_paths: Optional[Set[str]] = None
) -> str:
"""
Read file contents.
Args:
path: File path
max_size: Maximum file size to read in bytes
permission_manager: Permission manager instance
sandbox_path: Base path that defines the sandbox
additional_allowed_paths: Additional paths to allow
Returns:
str: File contents
"""
# Check permission if permission manager is provided
if permission_manager:
function_args = {"path": path, "max_size": max_size}
permission_granted = permission_manager.request_permission(
function_name="read_file",
args=function_args
)
if not permission_granted:
logger.warning(f"Permission denied for reading file: {path}")
return "Permission denied"
try:
# Validate the path for reading
validated_path = validate_path_for_reading(
path,
sandbox_path=sandbox_path,
additional_allowed_paths=additional_allowed_paths
)
# Check the file size if max_size is specified
if max_size is not None:
file_size = os.path.getsize(validated_path)
if file_size > max_size:
logger.error(f"File size ({file_size} bytes) exceeds maximum size ({max_size} bytes): {validated_path}")
return f"File size ({file_size} bytes) exceeds maximum size ({max_size} bytes)"
# Read the file
with open(validated_path, "r") as f:
content = f.read()
return content
except PathValidationError as e:
logger.error(f"Path validation error: {str(e)}")
return f"Path validation error: {str(e)}"
except Exception as e:
logger.error(f"Error reading file: {str(e)}")
return f"Error reading file: {str(e)}"
def write_file(
path: str,
content: str,
mode: str = 'w',
permission_manager = None,
sandbox_path: Optional[str] = None,
additional_allowed_paths: Optional[Set[str]] = None,
additional_protected_paths: Optional[Set[str]] = None
) -> bool:
"""
Write content to a file.
Args:
path: File path
content: Content to write
mode: Write mode ('w' for write, 'a' for append)
permission_manager: Permission manager instance
sandbox_path: Base path that defines the sandbox
additional_allowed_paths: Additional paths to allow
additional_protected_paths: Additional paths to protect
Returns:
bool: True if successful
"""
# Validate the mode
if mode not in ['w', 'a']:
logger.error(f"Invalid mode: {mode}")
return False
# Check permission if permission manager is provided
if permission_manager:
function_args = {"path": path, "content": content, "mode": mode}
permission_granted = permission_manager.request_permission(
function_name="write_file",
args=function_args
)
if not permission_granted:
logger.warning(f"Permission denied for writing to file: {path}")
return False
try:
# Validate the path for writing
validated_path = validate_path_for_writing(
path,
sandbox_path=sandbox_path,
additional_allowed_paths=additional_allowed_paths,
additional_protected_paths=additional_protected_paths
)
# Create the directory if it doesn't exist
os.makedirs(os.path.dirname(validated_path), exist_ok=True)
# Write to the file
with open(validated_path, mode) as f:
f.write(content)
return True
except PathValidationError as e:
logger.error(f"Path validation error: {str(e)}")
return False
except Exception as e:
logger.error(f"Error writing to file: {str(e)}")
return False
def find_files(
pattern: str,
directory: str = '.',
recursive: bool = True,
permission_manager = None,
sandbox_path: Optional[str] = None,
additional_allowed_paths: Optional[Set[str]] = None
) -> List[str]:
"""
Find files matching a pattern.
Args:
pattern: File pattern
directory: Search directory
recursive: Whether to search recursively
permission_manager: Permission manager instance
sandbox_path: Base path that defines the sandbox
additional_allowed_paths: Additional paths to allow
Returns:
List[str]: Matching file paths
"""
# Check permission if permission manager is provided
if permission_manager:
function_args = {"pattern": pattern, "directory": directory, "recursive": recursive}
permission_granted = permission_manager.request_permission(
function_name="find_files",
args=function_args
)
if not permission_granted:
logger.warning(f"Permission denied for finding files in: {directory}")
return []
try:
# Validate the directory and pattern for searching
validated_dir, validated_pattern = validate_path_for_searching(
directory,
pattern,
sandbox_path=sandbox_path,
additional_allowed_paths=additional_allowed_paths
)
# Construct the search pattern
search_pattern = os.path.join(validated_dir, '**', validated_pattern) if recursive else os.path.join(validated_dir, validated_pattern)
# Find files
files = glob.glob(search_pattern, recursive=recursive)
# Additional validation of found files (check if they're still in the sandbox)
validated_files = []
for file_path in files:
try:
# Only include files that are within the sandbox
abs_path = os.path.abspath(file_path)
if os.path.exists(abs_path) and (
is_path_in_sandbox(
abs_path,
sandbox_path=sandbox_path,
additional_allowed_paths=additional_allowed_paths
)
):
validated_files.append(abs_path)
except Exception as e:
logger.warning(f"Skipping file {file_path}: {str(e)}")
return validated_files
except PathValidationError as e:
logger.error(f"Path validation error: {str(e)}")
return []
except Exception as e:
logger.error(f"Error finding files: {str(e)}")
return []