strix-agent 0.4.0__py3-none-any.whl

strix/tools/terminal/terminal_session.py

@@ -0,0 +1,447 @@
+import logging
+import re
+import time
+import uuid
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+import libtmux
+
+
+logger = logging.getLogger(__name__)
+
+
+class BashCommandStatus(Enum):
+    CONTINUE = "continue"
+    COMPLETED = "completed"
+    NO_CHANGE_TIMEOUT = "no_change_timeout"
+    HARD_TIMEOUT = "hard_timeout"
+
+
+def _remove_command_prefix(command_output: str, command: str) -> str:
+    return command_output.lstrip().removeprefix(command.lstrip()).lstrip()
+
+
+class TerminalSession:
+    POLL_INTERVAL = 0.5
+    HISTORY_LIMIT = 10_000
+    PS1_END = "]$ "
+
+    def __init__(self, session_id: str, work_dir: str = "/workspace") -> None:
+        self.session_id = session_id
+        self.work_dir = str(Path(work_dir).resolve())
+        self._closed = False
+        self._cwd = self.work_dir
+
+        self.server: libtmux.Server | None = None
+        self.session: libtmux.Session | None = None
+        self.window: libtmux.Window | None = None
+        self.pane: libtmux.Pane | None = None
+
+        self.prev_status: BashCommandStatus | None = None
+        self.prev_output: str = ""
+        self._initialized = False
+
+        self.initialize()
+
+    @property
+    def PS1(self) -> str:  # noqa: N802
+        return r"[STRIX_$?]$ "
+
+    @property
+    def PS1_PATTERN(self) -> str:  # noqa: N802
+        return r"\[STRIX_(\d+)\]"
+
+    def initialize(self) -> None:
+        self.server = libtmux.Server()
+
+        session_name = f"strix-{self.session_id}-{uuid.uuid4()}"
+        self.session = self.server.new_session(
+            session_name=session_name,
+            start_directory=self.work_dir,
+            kill_session=True,
+            x=120,
+            y=30,
+        )
+
+        self.session.set_option("history-limit", str(self.HISTORY_LIMIT))
+        self.session.history_limit = self.HISTORY_LIMIT
+
+        _initial_window = self.session.active_window
+        self.window = self.session.new_window(
+            window_name="bash",
+            window_shell="/bin/bash",
+            start_directory=self.work_dir,
+        )
+        self.pane = self.window.active_pane
+        _initial_window.kill()
+
+        self.pane.send_keys(f'export PROMPT_COMMAND=\'export PS1="{self.PS1}"\'; export PS2=""')
+        time.sleep(0.1)
+        self._clear_screen()
+
+        self.prev_status = None
+        self.prev_output = ""
+        self._closed = False
+
+        self._cwd = str(Path(self.work_dir).resolve())
+        self._initialized = True
+
+        assert self.server is not None
+        assert self.session is not None
+        assert self.window is not None
+        assert self.pane is not None
+
+    def _get_pane_content(self) -> str:
+        if not self.pane:
+            raise RuntimeError("Terminal session not properly initialized")
+        return "\n".join(
+            line.rstrip() for line in self.pane.cmd("capture-pane", "-J", "-pS", "-").stdout
+        )
+
+    def _clear_screen(self) -> None:
+        if not self.pane:
+            raise RuntimeError("Terminal session not properly initialized")
+        self.pane.send_keys("C-l", enter=False)
+        time.sleep(0.1)
+        self.pane.cmd("clear-history")
+
+    def _is_control_key(self, command: str) -> bool:
+        return (
+            (command.startswith("C-") and len(command) >= 3)
+            or (command.startswith("^") and len(command) >= 2)
+            or (command.startswith("S-") and len(command) >= 3)
+            or (command.startswith("M-") and len(command) >= 3)
+        )
+
+    def _is_function_key(self, command: str) -> bool:
+        if not command.startswith("F") or len(command) > 3:
+            return False
+        try:
+            num_part = command[1:]
+            return num_part.isdigit() and 1 <= int(num_part) <= 12
+        except (ValueError, IndexError):
+            return False
+
+    def _is_navigation_or_special_key(self, command: str) -> bool:
+        navigation_keys = {"Up", "Down", "Left", "Right", "Home", "End"}
+        special_keys = {"BSpace", "BTab", "DC", "Enter", "Escape", "IC", "Space", "Tab"}
+        page_keys = {"NPage", "PageDown", "PgDn", "PPage", "PageUp", "PgUp"}
+
+        return command in navigation_keys or command in special_keys or command in page_keys
+
+    def _is_complex_modifier_key(self, command: str) -> bool:
+        return "-" in command and any(
+            command.startswith(prefix)
+            for prefix in ["C-S-", "C-M-", "S-M-", "M-S-", "M-C-", "S-C-"]
+        )
+
+    def _is_special_key(self, command: str) -> bool:
+        _command = command.strip()
+
+        if not _command:
+            return False
+
+        return (
+            self._is_control_key(_command)
+            or self._is_function_key(_command)
+            or self._is_navigation_or_special_key(_command)
+            or self._is_complex_modifier_key(_command)
+        )
+
+    def _matches_ps1_metadata(self, content: str) -> list[re.Match[str]]:
+        return list(re.finditer(self.PS1_PATTERN + r"\]\$ ", content))
+
+    def _get_command_output(
+        self,
+        command: str,
+        raw_command_output: str,
+        continue_prefix: str = "",
+    ) -> str:
+        if self.prev_output:
+            command_output = raw_command_output.removeprefix(self.prev_output)
+            if continue_prefix:
+                command_output = continue_prefix + command_output
+        else:
+            command_output = raw_command_output
+        self.prev_output = raw_command_output
+        command_output = _remove_command_prefix(command_output, command)
+        return command_output.rstrip()
+
+    def _combine_outputs_between_matches(
+        self,
+        pane_content: str,
+        ps1_matches: list[re.Match[str]],
+        get_content_before_last_match: bool = False,
+    ) -> str:
+        if len(ps1_matches) == 1:
+            if get_content_before_last_match:
+                return pane_content[: ps1_matches[0].start()]
+            return pane_content[ps1_matches[0].end() + 1 :]
+        if len(ps1_matches) == 0:
+            return pane_content
+
+        combined_output = ""
+        for i in range(len(ps1_matches) - 1):
+            output_segment = pane_content[ps1_matches[i].end() + 1 : ps1_matches[i + 1].start()]
+            combined_output += output_segment + "\n"
+        combined_output += pane_content[ps1_matches[-1].end() + 1 :]
+        return combined_output
+
+    def _extract_exit_code_from_matches(self, ps1_matches: list[re.Match[str]]) -> int | None:
+        if not ps1_matches:
+            return None
+
+        last_match = ps1_matches[-1]
+        try:
+            return int(last_match.group(1))
+        except (ValueError, IndexError):
+            return None
+
+    def _handle_empty_command(
+        self,
+        cur_pane_output: str,
+        ps1_matches: list[re.Match[str]],
+        is_command_running: bool,
+        timeout: float,
+    ) -> dict[str, Any]:
+        if not is_command_running:
+            raw_command_output = self._combine_outputs_between_matches(cur_pane_output, ps1_matches)
+            command_output = self._get_command_output("", raw_command_output)
+            return {
+                "content": command_output,
+                "status": "completed",
+                "exit_code": 0,
+                "working_dir": self._cwd,
+            }
+
+        start_time = time.time()
+        last_pane_output = cur_pane_output
+
+        while True:
+            cur_pane_output = self._get_pane_content()
+            ps1_matches = self._matches_ps1_metadata(cur_pane_output)
+
+            if cur_pane_output.rstrip().endswith(self.PS1_END.rstrip()) or len(ps1_matches) > 0:
+                exit_code = self._extract_exit_code_from_matches(ps1_matches)
+                raw_command_output = self._combine_outputs_between_matches(
+                    cur_pane_output, ps1_matches
+                )
+                command_output = self._get_command_output("", raw_command_output)
+                self.prev_status = BashCommandStatus.COMPLETED
+                self.prev_output = ""
+                self._ready_for_next_command()
+                return {
+                    "content": command_output,
+                    "status": "completed",
+                    "exit_code": exit_code or 0,
+                    "working_dir": self._cwd,
+                }
+
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                raw_command_output = self._combine_outputs_between_matches(
+                    cur_pane_output, ps1_matches
+                )
+                command_output = self._get_command_output("", raw_command_output)
+                return {
+                    "content": command_output
+                    + f"\n[Command still running after {timeout}s - showing output so far]",
+                    "status": "running",
+                    "exit_code": None,
+                    "working_dir": self._cwd,
+                }
+
+            if cur_pane_output != last_pane_output:
+                last_pane_output = cur_pane_output
+
+            time.sleep(self.POLL_INTERVAL)
+
+    def _handle_input_command(
+        self, command: str, no_enter: bool, is_command_running: bool
+    ) -> dict[str, Any]:
+        if not is_command_running:
+            return {
+                "content": "No command is currently running. Cannot send input.",
+                "status": "error",
+                "exit_code": None,
+                "working_dir": self._cwd,
+            }
+
+        if not self.pane:
+            raise RuntimeError("Terminal session not properly initialized")
+
+        is_special_key = self._is_special_key(command)
+        should_add_enter = not is_special_key and not no_enter
+        self.pane.send_keys(command, enter=should_add_enter)
+
+        time.sleep(2)
+        cur_pane_output = self._get_pane_content()
+        ps1_matches = self._matches_ps1_metadata(cur_pane_output)
+        raw_command_output = self._combine_outputs_between_matches(cur_pane_output, ps1_matches)
+        command_output = self._get_command_output(command, raw_command_output)
+
+        is_still_running = not (
+            cur_pane_output.rstrip().endswith(self.PS1_END.rstrip()) or len(ps1_matches) > 0
+        )
+
+        if is_still_running:
+            return {
+                "content": command_output,
+                "status": "running",
+                "exit_code": None,
+                "working_dir": self._cwd,
+            }
+
+        exit_code = self._extract_exit_code_from_matches(ps1_matches)
+        self.prev_status = BashCommandStatus.COMPLETED
+        self.prev_output = ""
+        self._ready_for_next_command()
+        return {
+            "content": command_output,
+            "status": "completed",
+            "exit_code": exit_code or 0,
+            "working_dir": self._cwd,
+        }
+
+    def _execute_new_command(self, command: str, no_enter: bool, timeout: float) -> dict[str, Any]:
+        if not self.pane:
+            raise RuntimeError("Terminal session not properly initialized")
+
+        initial_pane_output = self._get_pane_content()
+        initial_ps1_matches = self._matches_ps1_metadata(initial_pane_output)
+        initial_ps1_count = len(initial_ps1_matches)
+
+        start_time = time.time()
+        last_pane_output = initial_pane_output
+
+        is_special_key = self._is_special_key(command)
+        should_add_enter = not is_special_key and not no_enter
+        self.pane.send_keys(command, enter=should_add_enter)
+
+        while True:
+            cur_pane_output = self._get_pane_content()
+            ps1_matches = self._matches_ps1_metadata(cur_pane_output)
+            current_ps1_count = len(ps1_matches)
+
+            if cur_pane_output != last_pane_output:
+                last_pane_output = cur_pane_output
+
+            if current_ps1_count > initial_ps1_count or cur_pane_output.rstrip().endswith(
+                self.PS1_END.rstrip()
+            ):
+                exit_code = self._extract_exit_code_from_matches(ps1_matches)
+
+                get_content_before_last_match = bool(len(ps1_matches) == 1)
+                raw_command_output = self._combine_outputs_between_matches(
+                    cur_pane_output,
+                    ps1_matches,
+                    get_content_before_last_match=get_content_before_last_match,
+                )
+
+                command_output = self._get_command_output(command, raw_command_output)
+                self.prev_status = BashCommandStatus.COMPLETED
+                self.prev_output = ""
+                self._ready_for_next_command()
+
+                return {
+                    "content": command_output,
+                    "status": "completed",
+                    "exit_code": exit_code or 0,
+                    "working_dir": self._cwd,
+                }
+
+            elapsed_time = time.time() - start_time
+            if elapsed_time >= timeout:
+                raw_command_output = self._combine_outputs_between_matches(
+                    cur_pane_output, ps1_matches
+                )
+                command_output = self._get_command_output(
+                    command,
+                    raw_command_output,
+                    continue_prefix="[Below is the output of the previous command.]\n",
+                )
+                self.prev_status = BashCommandStatus.CONTINUE
+
+                timeout_msg = (
+                    f"\n[Command still running after {timeout}s - showing output so far. "
+                    "Use C-c to interrupt if needed.]"
+                )
+                return {
+                    "content": command_output + timeout_msg,
+                    "status": "running",
+                    "exit_code": None,
+                    "working_dir": self._cwd,
+                }
+
+            time.sleep(self.POLL_INTERVAL)
+
+    def execute(
+        self, command: str, is_input: bool = False, timeout: float = 10.0, no_enter: bool = False
+    ) -> dict[str, Any]:
+        if not self._initialized:
+            raise RuntimeError("Bash session is not initialized")
+
+        cur_pane_output = self._get_pane_content()
+        ps1_matches = self._matches_ps1_metadata(cur_pane_output)
+        is_command_running = not (
+            cur_pane_output.rstrip().endswith(self.PS1_END.rstrip()) or len(ps1_matches) > 0
+        )
+
+        if command.strip() == "":
+            return self._handle_empty_command(
+                cur_pane_output, ps1_matches, is_command_running, timeout
+            )
+
+        is_special_key = self._is_special_key(command)
+
+        if is_input:
+            return self._handle_input_command(command, no_enter, is_command_running)
+
+        if is_special_key and is_command_running:
+            return self._handle_input_command(command, no_enter, is_command_running)
+
+        if is_command_running:
+            return {
+                "content": (
+                    "A command is already running. Use is_input=true to send input to it, "
+                    "or interrupt it first (e.g., with C-c)."
+                ),
+                "status": "error",
+                "exit_code": None,
+                "working_dir": self._cwd,
+            }
+
+        return self._execute_new_command(command, no_enter, timeout)
+
+    def _ready_for_next_command(self) -> None:
+        self._clear_screen()
+
+    def is_running(self) -> bool:
+        if self._closed or not self.session:
+            return False
+        try:
+            return self.session.id in [s.id for s in self.server.sessions] if self.server else False
+        except (AttributeError, OSError) as e:
+            logger.debug("Error checking if session is running: %s", e)
+            return False
+
+    def get_working_dir(self) -> str:
+        return self._cwd
+
+    def close(self) -> None:
+        if self._closed:
+            return
+
+        if self.session:
+            try:
+                self.session.kill()
+            except (AttributeError, OSError) as e:
+                logger.debug("Error closing terminal session: %s", e)
+
+        self._closed = True
+        self.server = None
+        self.session = None
+        self.window = None
+        self.pane = None

