PikoAi 0.1.8__tar.gz → 0.1.10__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. {pikoai-0.1.8 → pikoai-0.1.10}/PKG-INFO +1 -1
  2. pikoai-0.1.10/Src/Agents/Executor/executor.py +185 -0
  3. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Agents/Executor/prompts.py +5 -21
  4. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/base_env.py +13 -8
  5. pikoai-0.1.10/Src/Env/python_executor.py +136 -0
  6. pikoai-0.1.10/Src/Env/shell.py +182 -0
  7. pikoai-0.1.10/Src/Env/tests/test_python_executor.py +70 -0
  8. pikoai-0.1.10/Src/Env/tests/test_shell_executor.py +29 -0
  9. {pikoai-0.1.8 → pikoai-0.1.10}/Src/OpenCopilot.py +45 -8
  10. {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/PKG-INFO +1 -1
  11. {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/SOURCES.txt +4 -0
  12. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/tool_dir.json +16 -0
  13. pikoai-0.1.10/Src/Tools/tool_manager.py +105 -0
  14. pikoai-0.1.10/Src/Utils/executor_utils.py +15 -0
  15. pikoai-0.1.10/Src/llm_interface/__init__.py +0 -0
  16. {pikoai-0.1.8 → pikoai-0.1.10}/setup.py +1 -1
  17. pikoai-0.1.10/test/test_opencopilot_file_integration.py +187 -0
  18. pikoai-0.1.8/Src/Agents/Executor/executor.py +0 -241
  19. pikoai-0.1.8/Src/Env/python_executor.py +0 -96
  20. pikoai-0.1.8/Src/Env/shell.py +0 -55
  21. pikoai-0.1.8/Src/Tools/tool_manager.py +0 -52
  22. pikoai-0.1.8/Src/Utils/executor_utils.py +0 -33
  23. {pikoai-0.1.8 → pikoai-0.1.10}/LICENSE +0 -0
  24. {pikoai-0.1.8 → pikoai-0.1.10}/README.md +0 -0
  25. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Agents/Executor/__init__.py +0 -0
  26. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Agents/__init__.py +0 -0
  27. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/__init__.py +0 -0
  28. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/base_executor.py +0 -0
  29. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/env.py +0 -0
  30. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/js_executor.py +0 -0
  31. {pikoai-0.1.8/Src/Tools → pikoai-0.1.10/Src/Env/tests}/__init__.py +0 -0
  32. {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/dependency_links.txt +0 -0
  33. {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/entry_points.txt +0 -0
  34. {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/requires.txt +0 -0
  35. {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/top_level.txt +0 -0
  36. {pikoai-0.1.8/Src/llm_interface → pikoai-0.1.10/Src/Tools}/__init__.py +0 -0
  37. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/file_task.py +0 -0
  38. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/system_details.py +0 -0
  39. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/userinp.py +0 -0
  40. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/web_loader.py +0 -0
  41. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/web_search.py +0 -0
  42. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Utils/__init__.py +0 -0
  43. {pikoai-0.1.8 → pikoai-0.1.10}/Src/Utils/ter_interface.py +0 -0
  44. {pikoai-0.1.8 → pikoai-0.1.10}/Src/cli.py +0 -0
  45. {pikoai-0.1.8 → pikoai-0.1.10}/Src/llm_interface/llm.py +0 -0
  46. {pikoai-0.1.8 → pikoai-0.1.10}/setup.cfg +0 -0
  47. {pikoai-0.1.8 → pikoai-0.1.10}/test/test.py +0 -0
  48. {pikoai-0.1.8 → pikoai-0.1.10}/test/test_file_task.py +0 -0
  49. {pikoai-0.1.8 → pikoai-0.1.10}/test/testjs.py +0 -0
  50. {pikoai-0.1.8 → pikoai-0.1.10}/test/testscript.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: PikoAi
3
- Version: 0.1.8
3
+ Version: 0.1.10
4
4
  Summary: An AI-powered task automation tool
5
5
  Home-page: https://github.com/nihaaaar22/OS-Assistant
6
6
  Author: Nihar S
@@ -0,0 +1,185 @@
1
+ # the change in this executor is the the tasks will not be iterated in a for loop and execution will not be done one by one
2
+ # instead it would be asked what is the next course of action
3
+
4
+ import os
5
+ import sys
6
+ import time
7
+ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))
8
+ from Utils.ter_interface import TerminalInterface
9
+ from Utils.executor_utils import parse_tool_call
10
+ from Agents.Executor.prompts import get_system_prompt, get_task_prompt # Import prompts
11
+
12
+ from typing import Optional
13
+ from mistralai.models.sdkerror import SDKError # This might be an issue if LiteLLM doesn't use SDKError
14
+ # LiteLLM maps exceptions to OpenAI exceptions.
15
+ # We'll keep it for now and see if errors arise during testing.
16
+ # from Env import python_executor # Will be replaced by BaseEnv
17
+ # from Env.shell import ShellExecutor # Will be replaced by BaseEnv
18
+ from Env.base_env import create_environment, BaseEnv # Added
19
+ from Env import python_executor # Keep for type hint in the old execute method if needed, or remove if execute is fully removed
20
+ from llm_interface.llm import LiteLLMInterface # Import LiteLLMInterface
21
+
22
+ from Tools import tool_manager
23
+
24
+ class RateLimiter:
25
+ def __init__(self, wait_time: float = 5.0, max_retries: int = 3):
26
+ self.wait_time = wait_time
27
+ self.max_retries = max_retries
28
+ self.last_call_time = None
29
+
30
+ def wait_if_needed(self):
31
+ if self.last_call_time is not None:
32
+ elapsed = time.time() - self.last_call_time
33
+ if elapsed < 1.0:
34
+ time.sleep(1.0 - elapsed)
35
+ self.last_call_time = time.time()
36
+
37
+ class executor:
38
+ def __init__(self, user_prompt, max_iter=10):
39
+ self.user_prompt = user_prompt
40
+ self.max_iter = max_iter
41
+ self.rate_limiter = RateLimiter(wait_time=5.0, max_retries=3)
42
+ self.executor_prompt_init() # Update system_prompt
43
+ # self.python_executor = python_executor.PythonExecutor() # Initialize PythonExecutor
44
+ # self.shell_executor = ShellExecutor() # Initialize ShellExecutor
45
+ self.message = [
46
+ {"role": "system", "content": self.system_prompt},
47
+ {"role": "user", "content": self.task_prompt}
48
+ ]
49
+ self.terminal = TerminalInterface()
50
+ self.initialize_llm()
51
+
52
+ def initialize_llm(self):
53
+ # Directly instantiate LiteLLMInterface.
54
+ # It handles its own configuration loading (including model_name from config.json).
55
+ self.llm = LiteLLMInterface()
56
+
57
+ def get_tool_dir(self):
58
+ import pkg_resources
59
+ tool_dir_path = pkg_resources.resource_filename('Tools', 'tool_dir.json')
60
+ with open(tool_dir_path, "r") as file:
61
+ return file.read()
62
+
63
+ def executor_prompt_init(self):
64
+ # Load tools details when initializing prompt
65
+ tools_details = self.get_tool_dir()
66
+
67
+ # Read working_directory from config.json
68
+ # This import needs to be here, or moved to the top if json is used elsewhere
69
+ import json
70
+ with open(os.path.join(os.path.dirname(__file__), '../../../config.json'), "r") as config_file:
71
+ config = json.load(config_file)
72
+ working_dir = config.get("working_directory", "")
73
+
74
+ self.system_prompt = get_system_prompt(self.user_prompt, working_dir, tools_details)
75
+ self.task_prompt = get_task_prompt()
76
+
77
+ def run_inference(self):
78
+ retries = 0
79
+ while retries <= self.rate_limiter.max_retries:
80
+ try:
81
+ self.rate_limiter.wait_if_needed()
82
+
83
+ response = self.llm.chat(self.message) # LiteLLMInterface.chat() returns the full response string
84
+
85
+ # Streaming is handled within LiteLLMInterface.chat()
86
+ # and TerminalInterface.process_markdown_chunk()
87
+ self.message.append({"role": "assistant", "content": response})
88
+ return response
89
+
90
+ except Exception as e: # Catching generic Exception as LiteLLM maps to OpenAI exceptions
91
+ # Check if the error message contains "429" for rate limiting
92
+ if "429" in str(e) and retries < self.rate_limiter.max_retries:
93
+ retries += 1
94
+ print(f"\nRate limit error detected. Waiting {self.rate_limiter.wait_time} seconds before retry {retries}/{self.rate_limiter.max_retries}")
95
+ time.sleep(self.rate_limiter.wait_time)
96
+ # Check if the error is an SDKError (though less likely with LiteLLM directly)
97
+ # or if it's any other exception that we should retry or raise.
98
+ elif isinstance(e, SDKError) and "429" in str(e) and retries < self.rate_limiter.max_retries: # Added SDKError check just in case
99
+ retries += 1
100
+ print(f"\nRate limit exceeded (SDKError). Waiting {self.rate_limiter.wait_time} seconds before retry {retries}/{self.rate_limiter.max_retries}")
101
+ time.sleep(self.rate_limiter.wait_time)
102
+ else:
103
+ print(f"\nError occurred during inference: {str(e)}")
104
+ # You might want to log the full traceback here for debugging
105
+ # import traceback
106
+ # print(traceback.format_exc())
107
+ raise
108
+ raise Exception("Failed to complete inference after maximum retries")
109
+
110
+ def run(self):
111
+
112
+ self.run_task()
113
+
114
+ def run_task(self):
115
+ # Remove tools_details parameter since it's in the prompt
116
+ task_message = self.task_prompt
117
+
118
+ self.message.append({"role": "user", "content": task_message})
119
+
120
+ iteration = 0
121
+ task_done = False
122
+
123
+ while iteration < self.max_iter and not task_done:
124
+ # Check for tool calls
125
+ response = self.run_inference()
126
+ tool_call = parse_tool_call(response)
127
+
128
+ if tool_call:
129
+ tool_name = tool_call['tool_name']
130
+ tool_input = tool_call['input']
131
+ print(f"\nIdentified tool call: {tool_name} with input {tool_input}")
132
+
133
+ # Call the tool and append the result (no confirmation or special logic)
134
+ try:
135
+ tool_output_result = tool_manager.call_tool(tool_name, tool_input)
136
+ self.terminal.tool_output_log(tool_output_result, tool_name)
137
+ print(tool_output_result)
138
+ self.message.append({"role": "user", "content": tool_output_result})
139
+ except ValueError as e:
140
+ error_msg = str(e)
141
+ print(f"Tool Error: {error_msg}")
142
+ self.message.append({"role": "user", "content": f"Tool Error: {error_msg}"})
143
+
144
+ else: # Not a tool call, could be a direct response or requires clarification
145
+ # This part handles responses that are not formatted as tool calls.
146
+ # It might be a final answer, a question, or just conversational text.
147
+ # The existing logic for TASK_DONE or asking for next step handles this.
148
+ # No specific code/shell parsing here anymore as they are tools.
149
+ pass # Explicitly pass if no tool call and no old code/shell logic.
150
+
151
+ # Check if task is done
152
+ if "TASK_DONE" in response:
153
+
154
+ task_done = True
155
+
156
+ else:
157
+ self.message.append({"role": "user", "content": "If the task i mentioned is complete then output TASK_DONE .If not then run another iteration."})
158
+ iteration += 1
159
+
160
+ if not task_done:
161
+ print(f"Task could not be completed within {self.max_iter} iterations.")
162
+
163
+ # This method is superseded by the BaseEnv approach in run_task
164
+ # def execute(self, code: str, exec_env: python_executor.PythonExecutor):
165
+ # """Executes the given Python code using the provided execution environment."""
166
+ # result = exec_env.execute(code)
167
+ # return result
168
+
169
+ if __name__ == "__main__":
170
+ # e1 = executor("") # Commenting out example usage for now as it might need adjustment
171
+ # user_prompt = input("Please enter your prompt: ")
172
+ # e1.user_prompt = user_prompt
173
+ # e1.executor_prompt_init() # Update system_prompt
174
+ # e1.message = [
175
+ # {"role": "system", "content": e1.system_prompt},
176
+ # {"role": "user", "content": e1.task_prompt}
177
+ # ] # Reset message list properly
178
+ # e1.run()
179
+
180
+ # while True:
181
+ # user_prompt = input("Please enter your prompt: ")
182
+ # e1.message.append({"role": "user", "content": user_prompt})
183
+ # # e1.message.append({"role":"user","content":e1.system_prompt})
184
+ # e1.run()
185
+ pass # Placeholder if main execution is commented out
@@ -17,9 +17,7 @@ You have access to the following tools:
17
17
 
18
18
  Your primary objective is to accomplish the user's goal by performing step-by-step actions. These actions can include:
19
19
  1. Calling a tool
20
- 2. Executing Python code
21
- 3. Executing Shell commands
22
- 4. Providing a direct response
20
+ 2. Providing a direct response
23
21
 
24
22
  You must break down the user's goal into smaller steps and perform one action at a time. After each action, carefully evaluate the output to determine the next step.
25
23
 
@@ -33,14 +31,9 @@ You must break down the user's goal into smaller steps and perform one action at
33
31
  }}
34
32
  }}
