| … | … | @@ -0,0 +1,106 @@ |
| 1 | #!/usr/bin/env python3 | |
| 2 | """ | |
| 3 | Example script demonstrating the enhanced permission system. | |
| 4 | ||
| 5 | This script shows how to use the interactive permission prompts | |
| 6 | and the session-based permission system. | |
| 7 | """ | |
| 8 | ||
| 9 | import sys | |
| 10 | import os | |
| 11 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| 12 | ||
| 13 | from ollamatools.security import PermissionManager | |
| 14 | ||
| 15 | ||
| 16 | def simulate_function_execution(function_name, args, is_sudo=False): | |
| 17 | """Simulate the execution of a function.""" | |
| 18 | if is_sudo: | |
| 19 | print(f"[SIMULATED] Executing sudo operation: {function_name}") | |
| 20 | else: | |
| 21 | print(f"[SIMULATED] Executing: {function_name}") | |
| 22 | ||
| 23 | print(f"[SIMULATED] Arguments: {args}") | |
| 24 | print(f"[SIMULATED] Operation completed successfully") | |
| 25 | print() | |
| 26 | ||
| 27 | ||
| 28 | def main(): | |
| 29 | # Create a permission manager with default settings | |
| 30 | pm = PermissionManager() | |
| 31 | ||
| 32 | print("\n=== Enhanced Permission System Example ===\n") | |
| 33 | print("This example demonstrates the interactive permission system.") | |
| 34 | print("You will be prompted to allow or deny various operations.") | |
| 35 | print("Try different responses (A, S, M, D) to see how they work.") | |
| 36 | print("\nPress Enter to continue...") | |
| 37 | input() | |
| 38 | ||
| 39 | # Example 1: Basic function permission | |
| 40 | print("\n--- Example 1: Basic Function Permission ---") | |
| 41 | print("A simple read operation that requires permission.") | |
| 42 | if pm.request_permission("read_file", {"path": "/etc/hosts"}): | |
| 43 | simulate_function_execution("read_file", {"path": "/etc/hosts"}) | |
| 44 | else: | |
| 45 | print("Operation denied by user.") | |
| 46 | ||
| 47 | # Example 2: Sudo command permission | |
| 48 | print("\n--- Example 2: Sudo Command Permission ---") | |
| 49 | print("An operation that requires elevated permissions (sudo).") | |
| 50 | if pm.request_permission("execute_command", {"command": "sudo apt update"}, is_sudo=True): | |
| 51 | simulate_function_execution("execute_command", {"command": "sudo apt update"}, is_sudo=True) | |
| 52 | else: | |
| 53 | print("Sudo operation denied by user.") | |
| 54 | ||
| 55 | # Example 3: Session permission | |
| 56 | print("\n--- Example 3: Session Permission ---") | |
| 57 | print("Try using the 'S' option to allow for session.") | |
| 58 | print("The same operation will then be allowed without prompting again.") | |
| 59 | ||
| 60 | # First request - will prompt | |
| 61 | print("\nFirst request:") | |
| 62 | if pm.request_permission("read_config", {"file": "settings.json"}): | |
| 63 | simulate_function_execution("read_config", {"file": "settings.json"}) | |
| 64 | else: | |
| 65 | print("Operation denied by user.") | |
| 66 | ||
| 67 | # Second request - may not prompt depending on user's choice | |
| 68 | print("\nSecond request (same operation):") | |
| 69 | if pm.request_permission("read_config", {"file": "settings.json"}): | |
| 70 | simulate_function_execution("read_config", {"file": "settings.json"}) | |
| 71 | else: | |
| 72 | print("Operation denied by user.") | |
| 73 | ||
| 74 | # Example 4: Modify operation | |
| 75 | print("\n--- Example 4: Modify Operation ---") | |
| 76 | print("Try using the 'M' option to modify the operation before execution.") | |
| 77 | ||
| 78 | if pm.request_permission("write_file", {"path": "/tmp/example.txt", "content": "Hello, World!"}): | |
| 79 | simulate_function_execution("write_file", {"path": "/tmp/example.txt", "content": "Hello, World!"}) | |
| 80 | else: | |
| 81 | print("Operation denied or modification cancelled by user.") | |
| 82 | ||
| 83 | # Example 5: Disabled prompts but sudo always prompts | |
| 84 | print("\n--- Example 5: Disabled Prompts but Sudo Always Prompts ---") | |
| 85 | print("Creating a permission manager with regular prompts disabled but sudo prompts enabled.") | |
| 86 | ||
| 87 | pm_custom = PermissionManager(prompt_enabled=False, sudo_always_prompt=True) | |
| 88 | ||
| 89 | # Regular operation - should not prompt | |
| 90 | print("\nRegular operation (should execute without prompting):") | |
| 91 | if pm_custom.request_permission("list_files", {"directory": "/tmp"}): | |
| 92 | simulate_function_execution("list_files", {"directory": "/tmp"}) | |
| 93 | else: | |
| 94 | print("Operation denied.") | |
| 95 | ||
| 96 | # Sudo operation - should still prompt | |
| 97 | print("\nSudo operation (should still prompt):") | |
| 98 | if pm_custom.request_permission("execute_command", {"command": "sudo systemctl status"}, is_sudo=True): | |
| 99 | simulate_function_execution("execute_command", {"command": "sudo systemctl status"}, is_sudo=True) | |
| 100 | else: | |
| 101 | print("Sudo operation denied by user.") | |
| 102 | ||
| 103 | ||
| 104 | if __name__ == "__main__": | |
| 105 | main() | |
| 106 |
| … | … | @@ -4,7 +4,7 @@ Permission management system for OllamaTools. |
| 4 | 4 | This module provides security controls for executing functions, |
| 5 | 5 | with a focus on a prompt-based permission system. |
| 6 | 6 | """ |
| 7 | from typing import Dict, Any, Optional | |
| 7 | from typing import Dict, Any, Optional, Tuple, Set | |
| 8 | 8 | |
| 9 | 9 | |
| 10 | 10 | class PermissionManager: |
| … | … | @@ -19,6 +19,7 @@ class PermissionManager: |
| 19 | 19 | Attributes: |
| 20 | 20 | prompt_enabled (bool): Whether to prompt for permission before operations |
| 21 | 21 | sudo_always_prompt (bool): Whether to always prompt for sudo commands |
| 22 | session_permissions (Set[str]): Set of operations allowed for the session | |
| 22 | 23 | """ |
| 23 | 24 | |
| 24 | 25 | def __init__( |
| … | … | @@ -35,6 +36,7 @@ class PermissionManager: |
| 35 | 36 | """ |
| 36 | 37 | self.prompt_enabled = prompt_enabled |
| 37 | 38 | self.sudo_always_prompt = sudo_always_prompt |
| 39 | self.session_permissions = set() # Store allowed operations for the session | |
| 38 | 40 | |
| 39 | 41 | def request_permission( |
| 40 | 42 | self, |
| … | … | @@ -56,20 +58,142 @@ class PermissionManager: |
| 56 | 58 | Returns: |
| 57 | 59 | bool: True if permission granted, False otherwise |
| 58 | 60 | """ |
| 59 | # For now, simply print the request and return True | |
| 60 | # Will be enhanced with actual prompting in the future | |
| 61 | ||
| 62 | 61 | # Format the arguments as a readable string |
| 63 | 62 | args_str = ', '.join(f"{k}={v!r}" for k, v in args.items()) |
| 64 | 63 | |
| 65 | # Print information about the permission request | |
| 66 | print(f"Permission request for: {function_name}({args_str})") | |
| 64 | # Generate a permission key for the operation | |
| 65 | permission_key = f"{function_name}:{args_str}" | |
| 66 | ||
| 67 | # Check if we need to prompt for permission | |
| 68 | should_prompt = self.prompt_enabled | |
| 69 | ||
| 70 | # Always prompt for sudo if sudo_always_prompt is True | |
| 71 | if is_sudo and self.sudo_always_prompt: | |
| 72 | should_prompt = True | |
| 73 | ||
| 74 | # Check if this operation is already allowed for the session | |
| 75 | if permission_key in self.session_permissions: | |
| 76 | return True | |
| 77 | ||
| 78 | # If prompting is disabled and not overridden by sudo, grant permission | |
| 79 | if not should_prompt: | |
| 80 | return True | |
| 81 | ||
| 82 | # Display the permission prompt | |
| 83 | print("\n" + "="*60) | |
| 84 | print("PERMISSION REQUEST") | |
| 85 | print("="*60) | |
| 86 | print(f"Function: {function_name}") | |
| 87 | print(f"Arguments: {args_str}") | |
| 88 | ||
| 67 | 89 | if is_sudo: |
| 68 | print("Note: This operation requires elevated permissions (sudo)") | |
| 90 | print("\nWARNING: This operation requires elevated permissions (sudo)") | |
| 91 | print(" Be cautious when granting permission for sudo operations") | |
| 92 | ||
| 93 | print("\nOptions:") | |
| 94 | print(" [A] Allow: Grant permission for this operation") | |
| 95 | print(" [S] Allow for Session: Grant permission for similar operations in this session") | |
| 96 | print(" [M] Modify: Modify the operation before execution") | |
| 97 | print(" [D] Deny: Reject the operation") | |
| 98 | ||
| 99 | while True: | |
| 100 | response = input("\nYour choice [A/S/M/D]: ").strip().upper() | |
| 101 | ||
| 102 | if response == 'A': # Allow once | |
| 103 | return True | |
| 104 | ||
| 105 | elif response == 'S': # Allow for session | |
| 106 | self.session_permissions.add(permission_key) | |
| 107 | return True | |
| 108 | ||
| 109 | elif response == 'M': # Modify | |
| 110 | modified_args = self.modify_command(function_name, args) | |
| 111 | if modified_args is not None: | |
| 112 | # Call request_permission again with modified arguments | |
| 113 | # but make sure we don't get into an infinite loop | |
| 114 | return self.request_permission( | |
| 115 | function_name, | |
| 116 | modified_args, | |
| 117 | self.detect_sudo(modified_args.get('command', '')) if 'command' in modified_args else is_sudo | |
| 118 | ) | |
| 119 | else: | |
| 120 | print("Modification cancelled.") | |
| 121 | return False | |
| 122 | ||
| 123 | elif response == 'D': # Deny | |
| 124 | print("Permission denied.") | |
| 125 | return False | |
| 126 | ||
| 127 | else: | |
| 128 | print("Invalid choice. Please enter A, S, M, or D.") | |
| 129 | ||
| 130 | def modify_command( | |
| 131 | self, | |
| 132 | function_name: str, | |
| 133 | args: Dict[str, Any] | |
| 134 | ) -> Optional[Dict[str, Any]]: | |
| 135 | """ | |
| 136 | Allow the user to modify command arguments. | |
| 137 | ||
| 138 | Args: | |
| 139 | function_name: Name of the function | |
| 140 | args: Current function arguments | |
| 141 | ||
| 142 | Returns: | |
| 143 | Optional[Dict[str, Any]]: Modified arguments or None if cancelled | |
| 144 | """ | |
| 145 | print("\n" + "="*60) | |
| 146 | print("MODIFY OPERATION") | |
| 147 | print("="*60) | |
| 148 | print(f"Function: {function_name}") | |
| 149 | print("Current arguments:") | |
| 150 | ||
| 151 | # Create a copy of the arguments that will be modified | |
| 152 | modified_args = args.copy() | |
| 153 | ||
| 154 | # Present each argument for modification | |
| 155 | for key, value in args.items(): | |
| 156 | while True: | |
| 157 | new_value = input(f" {key} [{value}] (Press Enter to keep, 'C' to cancel): ") | |
| 158 | ||
| 159 | if not new_value: # Keep current value | |
| 160 | break | |
| 161 | elif new_value.upper() == 'C': # Cancel the modification | |
| 162 | return None | |
| 163 | else: | |
| 164 | # Try to convert to the same type as the original value | |
| 165 | try: | |
| 166 | # Handle different types appropriately | |
| 167 | if isinstance(value, bool): | |
| 168 | # Convert string to boolean | |
| 169 | if new_value.lower() in ('true', 'yes', '1', 'y'): | |
| 170 | modified_args[key] = True | |
| 171 | elif new_value.lower() in ('false', 'no', '0', 'n'): | |
| 172 | modified_args[key] = False | |
| 173 | else: | |
| 174 | print(f"Invalid boolean value. Use 'true', 'false', 'yes', 'no', '1', '0', 'y', or 'n'.") | |
| 175 | continue | |
| 176 | elif isinstance(value, int): | |
| 177 | modified_args[key] = int(new_value) | |
| 178 | elif isinstance(value, float): | |
| 179 | modified_args[key] = float(new_value) | |
| 180 | else: | |
| 181 | # For strings and other types, use the new value as is | |
| 182 | modified_args[key] = new_value | |
| 183 | break | |
| 184 | except ValueError: | |
| 185 | print(f"Invalid value for {key}. Expected type: {type(value).__name__}") | |
| 186 | ||
| 187 | print("\nModified arguments:") | |
| 188 | for key, value in modified_args.items(): | |
| 189 | print(f" {key}: {value}") | |
| 69 | 190 | |
| 70 | # In this basic implementation, always grant permission | |
| 71 | # This will be replaced with actual prompting in step 2.2 | |
| 72 | return True | |
| 191 | # Confirm the modifications | |
| 192 | confirm = input("\nConfirm these modifications? [Y/n]: ").strip().lower() | |
| 193 | if confirm in ('', 'y', 'yes'): | |
| 194 | return modified_args | |
| 195 | else: | |
| 196 | return None | |
| 73 | 197 | |
| 74 | 198 | def detect_sudo(self, command: str) -> bool: |
| 75 | 199 | """ |
| … | … | @@ -2,59 +2,10 @@ |
| 2 | 2 | Tests for the permission management system. |
| 3 | 3 | """ |
| 4 | 4 | import pytest |
| 5 | from ollamatools.security import PermissionManager | |
| 6 | ||
| 7 | ||
| 8 | def test_permission_manager_init(): | |
| 9 | """Test PermissionManager initialization with default and custom parameters.""" | |
| 10 | # Test default initialization | |
| 11 | pm = PermissionManager() | |
| 12 | assert pm.prompt_enabled is True | |
| 13 | assert pm.sudo_always_prompt is True | |
| 14 | ||
| 15 | # Test custom initialization | |
| 16 | pm = PermissionManager(prompt_enabled=False, sudo_always_prompt=False) | |
| 17 | assert pm.prompt_enabled is False | |
| 18 | assert pm.sudo_always_prompt is False | |
| 5 | import sys | |
| 6 | import os | |
| 19 | 7 | |
| 8 | # Add the parent directory to the path so we can import the package | |
| 9 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| 20 | 10 | |
| 21 | def test_request_permission(capfd): | |
| 22 | """Test the basic request_permission method.""" | |
| 23 | pm = PermissionManager() | |
| 24 | ||
| 25 | # Test basic function | |
| 26 | result = pm.request_permission("test_function", {"arg1": "value1"}) | |
| 27 | assert result is True # In basic implementation, should always return True | |
| 28 | ||
| 29 | # Check if appropriate information was printed | |
| 30 | captured = capfd.readouterr() | |
| 31 | assert "Permission request for: test_function" in captured.out | |
| 32 | assert "arg1='value1'" in captured.out | |
| 33 | ||
| 34 | # Test with sudo | |
| 35 | result = pm.request_permission("test_function", {"arg1": "value1"}, is_sudo=True) | |
| 36 | assert result is True | |
| 37 | ||
| 38 | # Check if sudo warning was printed | |
| 39 | captured = capfd.readouterr() | |
| 40 | assert "elevated permissions (sudo)" in captured.out | |
| 41 | ||
| 42 | ||
| 43 | def test_detect_sudo(): | |
| 44 | """Test the detect_sudo method with various commands.""" | |
| 45 | pm = PermissionManager() | |
| 46 | ||
| 47 | # Test sudo commands | |
| 48 | assert pm.detect_sudo("sudo ls -la") is True | |
| 49 | assert pm.detect_sudo("sudo apt update") is True | |
| 50 | assert pm.detect_sudo("SUDO rm -rf /") is True # Case insensitive check | |
| 51 | ||
| 52 | # Test non-sudo commands | |
| 53 | assert pm.detect_sudo("ls -la") is False | |
| 54 | assert pm.detect_sudo("apt update") is False | |
| 55 | assert pm.detect_sudo("echo 'sudo is not at start'") is False | |
| 56 | ||
| 57 | # Test edge cases | |
| 58 | assert pm.detect_sudo("") is False | |
| 59 | assert pm.detect_sudo(" ") is False | |
| 60 | assert pm.detect_sudo(None) is False | |
| 11 | from ollamatools.security import PermissionManager |