Open-source release for Alpha version

This commit is contained in:
Ahmed Allam
2025-08-08 20:36:44 -07:00
commit 81ac98e8b9
105 changed files with 22125 additions and 0 deletions

View File

@@ -0,0 +1,4 @@
from .python_actions import python_action
__all__ = ["python_action"]

View 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}

View 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>

View 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

View 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