35
33
  <<END_TOOL_CALL>>
36
- - **Code Execution**: Write Python code when no tool is suitable or when custom logic is needed. Format:
37
- <<CODE>>
38
- your_python_code_here
39
- <<CODE>>
40
- - **Shell Command Execution**: Execute shell commands when needed. Format:
41
- <<SHELL_COMMAND>>
42
- your_shell_command_here
43
- <<END_SHELL_COMMAND>>
34
+ This includes executing Python code and shell commands:
35
+ `execute_python_code`: {{"code": "your_python_code_here"}}
36
+ `execute_shell_command`: {{"command": "your_shell_command_here"}}
44
37
  - **Direct Response**: Provide a direct answer if the task doesn't require tools or code.
45
38
 
46
39
  ### Important Notes:
@@ -68,16 +61,7 @@ Following are the things that you must read carefully and remember:
68
61
  }
69
62
  }
70
63
  <<END_TOOL_CALL>>
71
-
72
- - For code execution, use:
73
- <<CODE>>
74
- your_python_code_here
75
- <<CODE>>
76
-
77
- - For shell command execution, use:
78
- <<SHELL_COMMAND>>
79
- your_shell_command_here
80
- <<END_SHELL_COMMAND>>
64
+ Remember that executing Python code and shell commands is now done through specific tool calls (`execute_python_code` and `execute_shell_command`).
81
65
 
