Open-source release for Alpha version
This commit is contained in:
4
strix/tools/python/__init__.py
Normal file
4
strix/tools/python/__init__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
from .python_actions import python_action
|
||||
|
||||
|
||||
__all__ = ["python_action"]
|
||||
47
strix/tools/python/python_actions.py
Normal file
47
strix/tools/python/python_actions.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from typing import Any, Literal
|
||||
|
||||
from strix.tools.registry import register_tool
|
||||
|
||||
from .python_manager import get_python_session_manager
|
||||
|
||||
|
||||
PythonAction = Literal["new_session", "execute", "close", "list_sessions"]
|
||||
|
||||
|
||||
@register_tool
|
||||
def python_action(
|
||||
action: PythonAction,
|
||||
code: str | None = None,
|
||||
timeout: int = 30,
|
||||
session_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
def _validate_code(action_name: str, code: str | None) -> None:
|
||||
if not code:
|
||||
raise ValueError(f"code parameter is required for {action_name} action")
|
||||
|
||||
def _validate_action(action_name: str) -> None:
|
||||
raise ValueError(f"Unknown action: {action_name}")
|
||||
|
||||
manager = get_python_session_manager()
|
||||
|
||||
try:
|
||||
match action:
|
||||
case "new_session":
|
||||
return manager.create_session(session_id, code, timeout)
|
||||
|
||||
case "execute":
|
||||
_validate_code(action, code)
|
||||
assert code is not None
|
||||
return manager.execute_code(session_id, code, timeout)
|
||||
|
||||
case "close":
|
||||
return manager.close_session(session_id)
|
||||
|
||||
case "list_sessions":
|
||||
return manager.list_sessions()
|
||||
|
||||
case _:
|
||||
_validate_action(action) # type: ignore[unreachable]
|
||||
|
||||
except (ValueError, RuntimeError) as e:
|
||||
return {"stderr": str(e), "session_id": session_id, "stdout": "", "is_running": False}
|
||||
131
strix/tools/python/python_actions_schema.xml
Normal file
131
strix/tools/python/python_actions_schema.xml
Normal file
@@ -0,0 +1,131 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<tools>
|
||||
<tool name="python_action">
|
||||
<description>Perform Python actions using persistent interpreter sessions for cybersecurity tasks.</description>
|
||||
<details>Common Use Cases:
|
||||
- Security script development and testing (payload generation, exploit scripts)
|
||||
- Data analysis of security logs, network traffic, or vulnerability scans
|
||||
- Cryptographic operations and security tool automation
|
||||
- Interactive penetration testing workflows and proof-of-concept development
|
||||
- Processing security data formats (JSON, XML, CSV from security tools)
|
||||
- HTTP proxy interaction for web security testing (all proxy functions are pre-imported)
|
||||
|
||||
Each session instance is PERSISTENT and maintains its own global and local namespaces
|
||||
until explicitly closed, allowing for multi-step security workflows and stateful computations.
|
||||
|
||||
PROXY FUNCTIONS PRE-IMPORTED:
|
||||
All proxy action functions are automatically imported into every Python session, enabling
|
||||
seamless HTTP traffic analysis and web security testing
|
||||
|
||||
This is particularly useful for:
|
||||
- Analyzing captured HTTP traffic during web application testing
|
||||
- Automating request manipulation and replay attacks
|
||||
- Building custom security testing workflows combining proxy data with Python analysis
|
||||
- Correlating multiple requests for advanced attack scenarios</details>
|
||||
<parameters>
|
||||
<parameter name="action" type="string" required="true">
|
||||
<description>The Python action to perform: - new_session: Create a new Python interpreter session. This MUST be the first action for each session. - execute: Execute Python code in the specified session. - close: Close the specified session instance. - list_sessions: List all active Python sessions.</description>
|
||||
</parameter>
|
||||
<parameter name="code" type="string" required="false">
|
||||
<description>Required for 'new_session' (as initial code) and 'execute' actions. The Python code to execute.</description>
|
||||
</parameter>
|
||||
<parameter name="timeout" type="integer" required="false">
|
||||
<description>Maximum execution time in seconds for code execution. Applies to both 'new_session' (when initial code is provided) and 'execute' actions. Default is 30 seconds.</description>
|
||||
</parameter>
|
||||
<parameter name="session_id" type="string" required="false">
|
||||
<description>Unique identifier for the Python session. If not provided, uses the default session ID.</description>
|
||||
</parameter>
|
||||
</parameters>
|
||||
<returns type="Dict[str, Any]">
|
||||
<description>Response containing: - session_id: the ID of the session that was operated on - stdout: captured standard output from code execution (for execute action) - stderr: any error message if execution failed - result: string representation of the last expression result - execution_time: time taken to execute the code - message: status message about the action performed - Various session info depending on the action</description>
|
||||
</returns>
|
||||
<notes>
|
||||
Important usage rules:
|
||||
1. PERSISTENCE: Session instances remain active and maintain their state (variables,
|
||||
imports, function definitions) until explicitly closed with the 'close' action.
|
||||
This allows for multi-step workflows across multiple tool calls.
|
||||
2. MULTIPLE SESSIONS: You can run multiple Python sessions concurrently by using
|
||||
different session_id values. Each session operates independently with its own
|
||||
namespace.
|
||||
3. Session interaction MUST begin with 'new_session' action for each session instance.
|
||||
4. Only one action can be performed per call.
|
||||
5. CODE EXECUTION:
|
||||
- Both expressions and statements are supported
|
||||
- Expressions automatically return their result
|
||||
- Print statements and stdout are captured
|
||||
- Variables persist between executions in the same session
|
||||
- Imports, function definitions, etc. persist in the session
|
||||
- IPython magic commands are fully supported (%pip, %time, %whos, %%writefile, etc.)
|
||||
- Line magics (%) and cell magics (%%) work as expected
|
||||
6. CLOSE: Terminates the session completely and frees memory
|
||||
7. The Python sessions can operate concurrently with other tools. You may invoke
|
||||
terminal, browser, or other tools while maintaining active Python sessions.
|
||||
8. Each session has its own isolated namespace - variables in one session don't
|
||||
affect others.
|
||||
</notes>
|
||||
<examples>
|
||||
# Create new session for security analysis (default session)
|
||||
<function=python_action>
|
||||
<parameter=action>new_session</parameter>
|
||||
<parameter=code>import hashlib
|
||||
import base64
|
||||
import json
|
||||
print("Security analysis session started")</parameter>
|
||||
</function>
|
||||
|
||||
# Analyze security data in the default session
|
||||
<function=python_action>
|
||||
<parameter=action>execute</parameter>
|
||||
<parameter=code>vulnerability_data = {"cve": "CVE-2024-1234", "severity": "high"}
|
||||
encoded_payload = base64.b64encode(json.dumps(vulnerability_data).encode())
|
||||
print(f"Encoded: {encoded_payload.decode()}")</parameter>
|
||||
</function>
|
||||
|
||||
# Long running security scan with custom timeout
|
||||
<function=python_action>
|
||||
<parameter=action>execute</parameter>
|
||||
<parameter=code>import time
|
||||
# Simulate long-running vulnerability scan
|
||||
time.sleep(45)
|
||||
print('Security scan completed!')</parameter>
|
||||
<parameter=timeout>50</parameter>
|
||||
</function>
|
||||
|
||||
# Use IPython magic commands for package management and profiling
|
||||
<function=python_action>
|
||||
<parameter=action>execute</parameter>
|
||||
<parameter=code>%pip install requests
|
||||
%time response = requests.get('https://httpbin.org/json')
|
||||
%whos</parameter>
|
||||
|
||||
# Analyze requests for potential vulnerabilities
|
||||
<function=python_action>
|
||||
<parameter=action>execute</parameter>
|
||||
<parameter=code># Filter for POST requests that might contain sensitive data
|
||||
post_requests = list_requests(
|
||||
httpql_filter="req.method.eq:POST",
|
||||
page_size=20
|
||||
)
|
||||
|
||||
# Analyze each POST request for potential issues
|
||||
for req in post_requests.get('requests', []):
|
||||
request_id = req['id']
|
||||
# View the request details
|
||||
request_details = view_request(request_id, part="request")
|
||||
|
||||
# Check for potential SQL injection points
|
||||
body = request_details.get('body', '')
|
||||
if any(keyword in body.lower() for keyword in ['select', 'union', 'insert', 'update']):
|
||||
print(f"Potential SQL injection in request {request_id}")
|
||||
|
||||
# Repeat the request with a test payload
|
||||
test_payload = repeat_request(request_id, {
|
||||
'body': body + "' OR '1'='1"
|
||||
})
|
||||
print(f"Test response status: {test_payload.get('status_code')}")
|
||||
|
||||
print("Security analysis complete!")</parameter>
|
||||
</function>
|
||||
</examples>
|
||||
</tool>
|
||||
</tools>
|
||||
172
strix/tools/python/python_instance.py
Normal file
172
strix/tools/python/python_instance.py
Normal file
@@ -0,0 +1,172 @@
|
||||
import io
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
from IPython.core.interactiveshell import InteractiveShell
|
||||
|
||||
|
||||
MAX_STDOUT_LENGTH = 10_000
|
||||
MAX_STDERR_LENGTH = 5_000
|
||||
|
||||
|
||||
class PythonInstance:
|
||||
def __init__(self, session_id: str) -> None:
|
||||
self.session_id = session_id
|
||||
self.is_running = True
|
||||
self._execution_lock = threading.Lock()
|
||||
|
||||
import os
|
||||
|
||||
os.chdir("/workspace")
|
||||
|
||||
self.shell = InteractiveShell()
|
||||
self.shell.init_completer()
|
||||
self.shell.init_history()
|
||||
self.shell.init_logger()
|
||||
|
||||
self._setup_proxy_functions()
|
||||
|
||||
def _setup_proxy_functions(self) -> None:
|
||||
try:
|
||||
from strix.tools.proxy import proxy_actions
|
||||
|
||||
proxy_functions = [
|
||||
"list_requests",
|
||||
"list_sitemap",
|
||||
"repeat_request",
|
||||
"scope_rules",
|
||||
"send_request",
|
||||
"view_request",
|
||||
"view_sitemap_entry",
|
||||
]
|
||||
|
||||
proxy_dict = {name: getattr(proxy_actions, name) for name in proxy_functions}
|
||||
self.shell.user_ns.update(proxy_dict)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
def _validate_session(self) -> dict[str, Any] | None:
|
||||
if not self.is_running:
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"stdout": "",
|
||||
"stderr": "Session is not running",
|
||||
"result": None,
|
||||
}
|
||||
return None
|
||||
|
||||
def _setup_execution_environment(self, timeout: int) -> tuple[Any, io.StringIO, io.StringIO]:
|
||||
stdout_capture = io.StringIO()
|
||||
stderr_capture = io.StringIO()
|
||||
|
||||
def timeout_handler(signum: int, frame: Any) -> None:
|
||||
raise TimeoutError(f"Code execution timed out after {timeout} seconds")
|
||||
|
||||
old_handler = signal.signal(signal.SIGALRM, timeout_handler)
|
||||
signal.alarm(timeout)
|
||||
|
||||
sys.stdout = stdout_capture
|
||||
sys.stderr = stderr_capture
|
||||
|
||||
return old_handler, stdout_capture, stderr_capture
|
||||
|
||||
def _cleanup_execution_environment(
|
||||
self, old_handler: Any, old_stdout: Any, old_stderr: Any
|
||||
) -> None:
|
||||
signal.signal(signal.SIGALRM, old_handler)
|
||||
sys.stdout = old_stdout
|
||||
sys.stderr = old_stderr
|
||||
|
||||
def _truncate_output(self, content: str, max_length: int, suffix: str) -> str:
|
||||
if len(content) > max_length:
|
||||
return content[:max_length] + suffix
|
||||
return content
|
||||
|
||||
def _format_execution_result(
|
||||
self, execution_result: Any, stdout_content: str, stderr_content: str
|
||||
) -> dict[str, Any]:
|
||||
stdout = self._truncate_output(
|
||||
stdout_content, MAX_STDOUT_LENGTH, "... [stdout truncated at 10k chars]"
|
||||
)
|
||||
|
||||
if execution_result.result is not None:
|
||||
if stdout and not stdout.endswith("\n"):
|
||||
stdout += "\n"
|
||||
result_repr = repr(execution_result.result)
|
||||
result_repr = self._truncate_output(
|
||||
result_repr, MAX_STDOUT_LENGTH, "... [result truncated at 10k chars]"
|
||||
)
|
||||
stdout += result_repr
|
||||
|
||||
stdout = self._truncate_output(
|
||||
stdout, MAX_STDOUT_LENGTH, "... [output truncated at 10k chars]"
|
||||
)
|
||||
|
||||
stderr_content = stderr_content if stderr_content else ""
|
||||
stderr_content = self._truncate_output(
|
||||
stderr_content, MAX_STDERR_LENGTH, "... [stderr truncated at 5k chars]"
|
||||
)
|
||||
|
||||
if (
|
||||
execution_result.error_before_exec or execution_result.error_in_exec
|
||||
) and not stderr_content:
|
||||
stderr_content = "Execution error occurred"
|
||||
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"stdout": stdout,
|
||||
"stderr": stderr_content,
|
||||
"result": repr(execution_result.result)
|
||||
if execution_result.result is not None
|
||||
else None,
|
||||
}
|
||||
|
||||
def _handle_execution_error(self, error: BaseException) -> dict[str, Any]:
|
||||
error_msg = str(error)
|
||||
error_msg = self._truncate_output(
|
||||
error_msg, MAX_STDERR_LENGTH, "... [error truncated at 5k chars]"
|
||||
)
|
||||
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"stdout": "",
|
||||
"stderr": error_msg,
|
||||
"result": None,
|
||||
}
|
||||
|
||||
def execute_code(self, code: str, timeout: int = 30) -> dict[str, Any]:
|
||||
session_error = self._validate_session()
|
||||
if session_error:
|
||||
return session_error
|
||||
|
||||
with self._execution_lock:
|
||||
old_stdout, old_stderr = sys.stdout, sys.stderr
|
||||
|
||||
try:
|
||||
old_handler, stdout_capture, stderr_capture = self._setup_execution_environment(
|
||||
timeout
|
||||
)
|
||||
|
||||
try:
|
||||
execution_result = self.shell.run_cell(code, silent=False, store_history=True)
|
||||
signal.alarm(0)
|
||||
|
||||
return self._format_execution_result(
|
||||
execution_result, stdout_capture.getvalue(), stderr_capture.getvalue()
|
||||
)
|
||||
|
||||
except (TimeoutError, KeyboardInterrupt, SystemExit) as e:
|
||||
signal.alarm(0)
|
||||
return self._handle_execution_error(e)
|
||||
|
||||
finally:
|
||||
self._cleanup_execution_environment(old_handler, old_stdout, old_stderr)
|
||||
|
||||
def close(self) -> None:
|
||||
self.is_running = False
|
||||
self.shell.reset(new_session=False)
|
||||
|
||||
def is_alive(self) -> bool:
|
||||
return self.is_running
|
||||
131
strix/tools/python/python_manager.py
Normal file
131
strix/tools/python/python_manager.py
Normal file
@@ -0,0 +1,131 @@
|
||||
import atexit
|
||||
import contextlib
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
from .python_instance import PythonInstance
|
||||
|
||||
|
||||
class PythonSessionManager:
|
||||
def __init__(self) -> None:
|
||||
self.sessions: dict[str, PythonInstance] = {}
|
||||
self._lock = threading.Lock()
|
||||
self.default_session_id = "default"
|
||||
|
||||
self._register_cleanup_handlers()
|
||||
|
||||
def create_session(
|
||||
self, session_id: str | None = None, initial_code: str | None = None, timeout: int = 30
|
||||
) -> dict[str, Any]:
|
||||
if session_id is None:
|
||||
session_id = self.default_session_id
|
||||
|
||||
with self._lock:
|
||||
if session_id in self.sessions:
|
||||
raise ValueError(f"Python session '{session_id}' already exists")
|
||||
|
||||
session = PythonInstance(session_id)
|
||||
self.sessions[session_id] = session
|
||||
|
||||
if initial_code:
|
||||
result = session.execute_code(initial_code, timeout)
|
||||
result["message"] = (
|
||||
f"Python session '{session_id}' created successfully with initial code"
|
||||
)
|
||||
else:
|
||||
result = {
|
||||
"session_id": session_id,
|
||||
"message": f"Python session '{session_id}' created successfully",
|
||||
}
|
||||
|
||||
return result
|
||||
|
||||
def execute_code(
|
||||
self, session_id: str | None = None, code: str | None = None, timeout: int = 30
|
||||
) -> dict[str, Any]:
|
||||
if session_id is None:
|
||||
session_id = self.default_session_id
|
||||
|
||||
if not code:
|
||||
raise ValueError("No code provided for execution")
|
||||
|
||||
with self._lock:
|
||||
if session_id not in self.sessions:
|
||||
raise ValueError(f"Python session '{session_id}' not found")
|
||||
|
||||
session = self.sessions[session_id]
|
||||
|
||||
result = session.execute_code(code, timeout)
|
||||
result["message"] = f"Code executed in session '{session_id}'"
|
||||
return result
|
||||
|
||||
def close_session(self, session_id: str | None = None) -> dict[str, Any]:
|
||||
if session_id is None:
|
||||
session_id = self.default_session_id
|
||||
|
||||
with self._lock:
|
||||
if session_id not in self.sessions:
|
||||
raise ValueError(f"Python session '{session_id}' not found")
|
||||
|
||||
session = self.sessions.pop(session_id)
|
||||
|
||||
session.close()
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"message": f"Python session '{session_id}' closed successfully",
|
||||
"is_running": False,
|
||||
}
|
||||
|
||||
def list_sessions(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
session_info = {}
|
||||
for sid, session in self.sessions.items():
|
||||
session_info[sid] = {
|
||||
"is_running": session.is_running,
|
||||
"is_alive": session.is_alive(),
|
||||
}
|
||||
|
||||
return {"sessions": session_info, "total_count": len(session_info)}
|
||||
|
||||
def cleanup_dead_sessions(self) -> None:
|
||||
with self._lock:
|
||||
dead_sessions = []
|
||||
for sid, session in self.sessions.items():
|
||||
if not session.is_alive():
|
||||
dead_sessions.append(sid)
|
||||
|
||||
for sid in dead_sessions:
|
||||
session = self.sessions.pop(sid)
|
||||
with contextlib.suppress(Exception):
|
||||
session.close()
|
||||
|
||||
def close_all_sessions(self) -> None:
|
||||
with self._lock:
|
||||
sessions_to_close = list(self.sessions.values())
|
||||
self.sessions.clear()
|
||||
|
||||
for session in sessions_to_close:
|
||||
with contextlib.suppress(Exception):
|
||||
session.close()
|
||||
|
||||
def _register_cleanup_handlers(self) -> None:
|
||||
atexit.register(self.close_all_sessions)
|
||||
|
||||
signal.signal(signal.SIGTERM, self._signal_handler)
|
||||
signal.signal(signal.SIGINT, self._signal_handler)
|
||||
|
||||
if hasattr(signal, "SIGHUP"):
|
||||
signal.signal(signal.SIGHUP, self._signal_handler)
|
||||
|
||||
def _signal_handler(self, _signum: int, _frame: Any) -> None:
|
||||
self.close_all_sessions()
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
_python_session_manager = PythonSessionManager()
|
||||
|
||||
|
||||
def get_python_session_manager() -> PythonSessionManager:
|
||||
return _python_session_manager
|
||||
Reference in New Issue
Block a user