OverviewFilesCommitsIssuesDocsReleasesPeers

ollama_semantic_function_calling commits / ad4923a

adding in command execution func
ad4923a7cb33a01eca4d7f10acadb9776f459955
by fleeky.ashatron <fleeky@7rnx.net> · 2025-04-10T12:42:07.000Z

Changed files (4) +930 -15

ollamatools-project/examples/command_execution.py +96 -0
……@@ -0,0 +1,96 @@
1#!/usr/bin/env python3
2"""
3Example script demonstrating command execution with OllamaTools.
4
5This script shows how to use the command execution function with
6various examples and security features.
7"""
8
9import sys
10import os
11import time
12
13# Add the parent directory to the path so we can import the package
14sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
15
16from ollamatools.tools.system_tools import execute_command
17from ollamatools.security import PermissionManager
18
19
20def print_result(result):
21 """Print command execution result in a formatted way."""
22 print("\nCommand Result:")
23 print(f"Command: {result['command']}")
24 print(f"Exit Code: {result['exit_code']}")
25
26 if result['stdout']:
27 print("\nStandard Output:")
28 print("-" * 40)
29 print(result['stdout'])
30 print("-" * 40)
31
32 if result['stderr']:
33 print("\nStandard Error:")
34 print("-" * 40)
35 print(result['stderr'])
36 print("-" * 40)
37
38
39def main():
40 print("=== Command Execution Example ===\n")
41
42 # Create a permission manager for secure command execution
43 pm = PermissionManager()
44
45 try:
46 # Example 1: Simple command without arguments
47 print("\n--- Example 1: Simple Command ---")
48 print("Executing: 'pwd'")
49 result = execute_command("pwd", permission_manager=pm)
50 print_result(result)
51
52 # Example 2: Command with arguments
53 print("\n--- Example 2: Command with Arguments ---")
54 print("Executing: 'ls -l'")
55 result = execute_command("ls", ["-l"], permission_manager=pm)
56 print_result(result)
57
58 # Example 3: Command with an error
59 print("\n--- Example 3: Command with Error ---")
60 print("Executing: 'ls --invalid-option'")
61 result = execute_command("ls", ["--invalid-option"], permission_manager=pm)
62 print_result(result)
63
64 # Example 4: Command with timeout
65 print("\n--- Example 4: Command with Timeout ---")
66 print("Executing: 'sleep 10' with a 2 second timeout")
67 try:
68 result = execute_command("sleep", ["10"], timeout=2, permission_manager=pm)
69 print_result(result)
70 except TimeoutError as e:
71 print(f"Caught timeout error: {e}")
72
73 # Example 5: Sudo command
74 print("\n--- Example 5: Sudo Command ---")
75 print("Executing: 'sudo echo \"Hello from root\"'")
76 try:
77 result = execute_command("sudo", ["echo", "Hello from root"], permission_manager=pm)
78 print_result(result)
79 except PermissionError as e:
80 print(f"Permission error: {e}")
81
82 # Example 6: Potentially dangerous command (should be blocked)
83 print("\n--- Example 6: Dangerous Command ---")
84 print("Attempting to execute a dangerous command with shell redirection")
85 try:
86 result = execute_command("echo 'Hello' > /tmp/test.txt", permission_manager=pm)
87 print_result(result)
88 except ValueError as e:
89 print(f"Caught security error: {e}")
90
91 except Exception as e:
92 print(f"An error occurred: {e}")
93
94
95if __name__ == "__main__":
96 main()
ollamatools-project/ollamatools/tools/__init__.py +19 -2
……@@ -1,5 +1,22 @@
11"""
2Function groups for OllamaTools.
2OllamaTools function libraries.
33
4This package contains various function groups that can be used by the OllamaAgent.
4This package contains various function libraries that can be used
5with the OllamaAgent for semantic function detection and execution.
56"""
7
8# Import and expose system tools
9from ollamatools.tools.system_tools import (
10 execute_command,
11 read_file,
12 write_file,
13 find_files
14)
15
16# Define which functions should be exposed at the package level
17__all__ = [
18 "execute_command",
19 "read_file",
20 "write_file",
21 "find_files"
22]
ollamatools-project/ollamatools/tools/system_tools.py +323 -13
……@@ -1,32 +1,342 @@
11"""
2System tools for OllamaTools.
2System tools for command execution and file operations.
33
4This module contains functions for executing system commands and performing
5file operations.
4This module provides functions for executing system commands and
5performing file operations with security controls.
66"""
77
8from typing import Optional, List, Dict, Any
8import subprocess
9import shlex
10import os
11import re
12import glob
13from typing import Dict, Any, List, Optional, Union, Tuple
14
15from ollamatools.security import PermissionManager
916
1017
1118def execute_command(
1219 command: str,
1320 args: Optional[List[str]] = None,
14 timeout: int = 60
21 timeout: int = 60,
22 permission_manager: Optional[PermissionManager] = None
1523) -> Dict[str, Any]:
1624 """
1725 Execute a shell command with arguments.
1826
27 This function executes a shell command and returns the results.
28 It includes security measures such as sanitization and timeout.
29
1930 Args:
20 command: Command to execute
21 args: Command arguments
31 command: Command to execute (will be sanitized)
32 args: Optional list of command arguments
2233 timeout: Command timeout in seconds
34 permission_manager: Optional permission manager for security
2335
2436 Returns:
2537 dict: Result with stdout, stderr, and exit code
38
39 Raises:
40 ValueError: If command contains dangerous patterns
41 TimeoutError: If command execution times out
42 subprocess.SubprocessError: If command execution fails
43 PermissionError: If permission is denied
44 """
45 # Sanitize command
46 sanitized_command = _sanitize_command(command)
47
48 # Build full command with arguments
49 if args:
50 # Sanitize each argument
51 sanitized_args = [_sanitize_argument(arg) for arg in args]
52 full_command = [sanitized_command] + sanitized_args
53 else:
54 # Split the command into parts for better security (no shell=True)
55 full_command = shlex.split(sanitized_command)
56
57 # Check for sudo
58 is_sudo = False
59 if full_command and full_command[0].lower() == "sudo":
60 is_sudo = True
61
62 # Request permission if permission manager is provided
63 if permission_manager:
64 # Build command string for permission request
65 command_str = " ".join(full_command)
66 if not permission_manager.request_permission(
67 "execute_command",
68 {"command": command_str},
69 is_sudo=is_sudo
70 ):
71 raise PermissionError(f"Permission denied for command: {command_str}")
72
73 try:
74 # Execute the command with timeout
75 # Use subprocess.run instead of subprocess.Popen for simpler code
76 # Use check=False to handle errors in the return value
77 result = subprocess.run(
78 full_command,
79 capture_output=True,
80 text=True,
81 timeout=timeout,
82 check=False # Don't raise exception on non-zero exit codes
83 )
84
85 return {
86 "stdout": result.stdout,
87 "stderr": result.stderr,
88 "exit_code": result.returncode,
89 "command": " ".join(full_command)
90 }
91
92 except subprocess.TimeoutExpired:
93 raise TimeoutError(f"Command timed out after {timeout} seconds: {' '.join(full_command)}")
94
95 except (subprocess.SubprocessError, OSError) as e:
96 # Convert to a more informative error
97 raise subprocess.SubprocessError(f"Command execution failed: {e}")
98
99
100def read_file(
101 path: str,
102 max_size: Optional[int] = None,
103 permission_manager: Optional[PermissionManager] = None
104) -> str:
105 """
106 Read file contents.
107
108 Args:
109 path: File path
110 max_size: Maximum file size to read in bytes
111 permission_manager: Optional permission manager for security
112
113 Returns:
114 str: File contents
115
116 Raises:
117 FileNotFoundError: If file does not exist
118 PermissionError: If permission is denied
119 ValueError: If file exceeds max_size
120 """
121 # Sanitize and normalize path
122 path = os.path.normpath(os.path.abspath(path))
123
124 # Request permission if permission manager is provided
125 if permission_manager:
126 if not permission_manager.request_permission(
127 "read_file",
128 {"path": path}
129 ):
130 raise PermissionError(f"Permission denied for reading file: {path}")
131
132 # Check if file exists
133 if not os.path.isfile(path):
134 raise FileNotFoundError(f"File not found: {path}")
135
136 # Check file size if max_size is specified
137 if max_size is not None:
138 file_size = os.path.getsize(path)
139 if file_size > max_size:
140 raise ValueError(f"File size ({file_size} bytes) exceeds maximum allowed size ({max_size} bytes)")
141
142 # Read and return file contents
143 try:
144 with open(path, 'r', encoding='utf-8') as file:
145 return file.read()
146 except UnicodeDecodeError:
147 # Handle binary files by reading in binary mode and providing a warning
148 with open(path, 'rb') as file:
149 binary_content = file.read()
150 warning = f"[WARNING: File appears to be binary. Displaying as hex.]\n\n"
151 # Convert binary to hex representation for safe display
152 hex_content = binary_content.hex(' ', 16)
153 return warning + hex_content
154
155
156def write_file(
157 path: str,
158 content: str,
159 mode: str = 'w',
160 permission_manager: Optional[PermissionManager] = None
161) -> bool:
162 """
163 Write content to a file.
164
165 Args:
166 path: File path
167 content: Content to write
168 mode: Write mode ('w' for write, 'a' for append)
169 permission_manager: Optional permission manager for security
170
171 Returns:
172 bool: True if successful
173
174 Raises:
175 PermissionError: If permission is denied
176 ValueError: If mode is invalid
177 IOError: If file cannot be written
178 """
179 # Validate mode
180 if mode not in ('w', 'a'):
181 raise ValueError(f"Invalid mode: {mode}. Must be 'w' (write) or 'a' (append)")
182
183 # Sanitize and normalize path
184 path = os.path.normpath(os.path.abspath(path))
185
186 # Request permission if permission manager is provided
187 if permission_manager:
188 operation = "write_file" if mode == 'w' else "append_file"
189 if not permission_manager.request_permission(
190 operation,
191 {"path": path, "mode": mode}
192 ):
193 raise PermissionError(f"Permission denied for writing to file: {path}")
194
195 # Ensure directory exists
196 directory = os.path.dirname(path)
197 if directory and not os.path.exists(directory):
198 try:
199 os.makedirs(directory, exist_ok=True)
200 except OSError as e:
201 raise IOError(f"Failed to create directory '{directory}': {e}")
202
203 # Write content to file
204 try:
205 with open(path, mode, encoding='utf-8') as file:
206 file.write(content)
207 return True
208 except IOError as e:
209 raise IOError(f"Failed to write to file '{path}': {e}")
210
211
212def find_files(
213 pattern: str,
214 directory: str = '.',
215 recursive: bool = True,
216 permission_manager: Optional[PermissionManager] = None
217) -> List[str]:
26218 """
27 # Placeholder implementation
28 return {
29 "stdout": "Command execution placeholder",
30 "stderr": "",
31 "exit_code": 0
32 }
219 Find files matching a pattern.
220
221 Args:
222 pattern: File pattern (glob syntax)
223 directory: Search directory
224 recursive: Whether to search recursively
225 permission_manager: Optional permission manager for security
226
227 Returns:
228 List[str]: Matching file paths
229
230 Raises:
231 PermissionError: If permission is denied
232 ValueError: If directory does not exist
233 ValueError: If pattern is invalid
234 """
235 # Sanitize and normalize directory path
236 directory = os.path.normpath(os.path.abspath(directory))
237
238 # Validate directory
239 if not os.path.isdir(directory):
240 raise ValueError(f"Directory not found: {directory}")
241
242 # Request permission if permission manager is provided
243 if permission_manager:
244 if not permission_manager.request_permission(
245 "find_files",
246 {"pattern": pattern, "directory": directory, "recursive": recursive}
247 ):
248 raise PermissionError(f"Permission denied for finding files in: {directory}")
249
250 # Find matching files
251 try:
252 if recursive:
253 # Use ** for recursive search if not already in pattern
254 if '**' not in pattern:
255 search_pattern = os.path.join(directory, '**', pattern)
256 else:
257 search_pattern = os.path.join(directory, pattern)
258
259 # Use recursive glob for Python 3.5+
260 matching_files = glob.glob(search_pattern, recursive=True)
261 else:
262 # Non-recursive search
263 search_pattern = os.path.join(directory, pattern)
264 matching_files = glob.glob(search_pattern)
265
266 # Sort results for consistent output
267 matching_files.sort()
268
269 return matching_files
270
271 except (ValueError, re.error) as e:
272 raise ValueError(f"Invalid pattern '{pattern}': {e}")
273
274
275def _sanitize_command(command: str) -> str:
276 """
277 Sanitize a command for security.
278
279 Args:
280 command: Command to sanitize
281
282 Returns:
283 str: Sanitized command
284
285 Raises:
286 ValueError: If command contains dangerous patterns
287 """
288 if not command or not command.strip():
289 raise ValueError("Empty command")
290
291 # Remove leading/trailing whitespace
292 command = command.strip()
293
294 # Disallow certain dangerous patterns
295 dangerous_patterns = [
296 r'[|&;]', # Shell control operators
297 r'>\s*\w+', # Redirections
298 r'<\s*\w+', # Input redirections
299 r'\$\(', # Command substitution
300 r'`', # Backtick command substitution
301 ]
302
303 for pattern in dangerous_patterns:
304 if re.search(pattern, command):
305 raise ValueError(f"Command contains dangerous pattern: {pattern}")
306
307 return command
308
309
310def _sanitize_argument(arg: str) -> str:
311 """
312 Sanitize a command argument for security.
313
314 Args:
315 arg: Argument to sanitize
316
317 Returns:
318 str: Sanitized argument
319
320 Raises:
321 ValueError: If argument contains dangerous patterns
322 """
323 if not isinstance(arg, str):
324 arg = str(arg)
325
326 # Remove leading/trailing whitespace
327 arg = arg.strip()
328
329 # Disallow certain dangerous patterns (similar to command sanitization)
330 dangerous_patterns = [
331 r'[|&;]', # Shell control operators
332 r'>\s*\w+', # Redirections
333 r'<\s*\w+', # Input redirections
334 r'\$\(', # Command substitution
335 r'`', # Backtick command substitution
336 ]
337
338 for pattern in dangerous_patterns:
339 if re.search(pattern, arg):
340 raise ValueError(f"Argument contains dangerous pattern: {pattern}")
341
342 return arg
ollamatools-project/tests/test_system_tools.py +492 -0
……@@ -0,0 +1,492 @@
1"""
2Tests for the system tools module.
3"""
4import pytest
5import sys
6import os
7import re
8import tempfile
9import shutil
10from unittest.mock import patch, MagicMock
11import subprocess
12
13# Add the parent directory to the path so we can import the package
14sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
15
16from ollamatools.tools.system_tools import (
17 execute_command,
18 read_file,
19 write_file,
20 find_files,
21 _sanitize_command,
22 _sanitize_argument
23)
24from ollamatools.security import PermissionManager
25
26
27# Setup and teardown for temporary test directory
28@pytest.fixture
29def temp_dir():
30 """Create a temporary directory for file operations tests."""
31 tmp_dir = tempfile.mkdtemp()
32 try:
33 yield tmp_dir
34 finally:
35 shutil.rmtree(tmp_dir)
36
37
38def test_sanitize_command():
39 """Test the command sanitization function."""
40 # Valid commands
41 assert _sanitize_command("ls") == "ls"
42 assert _sanitize_command("ls -la") == "ls -la"
43 assert _sanitize_command("echo 'Hello World'") == "echo 'Hello World'"
44 assert _sanitize_command(" find /tmp -name '*.txt' ") == "find /tmp -name '*.txt'"
45
46 # Invalid commands
47 with pytest.raises(ValueError) as excinfo:
48 _sanitize_command("")
49 assert "Empty command" in str(excinfo.value)
50
51 with pytest.raises(ValueError) as excinfo:
52 _sanitize_command("ls | grep txt")
53 assert "dangerous pattern" in str(excinfo.value)
54
55 with pytest.raises(ValueError) as excinfo:
56 _sanitize_command("echo 'hello' > file.txt")
57 assert "dangerous pattern" in str(excinfo.value)
58
59 with pytest.raises(ValueError) as excinfo:
60 _sanitize_command("cat < file.txt")
61 assert "dangerous pattern" in str(excinfo.value)
62
63 with pytest.raises(ValueError) as excinfo:
64 _sanitize_command("echo $(whoami)")
65 assert "dangerous pattern" in str(excinfo.value)
66
67 with pytest.raises(ValueError) as excinfo:
68 _sanitize_command("echo `whoami`")
69 assert "dangerous pattern" in str(excinfo.value)
70
71
72def test_sanitize_argument():
73 """Test the argument sanitization function."""
74 # Valid arguments
75 assert _sanitize_argument("file.txt") == "file.txt"
76 assert _sanitize_argument("-la") == "-la"
77 assert _sanitize_argument("--help") == "--help"
78 assert _sanitize_argument("/tmp/file with spaces.txt") == "/tmp/file with spaces.txt"
79
80 # Test non-string arguments
81 assert _sanitize_argument(123) == "123"
82
83 # Invalid arguments
84 with pytest.raises(ValueError) as excinfo:
85 _sanitize_argument("file.txt | grep txt")
86 assert "dangerous pattern" in str(excinfo.value)
87
88 with pytest.raises(ValueError) as excinfo:
89 _sanitize_argument("$(whoami)")
90 assert "dangerous pattern" in str(excinfo.value)
91
92
93def test_execute_command_simple():
94 """Test executing a simple command."""
95 # Test a simple echo command
96 result = execute_command("echo", ["Hello", "World"])
97
98 assert result["stdout"].strip() == "Hello World"
99 assert result["stderr"] == ""
100 assert result["exit_code"] == 0
101 assert result["command"] == "echo Hello World"
102
103
104def test_execute_command_error():
105 """Test executing a command that returns an error."""
106 # Test a command that should fail (non-existent command)
107 result = execute_command("ls", ["--invalid-option"])
108
109 assert result["stderr"] != "" # Should have error output
110 assert result["exit_code"] != 0 # Non-zero exit code
111
112 # Test a command that doesn't exist
113 with pytest.raises(subprocess.SubprocessError):
114 execute_command("command_that_does_not_exist")
115
116
117def test_execute_command_timeout():
118 """Test command timeout handling."""
119 # Patch subprocess.run to simulate a timeout
120 with patch('subprocess.run') as mock_run:
121 mock_run.side_effect = subprocess.TimeoutExpired(cmd="sleep 10", timeout=1)
122
123 with pytest.raises(TimeoutError) as excinfo:
124 execute_command("sleep", ["10"], timeout=1)
125
126 assert "timed out" in str(excinfo.value)
127
128
129def test_execute_command_with_permission_manager():
130 """Test executing a command with the permission manager."""
131 # Create a mock permission manager
132 mock_pm = MagicMock(spec=PermissionManager)
133
134 # Test with permission granted
135 mock_pm.request_permission.return_value = True
136 result = execute_command("echo", ["Hello"], permission_manager=mock_pm)
137 assert result["stdout"].strip() == "Hello"
138 assert mock_pm.request_permission.called
139
140 # Test with permission denied
141 mock_pm.request_permission.return_value = False
142 with pytest.raises(PermissionError) as excinfo:
143 execute_command("echo", ["Hello"], permission_manager=mock_pm)
144 assert "Permission denied" in str(excinfo.value)
145 assert mock_pm.request_permission.called
146
147
148def test_execute_command_sudo():
149 """Test sudo command detection and handling."""
150 # Create a mock permission manager
151 mock_pm = MagicMock(spec=PermissionManager)
152 mock_pm.request_permission.return_value = True
153
154 # We can't actually execute sudo commands in tests, so we'll just check
155 # that the sudo flag is correctly passed to the permission manager
156 with patch('subprocess.run') as mock_run:
157 mock_run.return_value = subprocess.CompletedProcess(
158 args=["sudo", "ls"], returncode=0, stdout="", stderr=""
159 )
160
161 execute_command("sudo", ["ls"], permission_manager=mock_pm)
162
163 # Check that the is_sudo flag was set to True
164 mock_pm.request_permission.assert_called_with(
165 "execute_command",
166 {"command": "sudo ls"},
167 is_sudo=True
168 )
169
170
171def test_read_file(temp_dir):
172 """Test reading file contents."""
173 # Create a test file
174 test_file_path = os.path.join(temp_dir, "test_file.txt")
175 test_content = "Hello, world!\nThis is a test file."
176
177 with open(test_file_path, 'w') as f:
178 f.write(test_content)
179
180 # Test reading the file
181 content = read_file(test_file_path)
182 assert content == test_content
183
184 # Test with non-existent file
185 with pytest.raises(FileNotFoundError):
186 read_file(os.path.join(temp_dir, "nonexistent.txt"))
187
188
189def test_read_file_with_max_size(temp_dir):
190 """Test reading file with max size limit."""
191 # Create a test file
192 test_file_path = os.path.join(temp_dir, "large_file.txt")
193 test_content = "X" * 1000 # 1000 bytes
194
195 with open(test_file_path, 'w') as f:
196 f.write(test_content)
197
198 # Test reading with sufficient max_size
199 content = read_file(test_file_path, max_size=2000)
200 assert content == test_content
201
202 # Test reading with insufficient max_size
203 with pytest.raises(ValueError) as excinfo:
204 read_file(test_file_path, max_size=500)
205 assert "exceeds maximum allowed size" in str(excinfo.value)
206
207
208def test_read_file_with_permission_manager(temp_dir):
209 """Test reading file with permission manager."""
210 # Create a test file
211 test_file_path = os.path.join(temp_dir, "test_file.txt")
212 test_content = "Hello, world!"
213
214 with open(test_file_path, 'w') as f:
215 f.write(test_content)
216
217 # Create a mock permission manager
218 mock_pm = MagicMock(spec=PermissionManager)
219
220 # Test with permission granted
221 mock_pm.request_permission.return_value = True
222 content = read_file(test_file_path, permission_manager=mock_pm)
223 assert content == test_content
224 assert mock_pm.request_permission.called
225
226 # Test with permission denied
227 mock_pm.request_permission.return_value = False
228 with pytest.raises(PermissionError) as excinfo:
229 read_file(test_file_path, permission_manager=mock_pm)
230 assert "Permission denied" in str(excinfo.value)
231 assert mock_pm.request_permission.called
232
233
234def test_write_file(temp_dir):
235 """Test writing content to a file."""
236 # Define test file path and content
237 test_file_path = os.path.join(temp_dir, "write_test.txt")
238 test_content = "This is test content to write to a file."
239
240 # Test writing to file
241 result = write_file(test_file_path, test_content)
242 assert result is True
243
244 # Verify file content
245 with open(test_file_path, 'r') as f:
246 assert f.read() == test_content
247
248 # Test append mode
249 additional_content = "\nThis is additional content."
250 result = write_file(test_file_path, additional_content, mode='a')
251 assert result is True
252
253 # Verify appended content
254 with open(test_file_path, 'r') as f:
255 assert f.read() == test_content + additional_content
256
257 # Test invalid mode
258 with pytest.raises(ValueError) as excinfo:
259 write_file(test_file_path, test_content, mode='x')
260 assert "Invalid mode" in str(excinfo.value)
261
262
263def test_write_file_with_new_directory(temp_dir):
264 """Test writing to a file in a non-existent directory."""
265 # Define test file path in a new directory
266 new_dir = os.path.join(temp_dir, "new_directory")
267 test_file_path = os.path.join(new_dir, "test_file.txt")
268 test_content = "This is a test file in a new directory."
269
270 # Write to file (should create directory)
271 result = write_file(test_file_path, test_content)
272 assert result is True
273
274 # Verify directory and file were created
275 assert os.path.exists(new_dir)
276 assert os.path.isfile(test_file_path)
277
278 # Verify file content
279 with open(test_file_path, 'r') as f:
280 assert f.read() == test_content
281
282
283def test_write_file_with_permission_manager(temp_dir):
284 """Test writing file with permission manager."""
285 # Define test file path and content
286 test_file_path = os.path.join(temp_dir, "permission_test.txt")
287 test_content = "Test content for permission test."
288
289 # Create a mock permission manager
290 mock_pm = MagicMock(spec=PermissionManager)
291
292 # Test with permission granted
293 mock_pm.request_permission.return_value = True
294 result = write_file(test_file_path, test_content, permission_manager=mock_pm)
295 assert result is True
296 assert os.path.exists(test_file_path)
297 assert mock_pm.request_permission.called
298
299 # Test with permission denied
300 mock_pm.request_permission.return_value = False
301 with pytest.raises(PermissionError) as excinfo:
302 write_file(test_file_path, test_content, permission_manager=mock_pm)
303 assert "Permission denied" in str(excinfo.value)
304 assert mock_pm.request_permission.called
305
306
307def test_find_files(temp_dir):
308 """Test finding files matching a pattern."""
309 # Create test files
310 test_files = [
311 "file1.txt",
312 "file2.txt",
313 "image1.jpg",
314 "image2.jpg",
315 "nested/file3.txt",
316 "nested/file4.txt",
317 "nested/deep/file5.txt"
318 ]
319
320 # Create directories and files
321 for file_path in test_files:
322 full_path = os.path.join(temp_dir, file_path)
323 os.makedirs(os.path.dirname(full_path), exist_ok=True)
324 with open(full_path, 'w') as f:
325 f.write(f"Content of {file_path}")
326
327 # Test non-recursive search for txt files
328 txt_files = find_files("*.txt", directory=temp_dir, recursive=False)
329 assert len(txt_files) == 2 # Only top-level txt files
330 assert all(os.path.basename(f) in ["file1.txt", "file2.txt"] for f in txt_files)
331
332 # Test recursive search for txt files
333 all_txt_files = find_files("*.txt", directory=temp_dir, recursive=True)
334 assert len(all_txt_files) == 5 # All txt files
335
336 # Test search for jpg files
337 jpg_files = find_files("*.jpg", directory=temp_dir, recursive=True)
338 assert len(jpg_files) == 2
339 assert all(os.path.basename(f) in ["image1.jpg", "image2.jpg"] for f in jpg_files)
340
341 # Test search"""
342Tests for the system tools module.
343"""
344import pytest
345import sys
346import os
347import re
348from unittest.mock import patch, MagicMock
349import subprocess
350
351# Add the parent directory to the path so we can import the package
352sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
353
354from ollamatools.tools.system_tools import (
355 execute_command,
356 _sanitize_command,
357 _sanitize_argument
358)
359from ollamatools.security import PermissionManager
360
361
362def test_sanitize_command():
363 """Test the command sanitization function."""
364 # Valid commands
365 assert _sanitize_command("ls") == "ls"
366 assert _sanitize_command("ls -la") == "ls -la"
367 assert _sanitize_command("echo 'Hello World'") == "echo 'Hello World'"
368 assert _sanitize_command(" find /tmp -name '*.txt' ") == "find /tmp -name '*.txt'"
369
370 # Invalid commands
371 with pytest.raises(ValueError) as excinfo:
372 _sanitize_command("")
373 assert "Empty command" in str(excinfo.value)
374
375 with pytest.raises(ValueError) as excinfo:
376 _sanitize_command("ls | grep txt")
377 assert "dangerous pattern" in str(excinfo.value)
378
379 with pytest.raises(ValueError) as excinfo:
380 _sanitize_command("echo 'hello' > file.txt")
381 assert "dangerous pattern" in str(excinfo.value)
382
383 with pytest.raises(ValueError) as excinfo:
384 _sanitize_command("cat < file.txt")
385 assert "dangerous pattern" in str(excinfo.value)
386
387 with pytest.raises(ValueError) as excinfo:
388 _sanitize_command("echo $(whoami)")
389 assert "dangerous pattern" in str(excinfo.value)
390
391 with pytest.raises(ValueError) as excinfo:
392 _sanitize_command("echo `whoami`")
393 assert "dangerous pattern" in str(excinfo.value)
394
395
396def test_sanitize_argument():
397 """Test the argument sanitization function."""
398 # Valid arguments
399 assert _sanitize_argument("file.txt") == "file.txt"
400 assert _sanitize_argument("-la") == "-la"
401 assert _sanitize_argument("--help") == "--help"
402 assert _sanitize_argument("/tmp/file with spaces.txt") == "/tmp/file with spaces.txt"
403
404 # Test non-string arguments
405 assert _sanitize_argument(123) == "123"
406
407 # Invalid arguments
408 with pytest.raises(ValueError) as excinfo:
409 _sanitize_argument("file.txt | grep txt")
410 assert "dangerous pattern" in str(excinfo.value)
411
412 with pytest.raises(ValueError) as excinfo:
413 _sanitize_argument("$(whoami)")
414 assert "dangerous pattern" in str(excinfo.value)
415
416
417def test_execute_command_simple():
418 """Test executing a simple command."""
419 # Test a simple echo command
420 result = execute_command("echo", ["Hello", "World"])
421
422 assert result["stdout"].strip() == "Hello World"
423 assert result["stderr"] == ""
424 assert result["exit_code"] == 0
425 assert result["command"] == "echo Hello World"
426
427
428def test_execute_command_error():
429 """Test executing a command that returns an error."""
430 # Test a command that should fail (non-existent command)
431 result = execute_command("ls", ["--invalid-option"])
432
433 assert result["stderr"] != "" # Should have error output
434 assert result["exit_code"] != 0 # Non-zero exit code
435
436 # Test a command that doesn't exist
437 with pytest.raises(subprocess.SubprocessError):
438 execute_command("command_that_does_not_exist")
439
440
441def test_execute_command_timeout():
442 """Test command timeout handling."""
443 # Patch subprocess.run to simulate a timeout
444 with patch('subprocess.run') as mock_run:
445 mock_run.side_effect = subprocess.TimeoutExpired(cmd="sleep 10", timeout=1)
446
447 with pytest.raises(TimeoutError) as excinfo:
448 execute_command("sleep", ["10"], timeout=1)
449
450 assert "timed out" in str(excinfo.value)
451
452
453def test_execute_command_with_permission_manager():
454 """Test executing a command with the permission manager."""
455 # Create a mock permission manager
456 mock_pm = MagicMock(spec=PermissionManager)
457
458 # Test with permission granted
459 mock_pm.request_permission.return_value = True
460 result = execute_command("echo", ["Hello"], permission_manager=mock_pm)
461 assert result["stdout"].strip() == "Hello"
462 assert mock_pm.request_permission.called
463
464 # Test with permission denied
465 mock_pm.request_permission.return_value = False
466 with pytest.raises(PermissionError) as excinfo:
467 execute_command("echo", ["Hello"], permission_manager=mock_pm)
468 assert "Permission denied" in str(excinfo.value)
469 assert mock_pm.request_permission.called
470
471
472def test_execute_command_sudo():
473 """Test sudo command detection and handling."""
474 # Create a mock permission manager
475 mock_pm = MagicMock(spec=PermissionManager)
476 mock_pm.request_permission.return_value = True
477
478 # We can't actually execute sudo commands in tests, so we'll just check
479 # that the sudo flag is correctly passed to the permission manager
480 with patch('subprocess.run') as mock_run:
481 mock_run.return_value = subprocess.CompletedProcess(
482 args=["sudo", "ls"], returncode=0, stdout="", stderr=""
483 )
484
485 execute_command("sudo", ["ls"], permission_manager=mock_pm)
486
487 # Check that the is_sudo flag was set to True
488 mock_pm.request_permission.assert_called_with(
489 "execute_command",
490 {"command": "sudo ls"},
491 is_sudo=True
492 )