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.
- {pikoai-0.1.8 → pikoai-0.1.10}/PKG-INFO +1 -1
- pikoai-0.1.10/Src/Agents/Executor/executor.py +185 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Agents/Executor/prompts.py +5 -21
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/base_env.py +13 -8
- pikoai-0.1.10/Src/Env/python_executor.py +136 -0
- pikoai-0.1.10/Src/Env/shell.py +182 -0
- pikoai-0.1.10/Src/Env/tests/test_python_executor.py +70 -0
- pikoai-0.1.10/Src/Env/tests/test_shell_executor.py +29 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/OpenCopilot.py +45 -8
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/PKG-INFO +1 -1
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/SOURCES.txt +4 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/tool_dir.json +16 -0
- pikoai-0.1.10/Src/Tools/tool_manager.py +105 -0
- pikoai-0.1.10/Src/Utils/executor_utils.py +15 -0
- pikoai-0.1.10/Src/llm_interface/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/setup.py +1 -1
- pikoai-0.1.10/test/test_opencopilot_file_integration.py +187 -0
- pikoai-0.1.8/Src/Agents/Executor/executor.py +0 -241
- pikoai-0.1.8/Src/Env/python_executor.py +0 -96
- pikoai-0.1.8/Src/Env/shell.py +0 -55
- pikoai-0.1.8/Src/Tools/tool_manager.py +0 -52
- pikoai-0.1.8/Src/Utils/executor_utils.py +0 -33
- {pikoai-0.1.8 → pikoai-0.1.10}/LICENSE +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/README.md +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Agents/Executor/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Agents/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/base_executor.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/env.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Env/js_executor.py +0 -0
- {pikoai-0.1.8/Src/Tools → pikoai-0.1.10/Src/Env/tests}/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/dependency_links.txt +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/entry_points.txt +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/requires.txt +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/PikoAi.egg-info/top_level.txt +0 -0
- {pikoai-0.1.8/Src/llm_interface → pikoai-0.1.10/Src/Tools}/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/file_task.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/system_details.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/userinp.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/web_loader.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Tools/web_search.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Utils/__init__.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/Utils/ter_interface.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/cli.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/Src/llm_interface/llm.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/setup.cfg +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/test/test.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/test/test_file_task.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/test/testjs.py +0 -0
- {pikoai-0.1.8 → pikoai-0.1.10}/test/testscript.py +0 -0
@@ -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.
|
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
|
-
|
37
|
-
|
38
|
-
|
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
|
11
|
-
|
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
|
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 == "
|
22
|
-
return
|
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()
|