strix/tools/thinking/__init__.py

@@ -0,0 +1,4 @@
+from .thinking_actions import think
+
+
+__all__ = ["think"]

strix/tools/thinking/thinking_actions.py

@@ -0,0 +1,18 @@
+from typing import Any
+
+from strix.tools.registry import register_tool
+
+
+@register_tool(sandbox_execution=False)
+def think(thought: str) -> dict[str, Any]:
+    try:
+        if not thought or not thought.strip():
+            return {"success": False, "message": "Thought cannot be empty"}
+
+        return {
+            "success": True,
+            "message": f"Thought recorded successfully with {len(thought.strip())} characters",
+        }
+
+    except (ValueError, TypeError) as e:
+        return {"success": False, "message": f"Failed to record thought: {e!s}"}

strix/tools/thinking/thinking_actions_schema.xml

@@ -0,0 +1,52 @@
+<tools>
+  <tool name="think">
+    <description>Use the tool to think about something. It will not obtain new information or change the
+  database. Use it when complex reasoning or some cache memory is needed.</description>
+    <details>This tool creates dedicated space for structured thinking during complex tasks,
+  particularly useful for:
+  - Tool output analysis: When you need to carefully process the output of previous tool calls
+  - Policy-heavy environments: When you need to follow detailed guidelines and verify compliance
+  - Sequential decision making: When each action builds on previous ones and mistakes are costly
+  - Multi-step problem solving: When you need to break down complex problems into manageable steps</details>
+    <parameters>
+      <parameter name="thought" type="string" required="true">
+        <description>The thought or reasoning to record</description>
+      </parameter>
+    </parameters>
+    <returns type="Dict[str, Any]">
+      <description>Response containing: - success: Whether the thought was recorded successfully - message: Confirmation message with character count or error details</description>
+    </returns>
+    <examples>
+  # Planning and strategy
+  <function=think>
+  <parameter=thought>I need to analyze the scan results systematically. First, let me review
+  the open ports: 22 (SSH), 80 (HTTP), 443 (HTTPS), and 3306 (MySQL). The MySQL port being
+  externally accessible is a high priority security concern. I should check for default
+  credentials and version information. For the web services, I need to enumerate
+  directories and test for common web vulnerabilities.</parameter>
+  </function>
+
+  # Analysis of tool outputs
+  <function=think>
+  <parameter=thought>The Nmap scan revealed 15 open ports, but three stand out as concerning:
+  - Port 3306 (MySQL) - Database should not be exposed externally
+  - Port 5432 (PostgreSQL) - Another database port that's risky when public
+  - Port 6379 (Redis) - Often misconfigured and can lead to data exposure
+  I should prioritize testing these database services for authentication bypass and
+  information disclosure vulnerabilities.</parameter>
+  </function>
+
+  # Decision making and next steps
+  <function=think>
+  <parameter=thought>Based on the vulnerability scan results, I've identified several critical
+  issues that need immediate attention:
+  1. SQL injection in the login form (confirmed with sqlmap)
+  2. Reflected XSS in the search parameter
+  3. Directory traversal in the file upload function
+  I should document these findings with proof-of-concept exploits and assign appropriate
+  CVSS scores. The SQL injection poses the highest risk due to potential data
+  exfiltration.</parameter>
+  </function>
+    </examples>
+  </tool>
+</tools>

