strix-agent 0.1.1__py3-none-any.whl

strix/tools/python/__init__.py

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

strix/tools/python/python_actions.py

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

strix/tools/python/python_actions_schema.xml

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

strix/tools/python/python_instance.py

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

strix/tools/python/python_manager.py

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

strix/tools/registry.py

@@ -0,0 +1,196 @@
+import inspect
+import logging
+import os
+from collections.abc import Callable
+from functools import wraps
+from inspect import signature
+from pathlib import Path
+from typing import Any
+
+
+tools: list[dict[str, Any]] = []
+_tools_by_name: dict[str, Callable[..., Any]] = {}
+logger = logging.getLogger(__name__)
+
+
+class ImplementedInClientSideOnlyError(Exception):
+    def __init__(
+        self,
+        message: str = "This tool is implemented in the client side only",
+    ) -> None:
+        self.message = message
+        super().__init__(self.message)
+
+
+def _process_dynamic_content(content: str) -> str:
+    if "{{DYNAMIC_MODULES_DESCRIPTION}}" in content:
+        try:
+            from strix.prompts import generate_modules_description
+
+            modules_description = generate_modules_description()
+            content = content.replace("{{DYNAMIC_MODULES_DESCRIPTION}}", modules_description)
+        except ImportError:
+            logger.warning("Could not import prompts utilities for dynamic schema generation")
+            content = content.replace(
+                "{{DYNAMIC_MODULES_DESCRIPTION}}",
+                "List of prompt modules to load for this agent (max 3). Module discovery failed.",
+            )
+
+    return content
+
+
+def _load_xml_schema(path: Path) -> Any:
+    if not path.exists():
+        return None
+    try:
+        content = path.read_text()
+
+        content = _process_dynamic_content(content)
+
+        start_tag = '<tool name="'
+        end_tag = "</tool>"
+        tools_dict = {}
+
+        pos = 0
+        while True:
+            start_pos = content.find(start_tag, pos)
+            if start_pos == -1:
+                break
+
+            name_start = start_pos + len(start_tag)
+            name_end = content.find('"', name_start)
+            if name_end == -1:
+                break
+            tool_name = content[name_start:name_end]
+
+            end_pos = content.find(end_tag, name_end)
+            if end_pos == -1:
+                break
+            end_pos += len(end_tag)
+
+            tool_element = content[start_pos:end_pos]
+            tools_dict[tool_name] = tool_element
+
+            pos = end_pos
+
+            if pos >= len(content):
+                break
+    except (IndexError, ValueError, UnicodeError) as e:
+        logger.warning(f"Error loading schema file {path}: {e}")
+        return None
+    else:
+        return tools_dict
+
+
+def _get_module_name(func: Callable[..., Any]) -> str:
+    module = inspect.getmodule(func)
+    if not module:
+        return "unknown"
+
+    module_name = module.__name__
+    if ".tools." in module_name:
+        parts = module_name.split(".tools.")[-1].split(".")
+        if len(parts) >= 1:
+            return parts[0]
+    return "unknown"
+
+
+def register_tool(
+    func: Callable[..., Any] | None = None, *, sandbox_execution: bool = True
+) -> Callable[..., Any]:
+    def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        func_dict = {
+            "name": f.__name__,
+            "function": f,
+            "module": _get_module_name(f),
+            "sandbox_execution": sandbox_execution,
+        }
+
+        sandbox_mode = os.getenv("STRIX_SANDBOX_MODE", "false").lower() == "true"
+        if not sandbox_mode:
+            try:
+                module_path = Path(inspect.getfile(f))
+                schema_file_name = f"{module_path.stem}_schema.xml"
+                schema_path = module_path.parent / schema_file_name
+
+                xml_tools = _load_xml_schema(schema_path)
+
+                if xml_tools is not None and f.__name__ in xml_tools:
+                    func_dict["xml_schema"] = xml_tools[f.__name__]
+                else:
+                    func_dict["xml_schema"] = (
+                        f'<tool name="{f.__name__}">'
+                        "<description>Schema not found for tool.</description>"
+                        "</tool>"
+                    )
+            except (TypeError, FileNotFoundError) as e:
+                logger.warning(f"Error loading schema for {f.__name__}: {e}")
+                func_dict["xml_schema"] = (
+                    f'<tool name="{f.__name__}">'
+                    "<description>Error loading schema.</description>"
+                    "</tool>"
+                )
+
+        tools.append(func_dict)
+        _tools_by_name[str(func_dict["name"])] = f
+
+        @wraps(f)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return f(*args, **kwargs)
+
+        return wrapper
+
+    if func is None:
+        return decorator
+    return decorator(func)
+
+
+def get_tool_by_name(name: str) -> Callable[..., Any] | None:
+    return _tools_by_name.get(name)
+
+
+def get_tool_names() -> list[str]:
+    return list(_tools_by_name.keys())
+
+
+def needs_agent_state(tool_name: str) -> bool:
+    tool_func = get_tool_by_name(tool_name)
+    if not tool_func:
+        return False
+    sig = signature(tool_func)
+    return "agent_state" in sig.parameters
+
+
+def should_execute_in_sandbox(tool_name: str) -> bool:
+    for tool in tools:
+        if tool.get("name") == tool_name:
+            return bool(tool.get("sandbox_execution", True))
+    return True
+
+
+def get_tools_prompt() -> str:
+    tools_by_module: dict[str, list[dict[str, Any]]] = {}
+    for tool in tools:
+        module = tool.get("module", "unknown")
+        if module not in tools_by_module:
+            tools_by_module[module] = []
+        tools_by_module[module].append(tool)
+
+    xml_sections = []
+    for module, module_tools in sorted(tools_by_module.items()):
+        tag_name = f"{module}_tools"
+        section_parts = [f"<{tag_name}>"]
+        for tool in module_tools:
+            tool_xml = tool.get("xml_schema", "")
+            if tool_xml:
+                indented_tool = "\n".join(f"  {line}" for line in tool_xml.split("\n"))
+                section_parts.append(indented_tool)
+        section_parts.append(f"</{tag_name}>")
+        xml_sections.append("\n".join(section_parts))
+
+    return "\n\n".join(xml_sections)
+
+
+def clear_registry() -> None:
+    tools.clear()
+    _tools_by_name.clear()

strix/tools/reporting/__init__.py

@@ -0,0 +1,6 @@
+from .reporting_actions import create_vulnerability_report
+
+
+__all__ = [
+    "create_vulnerability_report",
+]