82
66
  After each action, always evaluate the output to decide your next step. Only include 'TASK_DONE'
83
67
  When the entire task is completed. Do not end the task immediately after a tool call or code execution without
@@ -1,25 +1,30 @@
1
- from Env.js_executor import JavaScriptExecutor #class import
2
- from Env.python_executor import PythonExecutor #class import
3
-
4
1
  # to perform funciton ask whether to execute code
5
2
 
6
3
 
7
4
  class BaseEnv:
8
5
 
9
6
 
10
- def __init__(self, language,code):
11
- self.language = language
7
+ def __init__(self):
8
+ pass
9
+
12
10
 
11
+ def execute(self, code_or_command: str):
12
+ raise NotImplementedError("This method should be overridden by subclasses")
13
13
 
14
- def execute(self):
14
+ def stop_execution(self):
15
15
  raise NotImplementedError("This method should be overridden by subclasses")
16
16
 
17
17
 
18
18
  def create_environment(language):
19
+ # Moved imports inside the function to avoid circular dependencies during testing
20
+ from Env.js_executor import JavaScriptExecutor #class import
21
+ from Env.python_executor import PythonExecutor #class import
22
+ from Env.shell import ShellExecutor #class import
23
+
19
24
  if language == "python":
20
25
  return PythonExecutor()