strix/tools/web_search/__init__.py

@@ -0,0 +1,4 @@
+from .web_search_actions import web_search
+
+
+__all__ = ["web_search"]

strix/tools/web_search/web_search_actions.py

@@ -0,0 +1,80 @@
+import os
+from typing import Any
+
+import requests
+
+from strix.tools.registry import register_tool
+
+
+SYSTEM_PROMPT = """You are assisting a cybersecurity agent specialized in vulnerability scanning
+and security assessment running on Kali Linux. When responding to search queries:
+
+1. Prioritize cybersecurity-relevant information including:
+   - Vulnerability details (CVEs, CVSS scores, impact)
+   - Security tools, techniques, and methodologies
+   - Exploit information and proof-of-concepts
+   - Security best practices and mitigations
+   - Penetration testing approaches
+   - Web application security findings
+
+2. Provide technical depth appropriate for security professionals
+3. Include specific versions, configurations, and technical details when available
+4. Focus on actionable intelligence for security assessment
+5. Cite reliable security sources (NIST, OWASP, CVE databases, security vendors)
+6. When providing commands or installation instructions, prioritize Kali Linux compatibility
+   and use apt package manager or tools pre-installed in Kali
+7. Be detailed and specific - avoid general answers. Always include concrete code examples,
+   command-line instructions, configuration snippets, or practical implementation steps
+   when applicable
+
+Structure your response to be comprehensive yet concise, emphasizing the most critical
+security implications and details."""
+
+
+@register_tool(sandbox_execution=False)
+def web_search(query: str) -> dict[str, Any]:
+    try:
+        api_key = os.getenv("PERPLEXITY_API_KEY")
+        if not api_key:
+            return {
+                "success": False,
+                "message": "PERPLEXITY_API_KEY environment variable not set",
+                "results": [],
+            }
+
+        url = "https://api.perplexity.ai/chat/completions"
+        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+
+        payload = {
+            "model": "sonar-reasoning",
+            "messages": [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": query},
+            ],
+        }
+
+        response = requests.post(url, headers=headers, json=payload, timeout=300)
+        response.raise_for_status()
+
+        response_data = response.json()
+        content = response_data["choices"][0]["message"]["content"]
+
+    except requests.exceptions.Timeout:
+        return {"success": False, "message": "Request timed out", "results": []}
+    except requests.exceptions.RequestException as e:
+        return {"success": False, "message": f"API request failed: {e!s}", "results": []}
+    except KeyError as e:
+        return {
+            "success": False,
+            "message": f"Unexpected API response format: missing {e!s}",
+            "results": [],
+        }
+    except Exception as e:  # noqa: BLE001
+        return {"success": False, "message": f"Web search failed: {e!s}", "results": []}
+    else:
+        return {
+            "success": True,
+            "query": query,
+            "content": content,
+            "message": "Web search completed successfully",
+        }