strix-agent 0.1.1__py3-none-any.whl

strix/llm/memory_compressor.py

@@ -0,0 +1,206 @@
+import logging
+import os
+from typing import Any
+
+import litellm
+
+
+logger = logging.getLogger(__name__)
+
+
+MAX_TOTAL_TOKENS = 100_000
+MIN_RECENT_MESSAGES = 15
+
+SUMMARY_PROMPT_TEMPLATE = """You are an agent performing context
+condensation for a security agent. Your job is to compress scan data while preserving
+ALL operationally critical information for continuing the security assessment.
+
+CRITICAL ELEMENTS TO PRESERVE:
+- Discovered vulnerabilities and potential attack vectors
+- Scan results and tool outputs (compressed but maintaining key findings)
+- Access credentials, tokens, or authentication details found
+- System architecture insights and potential weak points
+- Progress made in the assessment
+- Failed attempts and dead ends (to avoid duplication)
+- Any decisions made about the testing approach
+
+COMPRESSION GUIDELINES:
+- Preserve exact technical details (URLs, paths, parameters, payloads)
+- Summarize verbose tool outputs while keeping critical findings
+- Maintain version numbers, specific technologies identified
+- Keep exact error messages that might indicate vulnerabilities
+- Compress repetitive or similar findings into consolidated form
+
+Remember: Another security agent will use this summary to continue the assessment.
+They must be able to pick up exactly where you left off without losing any
+operational advantage or context needed to find vulnerabilities.
+
+CONVERSATION SEGMENT TO SUMMARIZE:
+{conversation}
+
+Provide a technically precise summary that preserves all operational security context while
+keeping the summary concise and to the point."""
+
+
+def _count_tokens(text: str, model: str) -> int:
+    try:
+        count = litellm.token_counter(model=model, text=text)
+        return int(count)
+    except Exception:
+        logger.exception("Failed to count tokens")
+        return len(text) // 4  # Rough estimate
+
+
+def _get_message_tokens(msg: dict[str, Any], model: str) -> int:
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return _count_tokens(content, model)
+    if isinstance(content, list):
+        return sum(
+            _count_tokens(item.get("text", ""), model)
+            for item in content
+            if isinstance(item, dict) and item.get("type") == "text"
+        )
+    return 0
+
+
+def _extract_message_text(msg: dict[str, Any]) -> str:
+    content = msg.get("content", "")
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list):
+        parts = []
+        for item in content:
+            if isinstance(item, dict):
+                if item.get("type") == "text":
+                    parts.append(item.get("text", ""))
+                elif item.get("type") == "image_url":
+                    parts.append("[IMAGE]")
+        return " ".join(parts)
+
+    return str(content)
+
+
+def _summarize_messages(
+    messages: list[dict[str, Any]],
+    model: str,
+) -> dict[str, Any]:
+    if not messages:
+        empty_summary = "<context_summary message_count='0'>{text}</context_summary>"
+        return {
+            "role": "assistant",
+            "content": empty_summary.format(text="No messages to summarize"),
+        }
+
+    formatted = []
+    for msg in messages:
+        role = msg.get("role", "unknown")
+        text = _extract_message_text(msg)
+        formatted.append(f"{role}: {text}")
+
+    conversation = "\n".join(formatted)
+    prompt = SUMMARY_PROMPT_TEMPLATE.format(conversation=conversation)
+
+    try:
+        completion_args = {
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}],
+        }
+
+        response = litellm.completion(**completion_args)
+        summary = response.choices[0].message.content
+        summary_msg = "<context_summary message_count='{count}'>{text}</context_summary>"
+        return {
+            "role": "assistant",
+            "content": summary_msg.format(count=len(messages), text=summary),
+        }
+    except Exception:
+        logger.exception("Failed to summarize messages")
+        return messages[0]
+
+
+def _handle_images(messages: list[dict[str, Any]], max_images: int) -> None:
+    image_count = 0
+    for msg in reversed(messages):
+        content = msg.get("content", [])
+        if isinstance(content, list):
+            for item in content:
+                if isinstance(item, dict) and item.get("type") == "image_url":
+                    if image_count >= max_images:
+                        item.update(
+                            {
+                                "type": "text",
+                                "text": "[Previously attached image removed to preserve context]",
+                            }
+                        )
+                    else:
+                        image_count += 1
+
+
+class MemoryCompressor:
+    def __init__(
+        self,
+        max_images: int = 3,
+        model_name: str | None = None,
+    ):
+        self.max_images = max_images
+        self.model_name = model_name or os.getenv("STRIX_LLM", "anthropic/claude-sonnet-4-20250514")
+
+        if not self.model_name:
+            raise ValueError("STRIX_LLM environment variable must be set and not empty")
+
+    def compress_history(
+        self,
+        messages: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Compress conversation history to stay within token limits.
+
+        Strategy:
+        1. Handle image limits first
+        2. Keep all system messages
+        3. Keep minimum recent messages
+        4. Summarize older messages when total tokens exceed limit
+
+        The compression preserves:
+        - All system messages unchanged
+        - Most recent messages intact
+        - Critical security context in summaries
+        - Recent images for visual context
+        - Technical details and findings
+        """
+        if not messages:
+            return messages
+
+        _handle_images(messages, self.max_images)
+
+        system_msgs = []
+        regular_msgs = []
+        for msg in messages:
+            if msg.get("role") == "system":
+                system_msgs.append(msg)
+            else:
+                regular_msgs.append(msg)
+
+        recent_msgs = regular_msgs[-MIN_RECENT_MESSAGES:]
+        old_msgs = regular_msgs[:-MIN_RECENT_MESSAGES]
+
+        # Type assertion since we ensure model_name is not None in __init__
+        model_name: str = self.model_name  # type: ignore[assignment]
+
+        total_tokens = sum(
+            _get_message_tokens(msg, model_name) for msg in system_msgs + regular_msgs
+        )
+
+        if total_tokens <= MAX_TOTAL_TOKENS * 0.9:
+            return messages
+
+        compressed = []
+        chunk_size = 10
+        for i in range(0, len(old_msgs), chunk_size):
+            chunk = old_msgs[i : i + chunk_size]
+            summary = _summarize_messages(chunk, model_name)
+            if summary:
+                compressed.append(summary)
+
+        return system_msgs + compressed + recent_msgs

strix/llm/request_queue.py

@@ -0,0 +1,63 @@
+import asyncio
+import logging
+import threading
+import time
+from typing import Any
+
+from litellm import ModelResponse, completion
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+
+logger = logging.getLogger(__name__)
+
+
+class LLMRequestQueue:
+    def __init__(self, max_concurrent: int = 6, delay_between_requests: float = 1.0):
+        self.max_concurrent = max_concurrent
+        self.delay_between_requests = delay_between_requests
+        self._semaphore = threading.BoundedSemaphore(max_concurrent)
+        self._last_request_time = 0.0
+        self._lock = threading.Lock()
+
+    async def make_request(self, completion_args: dict[str, Any]) -> ModelResponse:
+        try:
+            while not self._semaphore.acquire(timeout=0.2):
+                await asyncio.sleep(0.1)
+
+            with self._lock:
+                now = time.time()
+                time_since_last = now - self._last_request_time
+                sleep_needed = max(0, self.delay_between_requests - time_since_last)
+                self._last_request_time = now + sleep_needed
+
+            if sleep_needed > 0:
+                await asyncio.sleep(sleep_needed)
+
+            return await self._reliable_request(completion_args)
+        finally:
+            self._semaphore.release()
+
+    @retry(  # type: ignore[misc]
+        stop=stop_after_attempt(15),
+        wait=wait_exponential(multiplier=1.2, min=1, max=300),
+        reraise=True,
+    )
+    async def _reliable_request(self, completion_args: dict[str, Any]) -> ModelResponse:
+        response = completion(**completion_args, stream=False)
+        if isinstance(response, ModelResponse):
+            return response
+        self._raise_unexpected_response()
+        raise RuntimeError("Unreachable code")
+
+    def _raise_unexpected_response(self) -> None:
+        raise RuntimeError("Unexpected response type")
+
+
+_global_queue: LLMRequestQueue | None = None
+
+
+def get_global_queue() -> LLMRequestQueue:
+    global _global_queue  # noqa: PLW0603
+    if _global_queue is None:
+        _global_queue = LLMRequestQueue()
+    return _global_queue

strix/llm/utils.py

@@ -0,0 +1,84 @@
+import re
+from typing import Any
+
+
+def _truncate_to_first_function(content: str) -> str:
+    if not content:
+        return content
+
+    function_starts = [match.start() for match in re.finditer(r"<function=", content)]
+
+    if len(function_starts) >= 2:
+        second_function_start = function_starts[1]
+
+        return content[:second_function_start].rstrip()
+
+    return content
+
+
+def parse_tool_invocations(content: str) -> list[dict[str, Any]] | None:
+    content = _fix_stopword(content)
+
+    tool_invocations: list[dict[str, Any]] = []
+
+    fn_regex_pattern = r"<function=([^>]+)>\n?(.*?)</function>"
+    fn_param_regex_pattern = r"<parameter=([^>]+)>(.*?)</parameter>"
+
+    fn_matches = re.finditer(fn_regex_pattern, content, re.DOTALL)
+
+    for fn_match in fn_matches:
+        fn_name = fn_match.group(1)
+        fn_body = fn_match.group(2)
+
+        param_matches = re.finditer(fn_param_regex_pattern, fn_body, re.DOTALL)
+
+        args = {}
+        for param_match in param_matches:
+            param_name = param_match.group(1)
+            param_value = param_match.group(2).strip()
+            args[param_name] = param_value
+
+        tool_invocations.append({"toolName": fn_name, "args": args})
+
+    return tool_invocations if tool_invocations else None
+
+
+def _fix_stopword(content: str) -> str:
+    if "<function=" in content and content.count("<function=") == 1:
+        if content.endswith("</"):
+            content = content.rstrip() + "function>"
+        elif not content.rstrip().endswith("</function>"):
+            content = content + "\n</function>"
+    return content
+
+
+def format_tool_call(tool_name: str, args: dict[str, Any]) -> str:
+    xml_parts = [f"<function={tool_name}>"]
+
+    for key, value in args.items():
+        xml_parts.append(f"<parameter={key}>{value}</parameter>")
+
+    xml_parts.append("</function>")
+
+    return "\n".join(xml_parts)
+
+
+def clean_content(content: str) -> str:
+    if not content:
+        return ""
+
+    content = _fix_stopword(content)
+
+    tool_pattern = r"<function=[^>]+>.*?</function>"
+    cleaned = re.sub(tool_pattern, "", content, flags=re.DOTALL)
+
+    hidden_xml_patterns = [
+        r"<inter_agent_message>.*?</inter_agent_message>",
+        r"<agent_completion_report>.*?</agent_completion_report>",
+    ]
+    for pattern in hidden_xml_patterns:
+        cleaned = re.sub(pattern, "", cleaned, flags=re.DOTALL | re.IGNORECASE)
+
+    cleaned = re.sub(r"\n\s*\n", "\n\n", cleaned)
+
+    return cleaned.strip()

strix/prompts/__init__.py

@@ -0,0 +1,113 @@
+from pathlib import Path
+
+from jinja2 import Environment
+
+
+def get_available_prompt_modules() -> dict[str, list[str]]:
+    modules_dir = Path(__file__).parent
+    available_modules = {}
+
+    for category_dir in modules_dir.iterdir():
+        if category_dir.is_dir() and not category_dir.name.startswith("__"):
+            category_name = category_dir.name
+            modules = []
+
+            for file_path in category_dir.glob("*.jinja"):
+                module_name = file_path.stem
+                modules.append(module_name)
+
+            if modules:
+                available_modules[category_name] = sorted(modules)
+
+    return available_modules
+
+
+def get_all_module_names() -> set[str]:
+    all_modules = set()
+    for category_modules in get_available_prompt_modules().values():
+        all_modules.update(category_modules)
+    return all_modules
+
+
+def validate_module_names(module_names: list[str]) -> dict[str, list[str]]:
+    available_modules = get_all_module_names()
+    valid_modules = []
+    invalid_modules = []
+
+    for module_name in module_names:
+        if module_name in available_modules:
+            valid_modules.append(module_name)
+        else:
+            invalid_modules.append(module_name)
+
+    return {"valid": valid_modules, "invalid": invalid_modules}
+
+
+def generate_modules_description() -> str:
+    available_modules = get_available_prompt_modules()
+
+    if not available_modules:
+        return "No prompt modules available"
+
+    description_parts = []
+
+    for category, modules in available_modules.items():
+        modules_str = ", ".join(modules)
+        description_parts.append(f"{category} ({modules_str})")
+
+    description = (
+        f"List of prompt modules to load for this agent (max 3). "
+        f"Available modules: {', '.join(description_parts)}. "
+    )
+
+    example_modules = []
+    for modules in available_modules.values():
+        example_modules.extend(modules[:2])
+        if len(example_modules) >= 2:
+            break
+
+    if example_modules:
+        example = f"Example: {example_modules[:2]} for specialized agent"
+        description += example
+
+    return description
+
+
+def load_prompt_modules(module_names: list[str], jinja_env: Environment) -> dict[str, str]:
+    import logging
+
+    logger = logging.getLogger(__name__)
+    module_content = {}
+    prompts_dir = Path(__file__).parent
+
+    available_modules = get_available_prompt_modules()
+
+    for module_name in module_names:
+        try:
+            module_path = None
+
+            if "/" in module_name:
+                module_path = f"{module_name}.jinja"
+            else:
+                for category, modules in available_modules.items():
+                    if module_name in modules:
+                        module_path = f"{category}/{module_name}.jinja"
+                        break
+
+                if not module_path:
+                    root_candidate = f"{module_name}.jinja"
+                    if (prompts_dir / root_candidate).exists():
+                        module_path = root_candidate
+
+            if module_path and (prompts_dir / module_path).exists():
+                template = jinja_env.get_template(module_path)
+                var_name = module_name.split("/")[-1]
+                module_content[var_name] = template.render()
+                logger.info(f"Loaded prompt module: {module_name} -> {var_name}")
+            else:
+                logger.warning(f"Prompt module not found: {module_name}")
+
+        except (FileNotFoundError, OSError, ValueError) as e:
+            logger.warning(f"Failed to load prompt module {module_name}: {e}")
+
+    return module_content

strix/prompts/coordination/root_agent.jinja

@@ -0,0 +1,41 @@
+<coordination_role>
+You are a COORDINATION AGENT ONLY. You do NOT perform any security testing, vulnerability assessment, or technical work yourself.
+
+Your ONLY responsibilities:
+1. Create specialized agents for specific security tasks
+2. Monitor agent progress and coordinate between them
+3. Compile final scan reports from agent findings
+4. Manage agent communication and dependencies
+
+CRITICAL RESTRICTIONS:
+- NEVER perform vulnerability testing or security assessments
+- NEVER write detailed vulnerability reports (only compile final summaries)
+- ONLY use agent_graph and finish tools for coordination
+- You can create agents throughout the scan process, depending on the task and findings, not just at the beginning!
+</coordination_role>
+
+<agent_management>
+BEFORE CREATING AGENTS:
+1. Analyze the target scope and break into independent tasks
+2. Check existing agents to avoid duplication
+3. Create agents with clear, specific objectives to avoid duplication
+
+AGENT TYPES YOU CAN CREATE:
+- Reconnaissance: subdomain enum, port scanning, tech identification, etc.
+- Vulnerability Testing: SQL injection, XSS, auth bypass, IDOR, RCE, SSRF, etc. Can be black-box or white-box.
+    - Direct vulnerability testing agents to implement hierarchical workflow (per finding: discover, verify, report, fix): each one should create validation agents for findings verification, which spawn reporting agents for documentation, which create fix agents for remediation
+
+COORDINATION GUIDELINES:
+- Ensure clear task boundaries and success criteria
+- Terminate redundant agents when objectives overlap
+- Use message passing for agent communication
+</agent_management>
+
+<final_responsibilities>
+When all agents complete:
+1. Collect findings from all agents
+2. Compile a final scan summary report
+3. Use finish tool to complete the assessment
+
+Your value is in orchestration, not execution.
+</final_responsibilities>

strix/prompts/vulnerabilities/authentication_jwt.jinja

@@ -0,0 +1,129 @@
+<authentication_jwt_guide>
+<title>AUTHENTICATION & JWT VULNERABILITIES</title>
+
+<critical>Authentication flaws lead to complete account takeover. JWT misconfigurations are everywhere.</critical>
+
+<jwt_structure>
+header.payload.signature
+- Header: {"alg":"HS256","typ":"JWT"}
+- Payload: {"sub":"1234","name":"John","iat":1516239022}
+- Signature: HMACSHA256(base64UrlEncode(header) + "." + base64UrlEncode(payload), secret)
+</jwt_structure>
+
+<common_attacks>
+<algorithm_confusion>
+RS256 to HS256:
+- Change RS256 to HS256 in header
+- Use public key as HMAC secret
+- Sign token with public key (often in /jwks.json or /.well-known/)
+</algorithm_confusion>
+
+<none_algorithm>
+- Set "alg": "none" in header
+- Remove signature completely (keep the trailing dot)
+</none_algorithm>
+
+<weak_secrets>
+Common secrets: 'secret', 'password', '123456', 'key', 'jwt_secret', 'your-256-bit-secret'
+</weak_secrets>
+
+<kid_manipulation>
+- SQL Injection: "kid": "key' UNION SELECT 'secret'--"
+- Command injection: "kid": "|sleep 10"
+- Path traversal: "kid": "../../../../../../dev/null"
+</kid_manipulation>
+</common_attacks>
+
+<advanced_techniques>
+<jwk_injection>
+Embed public key in token header:
+{"jwk": {"kty": "RSA", "n": "your-public-key-n", "e": "AQAB"}}
+</jwk_injection>
+
+<jku_manipulation>
+Set jku/x5u to attacker-controlled URL hosting malicious JWKS
+</jku_manipulation>
+
+<timing_attacks>
+Extract signature byte-by-byte using verification timing differences
+</timing_attacks>
+</advanced_techniques>
+
+<oauth_vulnerabilities>
+<authorization_code_theft>
+- Exploit redirect_uri with open redirects, subdomain takeover, parameter pollution
+- Missing/predictable state parameter = CSRF
+- PKCE downgrade: remove code_challenge parameter
+</authorization_code_theft>
+</oauth_vulnerabilities>
+
+<saml_attacks>
+- Signature exclusion: remove signature element
+- Signature wrapping: inject assertions
+- XXE in SAML responses
+</saml_attacks>
+
+<session_attacks>
+- Session fixation: force known session ID
+- Session puzzling: mix different session objects
+- Race conditions in session generation
+</session_attacks>
+
+<password_reset_flaws>
+- Predictable tokens: MD5(timestamp), sequential numbers
+- Host header injection for reset link poisoning
+- Race condition resets
+</password_reset_flaws>
+
+<mfa_bypass>
+- Response manipulation: change success:false to true
+- Status code manipulation: 403 to 200
+- Brute force with no rate limiting
+- Backup code abuse
+</mfa_bypass>
+
+<advanced_bypasses>
+<unicode_normalization>
+Different representations: admin@exａmple.com (fullwidth), аdmin@example.com (Cyrillic)
+</unicode_normalization>
+
+<authentication_chaining>
+- JWT + SQLi: kid parameter with SQL injection
+- OAuth + XSS: steal tokens via XSS
+- SAML + XXE + SSRF: chain for internal access
+</authentication_chaining>
+</advanced_bypasses>
+
+<tools>
+- jwt_tool: Comprehensive JWT testing
+- Check endpoints: /login, /oauth/authorize, /saml/login, /.well-known/openid-configuration, /jwks.json
+</tools>
+
+<validation>
+To confirm authentication flaw:
+1. Demonstrate account access without credentials
+2. Show privilege escalation
+3. Prove token forgery works
+4. Bypass authentication/2FA requirements
+5. Maintain persistent access
+</validation>
+
+<false_positives>
+NOT a vulnerability if:
+- Requires valid credentials
+- Only affects own session
+- Proper signature validation
+- Token expiration enforced
+- Rate limiting prevents brute force
+</false_positives>
+
+<impact>
+- Account takeover: access other users' accounts
+- Privilege escalation: user to admin
+- Token forgery: create valid tokens
+- Bypass mechanisms: skip auth/2FA
+- Persistent access: survives logout
+</impact>
+
+<remember>Focus on RS256->HS256, weak secrets, and none algorithm first. Modern apps use multiple auth methods simultaneously - find gaps in integration.</remember>
+</authentication_jwt_guide>