21
- elif language == "javascript":
22
- return JavaScriptExecutor()
26
+ elif language == "shell":
27
+ return ShellExecutor()
23
28
  else:
24
29
  raise ValueError(f"Unsupported language: {language}")
25
30
 
@@ -0,0 +1,136 @@
1
+ # from .base_executor import BaseExecutor
2
+
3
+ # class PythonExecutor():
4
+ # def execute(self, code: str) -> str:
5
+ # """Executes Python code and returns the result or an error message."""
6
+
7
+ # # if not self.validate_code(code):
8
+ # # return "Code validation failed: Unsafe code detected."
9
+
10
+ # local_vars = {}
11
+ # try:
12
+ # exec(code, {}, local_vars) # Execute code in an isolated environment
13
+ # return local_vars.get("output", "Code executed successfully.")
14
+ # except Exception as e:
15
+ # # return self.handle_error(e)
16
+ # print("error in running python code", e)
17
+
18
+ import subprocess
19
+ import tempfile
20
+ import os
21
+ from typing import Dict
22
+ import textwrap
23
+ import sys
24
+ from Src.Env.base_env import BaseEnv
25
+ import time
26
+
27
+ class PythonExecutor(BaseEnv):
28
+ def __init__(self):
29
+ super().__init__()
30
+ self.process = None
31
+ self.forbidden_terms = [
32
+ 'import os', 'import sys', 'import subprocess',
33
+ 'open(', 'exec(', 'eval(',
34
+ ]
35
+
36
+ def basic_code_check(self, code: str) -> bool:
37
+ """Simple check for potentially dangerous code"""
38
+ code_lower = code.lower()
39
+ return not any(term.lower() in code_lower for term in self.forbidden_terms)
40
+
41
+ def execute(self, code_or_command: str) -> Dict[str, str]:
42
+ """Executes Python code in a separate process and returns the result"""
43
+
44
+ # Basic safety check
45
+ if not self.basic_code_check(code_or_command):
46
+ return {
47
+ 'success': False,
48
+ 'output': 'Error: Code contains potentially unsafe operations. You can try and use tools to achieve same functionality.',
49
+ 'error': 'Security check failed'
50
+ }
51
+
52
+ # Properly indent the code to fit inside the try block
53
+ indented_code = textwrap.indent(code_or_command, ' ')
54
+ # Wrap the indented code to capture output
55
+ wrapped_code = f"""
56
+ try:
57
+ {indented_code}
58
+ except Exception as e:
59
+ print(f"Error: {{str(e)}}")
60
+ """
61
+
62
+
63
+ try:
64
+ # Execute the code in a subprocess
65
+ self.process = subprocess.Popen(
66
+ [sys.executable, "-u", "-c", wrapped_code],
67
+ stdout=subprocess.PIPE,
68
+ stderr=subprocess.PIPE,
69
+ text=True,
70
+ bufsize=1, # Line buffered
71
+ universal_newlines=True
72
+ )
73
+
74
+ stdout_data = []
75
+ stderr_data = []
76
+ start_time = time.time()
77
+
78
+ # First read all stdout
79
+ for line in self.process.stdout:
80
+ # Check for timeout
81
+ if time.time() - start_time > 30:
82
+ self.process.kill()
83
+ return {
84
+ 'success': False,
85
+ 'output': 'Execution timed out after 30 seconds',
86
+ 'error': 'Timeout error'
87
+ }
88
+
89
+ stdout_data.append(line)
90
+ print(line, end='', flush=True) # Print in real-time
91
+
92
+ # Then read all stderr
93
+ for line in self.process.stderr:
94
+ # Check for timeout
95
+ if time.time() - start_time > 30:
96
+ self.process.kill()
97
+ return {
98
+ 'success': False,
99
+ 'output': 'Execution timed out after 30 seconds',
100
+ 'error': 'Timeout error'
101
+ }
102
+
103
+ stderr_data.append(line)
104
+ print(line, end='', file=sys.stderr, flush=True) # Print in real-time
105
+
106
+ # Wait for process to complete
107
+ returncode = self.process.wait()
108
+
109
+ return {
110
+ 'success': returncode == 0,
111
+ 'output': ''.join(stdout_data) if returncode == 0 else ''.join(stderr_data),
112
+ 'error': ''.join(stderr_data) if returncode != 0 else ''
113
+ }
114
+
115
+ except Exception as e:
116
+ return {
117
+ 'success': False,
118
+ 'output': f'Error: {str(e)}',
119
+ 'error': str(e)
120
+ }
121
+ finally:
122
+ self.process = None # Reset process
123
+
124
+
125
+ def stop_execution(self):
126
+ if self.process and hasattr(self.process, 'pid') and self.process.pid is not None:
127
+ try:
128
+ self.process.terminate()
129
+ print(f"Attempted to terminate Python process with PID: {self.process.pid}")
130
+ except Exception as e:
131
+ print(f"Error terminating Python process with PID {self.process.pid}: {e}")
132
+ finally:
133
+ self.process = None
134
+ else:
135
+ print("No active Python process to stop.")
136
+
@@ -0,0 +1,182 @@
1
+ import subprocess
2
+ import time
3
+ import sys
4
+ from Src.Env.base_env import BaseEnv
5
+ import re
6
+
7
+ class ShellExecutor(BaseEnv):
8
+ def __init__(self):
9
+ super().__init__()
10
+ self.process = None
11
+
12
+ def execute(self, code_or_command: str) -> dict:
13
+ """
14
+ Executes a shell command and streams its output in real-time.
15
+ Strictly prevents execution of harmful commands and access to sensitive directories.
16
+
17
+ Args:
18
+ code_or_command: The shell command to execute.
19
+
20
+ Returns:
21
+ A dictionary with the following keys:
22
+ - 'output': The captured standard output (string).
23
+ - 'error': The captured standard error (string).
24
+ - 'success': A boolean indicating whether the command executed successfully.
25
+ """
26
+ try:
27
+ # Strict security check for harmful commands and sensitive directories
28
+ forbidden_patterns = [
29
+ r'rm\s+-rf\s+/?(\s|$)',
30
+ r'rm\s+-rf\s+--no-preserve-root',
31
+ r'rm\s+-rf\s+/\\?',
32
+ r'shutdown(\s|$)',
33
+ r'reboot(\s|$)',
34
+ r'halt(\s|$)',
35
+ r':(){:|:&};:', # fork bomb
36
+ r'chmod\s+777\s+/(\s|$)',
37
+ r'chown\s+root',
38
+ r'\bmkfs\b',
39
+ r'\bdd\b.*\bif=\/dev\/zero\b',
40
+ r'\bdd\b.*\bof=\/dev\/sda',
41
+ r'\bpoweroff\b',
42
+ r'\binit\s+0\b',
43
+ r'\bsudo\s+rm\s+-rf\s+/?',
44
+ r'\bsudo\s+shutdown',
45
+ r'\bsudo\s+reboot',
46
+ r'\bsudo\s+halt',
47
+ r'\bsudo\s+mkfs',
48
+ r'\bsudo\s+dd',
49
+ r'\bsudo\s+init\s+0',
50
+ r'\bsudo\s+poweroff',
51
+ r'\bsudo\s+chmod\s+777\s+/',
52
+ r'\bsudo\s+chown\s+root',
53
+ r'\bdel\b.*\/s.*\/q.*\/f.*C:\\', # Windows
54
+ r'format\s+C:',
55
+ r'rd\s+/?s\s+/?q\s+C:\\',
56
+ r'\bshutdown\b.*\/s',
57
+ r'\bshutdown\b.*\/r',
58
+ r'\bshutdown\b.*\/f',
59
+ r'\bshutdown\b.*\/p',
60
+ r'\bshutdown\b.*\/t',
61
+ r'\bshutdown\b.*\/a',
62
+ r'\bnet\s+user\s+.*\s+/delete',
63
+ r'\bnet\s+user\s+administrator\s+/active:no',
64
+ r'\bnet\s+user\s+administrator\s+/active:yes',
65
+ r'\bnet\s+localgroup\s+administrators\s+.*\s+/delete',
66
+ r'\bnet\s+localgroup\s+administrators\s+.*\s+/add',
67
+ ]
68
+ sensitive_dirs = [
69
+ '/', '/etc', '/bin', '/usr', '/var', '/root', '/boot', '/dev', '/proc', '/sys', '/lib', '/lib64',
70
+ 'C:\\', 'C:/', 'C:\\Windows', 'C:/Windows', 'C:\\System32', 'C:/System32',
71
+ 'D:\\', 'D:/', 'E:\\', 'E:/'
72
+ ]
73
+ # Check for forbidden patterns
74
+ for pattern in forbidden_patterns:
75
+ if re.search(pattern, code_or_command, re.IGNORECASE):
76
+ return {
77
+ 'success': False,
78
+ 'output': 'Blocked potentially harmful command.',
79
+ 'error': f'Command matches forbidden pattern: {pattern}'
80
+ }
81
+ # Check for sensitive directory access
82
+ for sensitive_dir in sensitive_dirs:
83
+ # Only block if the command is trying to directly access or operate on the sensitive dir
84
+ if re.search(rf'\b{sensitive_dir}\b', code_or_command, re.IGNORECASE):
85
+ return {
86
+ 'success': False,
87
+ 'output': f'Blocked access to sensitive directory: {sensitive_dir}',
88
+ 'error': f'Attempted access to sensitive directory: {sensitive_dir}'
89
+ }
90
+
91
+ # Execute the command in a subprocess
92
+ self.process = subprocess.Popen(
93
+ code_or_command,
94
+ shell=True,
95
+ stdout=subprocess.PIPE,
96
+ stderr=subprocess.PIPE,
97
+ text=True,
98
+ bufsize=1, # Line buffered
99
+ universal_newlines=True
100
+ )
101
+
102
+ stdout_data = []
103
+ stderr_data = []
104
+ start_time = time.time()
105
+
106
+ # First read all stdout
107
+ for line in self.process.stdout:
108
+ # Check for timeout
109
+ if time.time() - start_time > 30:
110
+ self.process.kill()
111
+ return {
112
+ 'success': False,
113
+ 'output': 'Execution timed out after 30 seconds',
114
+ 'error': 'Timeout error'
115
+ }
116
+
117
+ stdout_data.append(line)
118
+ print(line, end='', flush=True) # Print in real-time
119
+
120
+ # Then read all stderr
121
+ for line in self.process.stderr:
122
+ # Check for timeout
123
+ if time.time() - start_time > 30:
124
+ self.process.kill()
125
+ return {
126
+ 'success': False,
127
+ 'output': 'Execution timed out after 30 seconds',
128
+ 'error': 'Timeout error'
129
+ }
130
+
131
+ stderr_data.append(line)
132
+ print(line, end='', file=sys.stderr, flush=True) # Print in real-time
133
+
134
+ # Wait for process to complete
135
+ returncode = self.process.wait()
136
+
137
+ return {
138
+ 'success': returncode == 0,
139
+ 'output': ''.join(stdout_data) if returncode == 0 else ''.join(stderr_data),
140
+ 'error': ''.join(stderr_data) if returncode != 0 else ''
141
+ }
142
+
143
+ except Exception as e:
144
+ return {
145
+ 'success': False,
146
+ 'output': f'Error: {str(e)}',
147
+ 'error': str(e)
148
+ }
149
+ finally:
150
+ self.process = None # Reset process
151
+
152
+ def stop_execution(self):
153
+ if self.process and hasattr(self.process, 'pid') and self.process.pid is not None:
154
+ try:
155
+ self.process.terminate()
156
+ print(f"Attempted to terminate shell process with PID: {self.process.pid}")
157
+ except Exception as e:
158
+ print(f"Error terminating shell process with PID {self.process.pid}: {e}")
159
+ finally:
160
+ self.process = None
161
+ else:
162
+ print("No active shell process to stop.")
163
+
164
+ if __name__ == '__main__':
165
+ # Example usage (optional, for testing)
166
+ executor = ShellExecutor()
167
+
168
+ # Test case 1: Successful command
169
+ result1 = executor.execute("echo 'Hello, World!'")
170
+ print(f"Test Case 1 Result: {result1}")
171
+
172
+ # Test case 2: Command with an error
173
+ result2 = executor.execute("ls non_existent_directory")
174
+ print(f"Test Case 2 Result: {result2}")
175
+
176
+ # Test case 3: Command that succeeds but writes to stderr (e.g. some warnings)
177
+ result3 = executor.execute("echo 'Error output' >&2")
178
+ print(f"Test Case 3 Result: {result3}")
179
+
180
+ # Test case 4: Command that produces no output
181
+ result4 = executor.execute(":") # The ':' command is a no-op in bash
182
+ print(f"Test Case 4 Result: {result4}")
@@ -0,0 +1,70 @@
1
+ import unittest
2
+ from Src.Env.python_executor import PythonExecutor
3
+ import time
4
+ import threading
5
+
6
+ class TestPythonExecutor(unittest.TestCase):
7
+
8
+ def test_execute_simple_script(self):
9
+ executor = PythonExecutor()
10
+ result = executor.execute("print('hello python')")
11
+ self.assertTrue(result["success"])
12
+ self.assertEqual(result["output"].strip(), "hello python")
13
+ self.assertEqual(result["error"], "")
14
+
15
+ def test_execute_script_with_error(self):
16
+ executor = PythonExecutor()
17
+ result = executor.execute("1/0")
18
+ self.assertFalse(result["success"])
19
+ self.assertTrue("ZeroDivisionError" in result["output"]) # Error message goes to stdout for python_executor
20
+ self.assertEqual(result["error"], "") # stderr should be empty as Popen merges stderr to stdout in this case
21
+
22
+ def test_stop_execution_long_script(self):
23
+ executor = PythonExecutor()
24
+ long_script = "import time; time.sleep(5); print('should not print')"
25
+
26
+ execute_result = {}
27
+ def target():
28
+ res = executor.execute(long_script)
29
+ execute_result.update(res)
30
+
31
+ thread = threading.Thread(target=target)
32
+ thread.start()
33
+
34
+ time.sleep(1) # Give the script time to start
35
+ executor.stop_execution()
36
+
37
+ thread.join(timeout=2) # Wait for the thread to finish (should be quick after stop)
38
+
39
+ self.assertFalse(execute_result.get("success", True), "Script execution should have failed or been stopped.")
40
+ # Depending on timing, the process might be killed before it produces output,
41
+ # or it might produce a timeout error, or a specific error from being terminated.
42
+ # We check if the output indicates it didn't complete normally.
43
+ output = execute_result.get("output", "")
44
+ error = execute_result.get("error", "")
45
+
46
+ # Check if 'should not print' is NOT in the output
47
+ self.assertNotIn("should not print", output, "Script should have been terminated before completion.")
48
+
49
+ # Check for signs of termination or timeout
50
+ # This part is a bit tricky as the exact message can vary.
51
+ # If basic_code_check fails, output is "Error: Code contains potentially unsafe operations..."
52
+ # If timeout in communicate(), output is "Execution timed out..."
53
+ # If process is terminated, output might be empty or contain partial error.
54
+ # For now, we'll accept that if "should not print" is not there, it's a good sign.
55
+ # A more robust check might involve looking for specific error messages related to termination if available.
56
+
57
+ # A simple check that process is no longer listed in executor
58
+ self.assertIsNone(executor.process, "Executor process should be None after stopping.")
59
+
60
+
61
+ def test_stop_execution_no_script(self):
62
+ executor = PythonExecutor()
63
+ try:
64
+ executor.stop_execution()
65
+ except Exception as e:
66
+ self.fail(f"stop_execution with no script raised an exception: {e}")
67
+ self.assertIsNone(executor.process, "Executor process should be None.")
68
+
69
+ if __name__ == '__main__':
70
+ unittest.main()