"""
Permission management system for OllamaTools.
This module provides security controls for executing functions,
with a focus on a prompt-based permission system.
"""
import os
from typing import Dict, Any, Optional, Tuple, Set
import logging
import re
from ollamatools.core.errors import PermissionError, ERROR_CODES
# Set up logging
logger = logging.getLogger(__name__)
class PermissionManager:
"""
Manages security permissions for function execution.
The PermissionManager controls access to potentially sensitive
operations such as executing system commands or making web requests.
It provides a prompt-based approach to security, allowing users to
approve or deny operations at runtime.
Attributes:
prompt_enabled (bool): Whether to prompt for permission before operations
sudo_always_prompt (bool): Whether to always prompt for sudo commands
session_permissions (Set[str]): Set of operations allowed for the session
"""
def __init__(
self,
prompt_enabled: bool = True,
sudo_always_prompt: bool = True
):
"""
Initialize the permission manager.
Args:
prompt_enabled: Whether to prompt for permissions
sudo_always_prompt: Whether to always prompt for sudo commands
"""
self.prompt_enabled = prompt_enabled
self.sudo_always_prompt = sudo_always_prompt
self.session_permissions = set() # Store allowed operations for the session
self.session_permission_patterns = {
"read_file": {}, # Directory-based permissions
"write_file": {}, # Directory-based permissions
"execute_command": {} # Command-based permissions
}
logger.debug("PermissionManager initialized with prompt_enabled=%s, sudo_always_prompt=%s",
prompt_enabled, sudo_always_prompt)
def _generate_permission_key(self, function_name: str, args: Dict[str, Any]) -> str:
"""
Generate a more flexible permission key that recognizes similar operations.
Args:
function_name: Name of the function
args: Function arguments
Returns:
str: Permission key
"""
# For commands, focus on the command name rather than all arguments
if function_name == "execute_command" and "command" in args:
command = args["command"]
# Just use the base command without all arguments for matching
base_command = command.split()[0] if isinstance(command, str) and " " in command else command
return f"{function_name}:command={base_command}"
# For file operations, focus on the operation type and parent directory
if function_name in ("read_file", "write_file") and "path" in args:
path = args["path"]
# Use parent directory for path-based operations to match similar files
parent_dir = os.path.dirname(os.path.abspath(path))
return f"{function_name}:dir={parent_dir}"
# For other operations, use the full args as before
args_str = ', '.join(f"{k}={v!r}" for k, v in args.items())
return f"{function_name}:{args_str}"
def request_permission(
self,
function_name: str,
args: Dict[str, Any],
is_sudo: bool = False
) -> bool:
"""
Request permission to execute a function.
This method handles the permission request workflow, which may involve
prompting the user for approval depending on configuration settings.
Args:
function_name: Name of the function
args: Function arguments
is_sudo: Whether the command uses sudo
Returns:
bool: True if permission granted, False otherwise
Raises:
PermissionError: If there's an issue with the permission request
"""
# Format the arguments as a readable string
args_str = ', '.join(f"{k}={v!r}" for k, v in args.items())
# Generate a permission key for the operation
permission_key = self._generate_permission_key(function_name, args)
# Check if we need to prompt for permission
should_prompt = self.prompt_enabled
# Always prompt for sudo if sudo_always_prompt is True
if is_sudo and self.sudo_always_prompt:
should_prompt = True
logger.info("Forcing permission prompt for sudo command")
# Check if this operation is already allowed for the session
if permission_key in self.session_permissions:
logger.debug("Permission already granted for session: %s", permission_key)
return True
# Check specific session permission patterns
if function_name == "execute_command" and "command" in args:
command = args["command"]
base_command = command.split()[0] if isinstance(command, str) and " " in command else command
if base_command in self.session_permission_patterns["execute_command"]:
logger.debug("Permission already granted for command type: %s", base_command)
return True
elif function_name in ("read_file", "write_file") and "path" in args:
path = args["path"]
parent_dir = os.path.dirname(os.path.abspath(path))
if parent_dir in self.session_permission_patterns[function_name]:
logger.debug("Permission already granted for directory: %s", parent_dir)
return True
# If prompting is disabled and not overridden by sudo, grant permission
if not should_prompt:
logger.debug("Permission prompts disabled, automatically granting permission for: %s", permission_key)
return True
# Display the permission prompt
print("\n" + "="*60)
print("PERMISSION REQUEST")
print("="*60)
print(f"Function: {function_name}")
print(f"Arguments: {args_str}")
if is_sudo:
print("\nWARNING: This operation requires elevated permissions (sudo)")
print(" Be cautious when granting permission for sudo operations")
print("\nOptions:")
print(" [A] Allow: Grant permission for this operation")
print(" [S] Allow for Session: Grant permission for similar operations in this session")
print(" [M] Modify: Modify the operation before execution")
print(" [D] Deny: Reject the operation")
while True:
try:
response = input("\nYour choice [A/S/M/D]: ").strip().upper()
if response == 'A': # Allow once
logger.info("Permission granted for: %s", permission_key)
return True
elif response == 'S': # Allow for session
self.session_permissions.add(permission_key)
# Store more specific pattern matches for better session handling
if function_name == "execute_command" and "command" in args:
command = args["command"]
base_command = command.split()[0] if isinstance(command, str) and " " in command else command
self.session_permission_patterns["execute_command"][base_command] = True
logger.info(f"Command '{base_command}' allowed for session")
elif function_name in ("read_file", "write_file") and "path" in args:
path = args["path"]
parent_dir = os.path.dirname(os.path.abspath(path))
self.session_permission_patterns[function_name][parent_dir] = True
logger.info(f"{function_name} allowed for directory: {parent_dir}")
logger.info("Permission granted for session: %s", permission_key)
return True
elif response == 'M': # Modify
modified_args = self.modify_command(function_name, args)
if modified_args is not None:
# Generate a new permission key for the modified args
modified_args_str = ', '.join(f"{k}={v!r}" for k, v in modified_args.items())
modified_key = f"{function_name}:{modified_args_str}"
# Add to session permissions if it's a modification of an existing permission
if permission_key in self.session_permissions:
self.session_permissions.add(modified_key)
# Check if sudo status has changed
new_is_sudo = False
if function_name == "execute_command" and "command" in modified_args:
new_is_sudo = self.detect_sudo(modified_args["command"])
# Call request_permission again with modified arguments
logger.info("Requesting permission for modified operation: %s", modified_key)
return self.request_permission(
function_name,
modified_args,
is_sudo=new_is_sudo
)
else:
logger.info("Modification cancelled for: %s", permission_key)
print("Modification cancelled.")
return False
elif response == 'D': # Deny
logger.info("Permission denied for: %s", permission_key)
print("Permission denied.")
return False
else:
print("Invalid choice. Please enter A, S, M, or D.")
except KeyboardInterrupt:
logger.warning("Permission request interrupted")
print("\nOperation cancelled.")
raise PermissionError(
"Permission request interrupted by user",
function_name=function_name,
args=args,
code=ERROR_CODES.get('PERMISSION_PROMPT_CANCELLED')
)
def modify_command(
self,
function_name: str,
args: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
"""
Allow the user to modify command arguments.
Args:
function_name: Name of the function
args: Current function arguments
Returns:
Optional[Dict[str, Any]]: Modified arguments or None if cancelled
Raises:
PermissionError: If there's an issue with the modification process
"""
print("\n" + "="*60)
print("MODIFY OPERATION")
print("="*60)
print(f"Function: {function_name}")
print("Current arguments:")
# Create a copy of the arguments that will be modified
modified_args = args.copy()
try:
# Present each argument for modification
for key, value in args.items():
while True:
new_value = input(f" {key} [{value}] (Press Enter to keep, 'C' to cancel): ")
if not new_value: # Keep current value
break
elif new_value.upper() == 'C': # Cancel the modification
return None
else:
# Try to convert to the same type as the original value
try:
# Handle different types appropriately
if isinstance(value, bool):
# Convert string to boolean
if new_value.lower() in ('true', 'yes', '1', 'y'):
modified_args[key] = True
elif new_value.lower() in ('false', 'no', '0', 'n'):
modified_args[key] = False
else:
print(f"Invalid boolean value. Use 'true', 'false', 'yes', 'no', '1', '0', 'y', or 'n'.")
continue
elif isinstance(value, int):
modified_args[key] = int(new_value)
elif isinstance(value, float):
modified_args[key] = float(new_value)
else:
# For strings and other types, use the new value as is
modified_args[key] = new_value
break
except ValueError:
print(f"Invalid value for {key}. Expected type: {type(value).__name__}")
# Show the modified arguments for confirmation
print("\nModified arguments:")
for key, value in modified_args.items():
print(f" {key}: {value}")
# Confirm the modifications
confirm = input("\nConfirm these modifications? [Y/n]: ").strip().lower()
if confirm in ('', 'y', 'yes'):
logger.info("Modified args confirmed for function: %s", function_name)
return modified_args
else:
logger.info("Modified args rejected for function: %s", function_name)
return None
except KeyboardInterrupt:
logger.warning("Command modification interrupted")
print("\nModification cancelled.")
return None
except Exception as e:
logger.error("Error modifying command: %s", str(e))
raise PermissionError(
f"Error modifying command: {str(e)}",
function_name=function_name,
args=args
)
def detect_sudo(self, command: str) -> bool:
"""
Detect if a command uses sudo.
Args:
command: Command string
Returns:
bool: True if command uses sudo
"""
# Basic detection for sudo command
if not command:
return False
# Split the command to handle commands with arguments
cmd_parts = command.strip().split()
if not cmd_parts:
return False
# Check if the command starts with sudo
if cmd_parts[0].lower() == "sudo":
logger.info("Sudo command detected: %s", command)
return True
# More advanced detection - check for sudo within a command string
sudo_pattern = r'(^|\s+)sudo(\s+|$)'
if re.search(sudo_pattern, command):
logger.info("Sudo pattern detected in command: %s", command)
return True
return False
def clear_session_permissions(self):
"""
Clear all session permissions.
"""
num_permissions = len(self.session_permissions)
self.session_permissions.clear()
# Clear pattern-based permissions too
for key in self.session_permission_patterns:
self.session_permission_patterns[key] = {}
logger.info("Cleared %d session permissions", num_permissions)
def add_session_permission(self, function_name: str, args: Dict[str, Any]):
"""
Add a specific operation to the session permissions.
Args:
function_name: Name of the function
args: Function arguments
"""
permission_key = self._generate_permission_key(function_name, args)
self.session_permissions.add(permission_key)
# Add to pattern-based permissions
if function_name == "execute_command" and "command" in args:
command = args["command"]
base_command = command.split()[0] if isinstance(command, str) and " " in command else command
self.session_permission_patterns["execute_command"][base_command] = True
elif function_name in ("read_file", "write_file") and "path" in args:
path = args["path"]
parent_dir = os.path.dirname(os.path.abspath(path))
self.session_permission_patterns[function_name][parent_dir] = True
logger.info("Added session permission: %s", permission_key)
def has_session_permission(self, function_name: str, args: Dict[str, Any]) -> bool:
"""
Check if an operation has session permission.
Args:
function_name: Name of the function
args: Function arguments
Returns:
bool: True if the operation has session permission
"""
# Check specific permission key
permission_key = self._generate_permission_key(function_name, args)
if permission_key in self.session_permissions:
return True
# Check pattern-based permissions
if function_name == "execute_command" and "command" in args:
command = args["command"]
base_command = command.split()[0] if isinstance(command, str) and " " in command else command
if base_command in self.session_permission_patterns["execute_command"]:
return True
elif function_name in ("read_file", "write_file") and "path" in args:
path = args["path"]
parent_dir = os.path.dirname(os.path.abspath(path))
if parent_dir in self.session_permission_patterns[function_name]:
return True
return False