From d21bc45702eb5c03c7042a9e1a408f88b99d5895 Mon Sep 17 00:00:00 2001 From: yokoszn Date: Sun, 23 Nov 2025 04:31:21 +1100 Subject: [PATCH 1/3] feat: add scope configuration system and role-based tool enforcement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scope Configuration System: - Add strix/scope/ module with Pydantic models for scope configuration - Support YAML/JSON scope files with networks, targets, exclusions, domains - Add --scope, --filter, --validate CLI arguments - Inject scope_context and exclusion_rules into agent task descriptions - Add scope awareness section to root_agent.jinja prompt Role-Based Tool Enforcement: - Add TOOL_PROFILES with 6 roles: root, recon, testing, validation, reporting, fixing - Implement is_tool_allowed_for_role() in registry.py - Add runtime enforcement in executor.py validate_tool_availability() - Track agent_role in AgentState and LLMConfig Additional Features: - Add Proxmox VE prompt module (proxmox_ve.jinja) - Add documentation: SCOPE_CONFIGURATION.md, PROMPTING_GUIDE.md - Add scope templates: scope.yaml, scope.json, proxmox-cluster.yaml - Add development configs: Dockerfile.dev, docker-compose.dev.yml, .env.example - Add cache-implementation-project.md planning document πŸ€– Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- .env.example | 15 + Dockerfile.dev | 41 ++ cache-implementation-project.md | 429 ++++++++++++++++ docker-compose.dev.yml | 21 + docs/PROMPTING_GUIDE.md | 421 ++++++++++++++++ docs/SCOPE_CONFIGURATION.md | 407 ++++++++++++++++ strix/agents/StrixAgent/strix_agent.py | 47 +- strix/agents/state.py | 1 + strix/interface/cli.py | 5 + strix/interface/main.py | 190 +++++++- strix/interface/tui.py | 9 +- strix/llm/config.py | 2 + strix/llm/llm.py | 7 +- strix/prompts/coordination/root_agent.jinja | 388 ++++++++++++++- strix/prompts/technologies/proxmox_ve.jinja | 274 +++++++++++ strix/scope/__init__.py | 34 ++ strix/scope/models.py | 167 +++++++ strix/scope/parser.py | 457 ++++++++++++++++++ strix/scope/validator.py | 214 ++++++++ strix/tools/__init__.py | 4 + .../agents_graph/agents_graph_actions.py | 30 +- .../agents_graph_actions_schema.xml | 45 +- strix/tools/executor.py | 16 +- strix/tools/registry.py | 115 ++++- templates/scope/proxmox-cluster.yaml | 102 ++++ templates/scope/scope-simple.csv | 8 + templates/scope/scope.json | 86 ++++ templates/scope/scope.yaml | 100 ++++ 28 files changed, 3563 insertions(+), 72 deletions(-) create mode 100644 .env.example create mode 100644 Dockerfile.dev create mode 100644 cache-implementation-project.md create mode 100644 docker-compose.dev.yml create mode 100644 docs/PROMPTING_GUIDE.md create mode 100644 docs/SCOPE_CONFIGURATION.md create mode 100644 strix/prompts/technologies/proxmox_ve.jinja create mode 100644 strix/scope/__init__.py create mode 100644 strix/scope/models.py create mode 100644 strix/scope/parser.py create mode 100644 strix/scope/validator.py create mode 100644 templates/scope/proxmox-cluster.yaml create mode 100644 templates/scope/scope-simple.csv create mode 100644 templates/scope/scope.json create mode 100644 templates/scope/scope.yaml diff --git a/.env.example b/.env.example new file mode 100644 index 000000000..a3be81e7a --- /dev/null +++ b/.env.example @@ -0,0 +1,15 @@ +# Strix Configuration +# Copy this file to .env and fill in your values + +# LLM Provider Configuration +# Format: provider/model (e.g., openai/gpt-4o, anthropic/claude-3-5-sonnet, zhipu/glm-4) +STRIX_LLM=zhipu/glm-4 + +# Your API key from the LLM provider +LLM_API_KEY=your-api-key-here + +# Optional: Custom API base URL (for local models or proxies) +# LLM_API_BASE=http://localhost:8000/v1 + +# Optional: Request timeout in seconds (default: 600) +# LLM_TIMEOUT=600 diff --git a/Dockerfile.dev b/Dockerfile.dev new file mode 100644 index 000000000..2d0eb18e0 --- /dev/null +++ b/Dockerfile.dev @@ -0,0 +1,41 @@ +FROM python:3.12-slim + +LABEL description="Strix Development Environment" + +# Install system dependencies and Docker CLI +RUN apt-get update && apt-get install -y --no-install-recommends \ + git \ + curl \ + build-essential \ + ca-certificates \ + gnupg \ + && install -m 0755 -d /etc/apt/keyrings \ + && curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg \ + && chmod a+r /etc/apt/keyrings/docker.gpg \ + && echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/debian bookworm stable" > /etc/apt/sources.list.d/docker.list \ + && apt-get update \ + && apt-get install -y --no-install-recommends docker-ce-cli \ + && rm -rf /var/lib/apt/lists/* + +# Install Poetry +RUN curl -sSL https://install.python-poetry.org | python3 - && \ + ln -s /root/.local/bin/poetry /usr/local/bin/poetry + +# Set working directory +WORKDIR /app + +# Copy dependency files first for layer caching +COPY pyproject.toml poetry.lock ./ + +# Install dependencies (no dev deps for running, add --with=dev if needed) +RUN poetry config virtualenvs.create false && \ + poetry install --no-root --no-interaction + +# Copy the rest of the application +COPY . . + +# Install the package itself +RUN poetry install --no-interaction + +# Default command +ENTRYPOINT ["strix"] diff --git a/cache-implementation-project.md b/cache-implementation-project.md new file mode 100644 index 000000000..3b22007e9 --- /dev/null +++ b/cache-implementation-project.md @@ -0,0 +1,429 @@ +# Cache Implementation Project + +## Overview + +This document outlines the implementation of a shared cache layer for Strix's multi-agent system to enable real-time context synchronization, shared mutable state, and improved agent coordination. + +## Current State + +### How Agents Share Data Today +- **Context inheritance**: One-time copy from parent to child at agent creation +- **Message passing**: Serial `send_message_to_agent()` for communication +- **Completion reports**: XML reports via `agent_finish()` when child completes +- **Shared filesystem**: All agents read/write `/workspace` directory +- **Shared proxy**: Caido proxy history visible to all agents + +### Limitations +- No real-time state synchronization between agents +- Parent doesn't see child's discoveries until completion +- No shared variable space for live collaboration +- Agents can duplicate work without knowing what others found +- No centralized findings registry during scan + +--- + +## Proposed Architecture + +### Option A: Redis (Recommended for Production) + +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ Docker Network β”‚ +β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ +β”‚ β”‚ Root β”‚ β”‚ Recon β”‚ β”‚ Testing β”‚ β”‚ Validation β”‚ β”‚ +β”‚ β”‚ Agent β”‚ β”‚ Agent β”‚ β”‚ Agent β”‚ β”‚ Agent β”‚ β”‚ +β”‚ β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”˜ β”‚ +β”‚ β”‚ β”‚ β”‚ β”‚ β”‚ +β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ +β”‚ β”‚ β”‚ +β”‚ β”Œβ”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β” β”‚ +β”‚ β”‚ Redis β”‚ β”‚ +β”‚ β”‚ Cache β”‚ β”‚ +β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +**Pros:** +- Production-ready, battle-tested +- Pub/sub for real-time agent notifications +- TTL support for automatic cleanup +- Persistence options if needed +- Excellent Python support (redis-py, aioredis) + +**Cons:** +- Additional service to manage +- Overkill for single-machine deployments +- Network overhead (minimal) + +### Option B: SQLite with WAL Mode + +```python +# Lightweight, no additional services +# Good for single-machine, moderate concurrency + +import sqlite3 +conn = sqlite3.connect('file:strix_cache?mode=memory&cache=shared', uri=True) +conn.execute('PRAGMA journal_mode=WAL') +``` + +**Pros:** +- Zero additional dependencies (stdlib) +- No separate service +- Works in shared memory mode +- Familiar SQL interface + +**Cons:** +- No pub/sub (requires polling) +- Write contention under high concurrency +- Not designed for cache use case + +### Option C: In-Memory Dict with File-Backed Sync (MVP) + +```python +# Start simple, evolve later +# Uses existing shared /workspace for persistence + +class AgentCache: + def __init__(self, workspace_path="/workspace/.strix_cache"): + self._local = {} + self._cache_file = workspace_path + self._lock = threading.Lock() + + def set(self, key: str, value: Any, namespace: str = "global") -> None: + ... + + def get(self, key: str, namespace: str = "global") -> Any: + ... + + def publish(self, channel: str, message: dict) -> None: + # Write to file-based message queue + ... +``` + +**Pros:** +- No new dependencies +- Easy to implement +- Uses existing /workspace sharing +- Good enough for MVP + +**Cons:** +- File I/O overhead +- No true pub/sub +- Scaling limitations +- Race conditions possible without careful locking + +--- + +## Recommended Phased Approach + +### Phase 1: File-Backed Shared State (MVP) + +**Goal:** Enable agents to share findings in real-time without waiting for completion. + +**Implementation:** +``` +/workspace/.strix/ +β”œβ”€β”€ cache/ +β”‚ β”œβ”€β”€ findings.json # Deduplicated findings registry +β”‚ β”œβ”€β”€ targets.json # Discovered targets/endpoints +β”‚ └── state.json # Scan-level state +β”œβ”€β”€ messages/ +β”‚ β”œβ”€β”€ {agent_id}/ # Per-agent message inbox +β”‚ β”‚ └── *.json # Individual messages +β”‚ └── broadcast/ # Broadcast messages +└── locks/ + └── *.lock # File-based locking +``` + +**New Tools:** +```python +@register_tool(sandbox_execution=False) +def cache_set(agent_state: Any, key: str, value: str, namespace: str = "global") -> dict: + """Store a value in the shared agent cache.""" + ... + +@register_tool(sandbox_execution=False) +def cache_get(agent_state: Any, key: str, namespace: str = "global") -> dict: + """Retrieve a value from the shared agent cache.""" + ... + +@register_tool(sandbox_execution=False) +def register_finding(agent_state: Any, finding: dict) -> dict: + """Register a finding in the shared findings registry (auto-deduplicates).""" + ... + +@register_tool(sandbox_execution=False) +def get_findings(agent_state: Any, severity: str = None, category: str = None) -> dict: + """Query registered findings from all agents.""" + ... +``` + +**Effort:** 1-2 days +**Risk:** Low + +### Phase 2: Redis Integration (Production) + +**Goal:** Replace file-backed cache with Redis for better performance and pub/sub. + +**Implementation:** +- Add redis container to docker-compose +- Create `strix/cache/redis_cache.py` adapter +- Add pub/sub for agent notifications +- Implement cache namespacing per scan + +**New Capabilities:** +```python +# Real-time notifications +await cache.subscribe("findings", on_new_finding) + +# Atomic operations +await cache.increment("stats:endpoints_discovered") + +# Automatic expiry +await cache.set("temp:scan_token", token, ttl=3600) +``` + +**Effort:** 2-3 days +**Risk:** Medium (new dependency) + +### Phase 3: Advanced Features + +**Goal:** Sophisticated multi-agent coordination. + +**Features:** +- Distributed locking for exclusive access +- Agent presence/heartbeat tracking +- Scan-level metrics aggregation +- Finding deduplication with similarity scoring +- Dependency graph for agent task ordering + +--- + +## Data Structures + +### Shared Findings Registry + +```json +{ + "findings": { + "finding_abc123": { + "id": "finding_abc123", + "title": "SQL Injection in login endpoint", + "severity": "critical", + "category": "injection", + "target": "/api/login", + "discovered_by": "agent_sqli_001", + "discovered_at": "2024-01-15T10:30:00Z", + "validated": false, + "validated_by": null, + "poc_available": true, + "hash": "sha256:...", // For deduplication + "related_findings": ["finding_def456"] + } + }, + "index": { + "by_severity": { + "critical": ["finding_abc123"], + "high": [], + "medium": [], + "low": [], + "info": [] + }, + "by_target": { + "/api/login": ["finding_abc123"] + }, + "by_agent": { + "agent_sqli_001": ["finding_abc123"] + } + } +} +``` + +### Agent State Registry + +```json +{ + "agents": { + "agent_root_001": { + "id": "agent_root_001", + "name": "Root Coordinator", + "role": "root", + "status": "running", + "task": "Coordinate Proxmox security assessment", + "started_at": "2024-01-15T10:00:00Z", + "last_heartbeat": "2024-01-15T10:35:00Z", + "children": ["agent_recon_001", "agent_recon_002"], + "findings_count": 0, + "current_phase": "reconnaissance" + } + }, + "hierarchy": { + "agent_root_001": { + "agent_recon_001": { + "agent_sqli_001": {}, + "agent_xss_001": {} + }, + "agent_recon_002": {} + } + } +} +``` + +### Target/Endpoint Registry + +```json +{ + "targets": { + "https://10.0.101.2:8006": { + "type": "infrastructure", + "service": "proxmox-ve", + "discovered_by": "agent_recon_001", + "endpoints": [ + { + "path": "/api2/json/access/ticket", + "method": "POST", + "params": ["username", "password", "realm"], + "auth_required": false, + "tested_by": [] + } + ], + "technologies": ["Proxmox VE 8.x", "pveproxy"], + "open_ports": [22, 8006, 3128] + } + } +} +``` + +--- + +## Integration Points + +### 1. Agent Creation +```python +# In agents_graph_actions.py create_agent() +async def create_agent(...): + ... + # Initialize agent's cache namespace + cache = get_scan_cache(scan_id) + cache.register_agent(state.agent_id, { + "name": name, + "role": agent_role, + "task": task, + "parent": parent_id + }) +``` + +### 2. Finding Discovery +```python +# In vulnerability testing agents +# Instead of just sending message to parent: +cache.register_finding({ + "title": "SQL Injection found", + "severity": "critical", + ... +}) +# All agents can now query this immediately +``` + +### 3. Deduplication +```python +# Before reporting a finding +existing = cache.find_similar_finding(new_finding) +if existing: + # Link as related, don't duplicate + cache.link_findings(existing["id"], new_finding["id"]) +else: + cache.register_finding(new_finding) +``` + +### 4. Real-Time Coordination (Phase 2+) +```python +# Agent subscribes to relevant channels +await cache.subscribe("findings:critical", handle_critical_finding) +await cache.subscribe(f"agent:{agent_id}:messages", handle_message) + +# Coordinator can broadcast +await cache.publish("broadcast", {"type": "pause", "reason": "user requested"}) +``` + +--- + +## Migration Path + +### From Current System +1. **No breaking changes** - Cache is additive +2. Existing message passing continues to work +3. Cache provides optional enhancement +4. Gradual adoption by updating prompt modules + +### Deprecation Timeline +- Phase 1: Both systems coexist +- Phase 2: Prefer cache for findings, messages for commands +- Phase 3: Consider deprecating direct message passing for data sharing + +--- + +## Performance Considerations + +### File-Based (Phase 1) +- Read: ~1-5ms (SSD) +- Write: ~5-20ms (with fsync) +- Acceptable for 10-50 agents +- Bottleneck: Write contention on findings.json + +### Redis (Phase 2) +- Read: ~0.1-0.5ms +- Write: ~0.1-0.5ms +- Supports 1000+ agents +- Bottleneck: Network (negligible on localhost) + +### Recommendations +- Phase 1: Use write-behind caching (batch writes) +- Phase 2: Use Redis pipelining for bulk operations +- All phases: Namespace by scan_id to isolate concurrent scans + +--- + +## Security Considerations + +1. **Cache Poisoning**: Validate all cache entries before use +2. **Information Leakage**: Clear cache between scans +3. **Denial of Service**: Implement size limits per namespace +4. **Access Control**: In Phase 2+, consider per-agent permissions + +--- + +## Open Questions + +1. **Persistence**: Should findings survive container restart? + - Current: No (container-scoped) + - Consider: Optional persistence for long scans + +2. **Multi-Scan Isolation**: How to handle concurrent scans? + - Proposal: Namespace everything by `scan_id` + +3. **Cache Invalidation**: When should cached data expire? + - Proposal: TTL based on data type (findings: never, temp state: 1 hour) + +4. **Conflict Resolution**: What if two agents find the same vuln? + - Proposal: First-write-wins with similarity linking + +--- + +## Next Steps + +1. [ ] Review and approve Phase 1 design +2. [ ] Implement file-backed cache in `strix/cache/` +3. [ ] Add cache tools to tool registry +4. [ ] Update root_agent.jinja to document cache usage +5. [ ] Test with multi-agent Proxmox scan +6. [ ] Evaluate need for Phase 2 based on performance + +--- + +## References + +- [Redis Pub/Sub](https://redis.io/docs/manual/pubsub/) +- [SQLite Shared Cache](https://www.sqlite.org/sharedcache.html) +- [Python threading locks](https://docs.python.org/3/library/threading.html#lock-objects) +- [File locking in Python](https://docs.python.org/3/library/fcntl.html#fcntl.flock) diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml new file mode 100644 index 000000000..ada40c2df --- /dev/null +++ b/docker-compose.dev.yml @@ -0,0 +1,21 @@ +services: + strix: + build: + context: . + dockerfile: Dockerfile.dev + env_file: + - .env + volumes: + # Mount source code for live editing + - ./strix:/app/strix:ro + # Mount target directory (change as needed) + - ./target:/workspace:ro + # Docker socket for spawning sandbox containers + - /var/run/docker.sock:/var/run/docker.sock + environment: + - STRIX_LLM=${STRIX_LLM:-openai/gpt-4o} + - LLM_API_KEY=${LLM_API_KEY} + - DOCKER_HOST=unix:///var/run/docker.sock + stdin_open: true + tty: true + network_mode: host diff --git a/docs/PROMPTING_GUIDE.md b/docs/PROMPTING_GUIDE.md new file mode 100644 index 000000000..8830bd57a --- /dev/null +++ b/docs/PROMPTING_GUIDE.md @@ -0,0 +1,421 @@ +# Strix Prompting Guide + +## Overview + +Strix operates with a multi-agent architecture where a **root coordinator** orchestrates specialized child agents. Understanding how to prompt Strix effectively is key to getting useful results. + +This guide covers both interface modes: +- **TUI Mode** (default): Interactive terminal UI for real-time collaboration +- **CLI Mode** (`-n`): Headless/non-interactive for automation and CI/CD + +--- + +## Architecture Summary + +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ ROOT COORDINATOR β”‚ +β”‚ Role: root | Tools: create_agent, messaging, think β”‚ +β”‚ Does NOT perform security testing directly β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ creates + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β–Ό β–Ό β–Ό +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ Recon β”‚ β”‚ Testing β”‚ β”‚ Validationβ”‚ +β”‚ Agent β”‚ β”‚ Agent β”‚ β”‚ Agent β”‚ +β”‚ role:reconβ”‚ β”‚role:testingβ”‚ β”‚role:valid β”‚ +β””β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ β”‚ + β–Ό β–Ό + Children Children +``` + +**Key Points:** +- Root coordinator parses your instructions and delegates work +- Child agents have specialized tools based on their role +- Each agent can load up to 5 prompt modules for domain expertise +- Runtime role enforcement prevents agents from using unauthorized tools + +--- + +## Operational Modes + +The root coordinator detects your intended mode from keywords or natural language: + +| Mode | Keywords | Behavior | +|------|----------|----------| +| **RECON-ONLY** | `recon only`, `reconnaissance only`, `no exploitation`, `passive` | Discovery only, generates PoCs but doesn't execute exploits | +| **POC-ONLY** | `poc only`, `proof of concept only` | Discovery + validation, runs PoCs in sandbox, no active exploitation | +| **FULL PENTEST** | `full pentest`, `full test`, `exploitation allowed` | Complete testing including active exploitation within scope | +| **DEFAULT** | (no keywords) | Interprets intent; defaults to POC-ONLY if unclear | + +--- + +## TUI Mode (Interactive) + +### Starting a Scan + +```bash +# Basic scan (TUI mode is default) +strix --target 10.0.101.2 + +# Short form +strix -t 10.0.101.2 + +# With initial instructions +strix --target https://example.com --instruction "Focus on authentication" + +# Multiple targets +strix --target ./local-code --target https://staging.example.com + +# With custom run name +strix --target example.com --run-name "my-pentest-run" +``` + +### During the Scan + +In TUI mode, you can interact with the agent in real-time: + +**Providing Additional Context:** +``` +The admin panel is at /admin and uses basic auth +``` + +**Adjusting Scope:** +``` +Skip the /api/v1/legacy endpoints, they're deprecated +``` + +**Answering Agent Questions:** +When the root coordinator needs clarification (e.g., before destructive actions), it will ask. Simply respond in the chat. + +**Stopping:** +- `ESC` - Stop the current agent gracefully +- `Ctrl+C` - Quit and save partial results + +### Effective TUI Prompts + +**Good: Specific with context** +``` +Proxmox VE 8.1 server. Test API authentication at port 8006. +Recon only, no destructive actions. I have root credentials if needed: root / [redacted] +``` + +**Good: Phased approach** +``` +Start with reconnaissance. After you show me the findings, I'll decide what to test further. +``` + +**Less Effective: Vague** +``` +hack it +``` + +--- + +## CLI Mode (Headless) + +### Starting a Scan + +```bash +# Non-interactive mode +strix --target 10.0.101.2 -n --instruction "recon only, generate PoCs" + +# With custom run name +strix --target https://api.example.com -n \ + --instruction "full pentest, focus on IDOR and auth bypass" \ + --run-name "api-pentest-2024-01" +``` + +### Instruction Design for CLI + +Since there's no interaction, your `--instruction` must be comprehensive: + +**Template:** +``` +[TARGET CONTEXT] + [OPERATIONAL MODE] + [FOCUS AREAS] + [CONSTRAINTS] +``` + +**Examples:** + +```bash +# Infrastructure recon +--instruction "Proxmox VE server at 10.0.101.2:8006. \ +Recon only, no exploitation. Generate PoCs for any CVEs found. \ +Focus on API authentication and VM escape vectors." + +# Web application pentest +--instruction "E-commerce Django app. Full pentest. \ +Focus on payment flow, session management, and IDOR. \ +Test credentials: testuser@example.com / TestPass123" + +# Code review +--instruction "Python FastAPI backend. Static analysis only. \ +Look for SQL injection, auth bypasses, and hardcoded secrets. \ +Critical paths: /api/auth/*, /api/admin/*" +``` + +### Exit Codes + +| Code | Meaning | +|------|---------| +| 0 | Scan completed, no vulnerabilities | +| 1 | Error during scan | +| 2 | Scan completed, vulnerabilities found | + +--- + +## Instruction Keywords Reference + +### Operational Mode Keywords + +| Keyword | Effect | +|---------|--------| +| `recon only` | Reconnaissance only, no exploitation | +| `poc only` | Generate and validate PoCs, no active exploitation | +| `full pentest` | Full testing including exploitation | +| `passive` | Same as recon only | +| `no exploitation` | Same as recon only | + +### Safety Keywords + +| Keyword | Effect | +|---------|--------| +| `no destructive actions` | Extra confirmation before any state changes | +| `read only` | No writes to target system | +| `safe mode` | Maximum caution, ask before each action | + +### Focus Keywords + +The agent recognizes technology and vulnerability type mentions: + +**Technologies:** +- `Proxmox`, `VMware`, `ESXi` β†’ Infrastructure testing focus +- `AWS`, `Azure`, `GCP`, `kubernetes`, `k8s` β†’ Cloud testing focus +- `Django`, `FastAPI`, `Express`, `Next.js` β†’ Framework-specific modules +- `Firebase`, `Supabase`, `Auth0` β†’ Service-specific modules + +**Vulnerability Types:** +- `SQL injection`, `SQLi` β†’ sql_injection module +- `XSS`, `cross-site scripting` β†’ xss module +- `authentication`, `auth`, `JWT` β†’ authentication_jwt module +- `IDOR`, `insecure direct object` β†’ idor module +- `SSRF`, `server-side request` β†’ ssrf module +- `file upload` β†’ insecure_file_uploads module + +--- + +## Target Type Detection + +The agent auto-detects target types: + +| Input | Detected As | +|-------|-------------| +| `https://example.com` | Web Application | +| `192.168.1.1` | Infrastructure | +| `10.0.101.2:8006` | Infrastructure (Proxmox likely) | +| `https://github.com/user/repo` | Repository | +| `./my-project` | Local Code | +| `s3://bucket-name` | Cloud (AWS) | +| `app.apk` | Mobile | + +You can override or clarify: +``` +--instruction "This IP runs a web application, not infrastructure" +``` + +--- + +## Multi-Model Consensus + +For high-confidence assessments, request consensus validation: + +```bash +--instruction "Critical production API. Use multi-model consensus for all high/critical findings." +``` + +This spawns advisor agents with different LLM models to independently validate findings. + +--- + +## Agent Roles and Tools + +| Role | Purpose | Key Tools | +|------|---------|-----------| +| `root` | Coordination only | create_agent, view_agent_graph, finish_scan, send_message_to_agent, wait_for_message, think | +| `recon` | Discovery and enumeration | terminal, python, browser, proxy, think, agent_finish, create_agent, view_agent_graph, send_message_to_agent, wait_for_message, read_file, write_file, list_directory, web_search | +| `testing` | Vulnerability testing | terminal, python, browser, proxy, think, agent_finish, create_agent, view_agent_graph, send_message_to_agent, wait_for_message, read_file, write_file, web_search | +| `validation` | PoC validation | terminal, python, browser, proxy, think, agent_finish, read_file, send_message_to_agent | +| `reporting` | Report generation | create_vulnerability_report, read_file, write_file, think, agent_finish, send_message_to_agent | +| `fixing` | Code remediation | read_file, write_file, terminal, python, think, agent_finish, send_message_to_agent | + +Runtime enforcement prevents agents from using tools outside their role. + +--- + +## Prompt Modules + +Agents can load specialized knowledge modules (max 5 per agent). Available modules: + +**Vulnerabilities:** +| Module | Use Case | +|--------|----------| +| `sql_injection` | SQL injection testing techniques | +| `xss` | Cross-site scripting testing | +| `ssrf` | Server-side request forgery | +| `xxe` | XML external entity injection | +| `rce` | Remote code execution | +| `csrf` | Cross-site request forgery | +| `idor` | Insecure direct object references | +| `authentication_jwt` | JWT and auth mechanism testing | +| `business_logic` | Business logic flaw testing | +| `insecure_file_uploads` | File upload vulnerabilities | +| `path_traversal_lfi_rfi` | Path traversal and file inclusion | +| `race_conditions` | Race condition vulnerabilities | +| `mass_assignment` | Mass assignment flaws | +| `broken_function_level_authorization` | Authorization bypass | + +**Technologies:** +| Module | Use Case | +|--------|----------| +| `proxmox_ve` | Proxmox VE infrastructure testing | +| `firebase_firestore` | Firebase security testing | +| `supabase` | Supabase security testing | + +**Frameworks:** +| Module | Use Case | +|--------|----------| +| `fastapi` | FastAPI application testing | +| `nextjs` | Next.js application testing | + +**Protocols:** +| Module | Use Case | +|--------|----------| +| `graphql` | GraphQL API testing | + +**Coordination:** +| Module | Use Case | +|--------|----------| +| `root_agent` | Root coordinator behavior (auto-loaded for root) | + +Modules are auto-selected based on target and focus areas, or explicitly requested. + +--- + +## Output and Reports + +Results are saved to `agent_runs//`: + +``` +agent_runs/ +└── scan-2024-01-15-abc123/ + β”œβ”€β”€ penetration_test_report.md # Final penetration test report + β”œβ”€β”€ vulnerabilities.csv # Vulnerability index (id, title, severity, timestamp) + └── vulnerabilities/ # Individual vulnerability reports + β”œβ”€β”€ vuln-abc123.md + └── vuln-def456.md +``` + +### Vulnerability Report Format + +Each vulnerability in `vulnerabilities/` contains: +```markdown +# [Vulnerability Title] + +**ID:** vuln-abc123 +**Severity:** CRITICAL +**Found:** 2024-01-15T10:30:00Z + +## Description + +[Detailed vulnerability description and PoC] +``` + +--- + +## Examples + +### Running with Docker Compose (Development) + +If using the development Docker setup: + +```bash +# Build the container +docker compose -f docker-compose.dev.yml build + +# Run with docker compose (add arguments after 'strix') +docker compose -f docker-compose.dev.yml run --rm strix \ + --target 10.0.101.2 \ + --instruction "your instructions here" +``` + +### Example 1: Proxmox Recon (CLI) + +```bash +strix --target 10.0.101.2 -n \ + --instruction "Proxmox VE server. Recon only, no destructive actions. \ +Generate PoCs for any vulnerabilities found. Focus on API auth and known CVEs." +``` + +### Example 2: Web App Full Pentest (TUI) + +```bash +strix --target https://staging.myapp.com \ + --instruction "Full pentest. Django backend with React frontend. \ +Test credentials: admin@test.com / AdminTest123" +``` + +Then interact in TUI to guide testing. + +### Example 3: Multi-Target White-Box (CLI) + +```bash +strix \ + --target ./backend \ + --target https://api.staging.myapp.com \ + -n \ + --instruction "White-box test. Code in ./backend matches deployed API. \ +Focus on auth endpoints and data validation. Generate fixes for confirmed vulns." +``` + +### Example 4: Cloud Infrastructure (TUI) + +```bash +strix --target arn:aws:s3:::my-bucket \ + --instruction "AWS environment audit. Check S3 permissions, IAM policies. \ +Recon only." +``` + +--- + +## Troubleshooting + +### Agent Not Doing What I Expected + +1. Check if operational mode was detected correctly +2. Be more explicit with keywords: `recon only`, `full pentest` +3. In TUI mode, provide clarification when asked + +### Scan Taking Too Long + +1. Narrow the scope: `Focus only on /api/auth/*` +2. Use recon-only mode first, then targeted testing +3. Specify what to skip: `Skip /static and /assets` + +### No Vulnerabilities Found + +1. Provide credentials if needed +2. Specify entry points: `Start with the login form at /login` +3. Mention known weak areas: `The password reset flow might be vulnerable` + +--- + +## Best Practices + +1. **Start with recon** - Understand the target before deep testing +2. **Be specific** - Vague prompts lead to unfocused testing +3. **Provide context** - Technology stack, credentials, known issues +4. **Use TUI for exploration** - CLI for automation +5. **Review partial results** - In TUI, guide based on findings +6. **Set clear boundaries** - What's in/out of scope diff --git a/docs/SCOPE_CONFIGURATION.md b/docs/SCOPE_CONFIGURATION.md new file mode 100644 index 000000000..ee9d55312 --- /dev/null +++ b/docs/SCOPE_CONFIGURATION.md @@ -0,0 +1,407 @@ +# Scope Configuration + +## Overview + +Strix supports structured scope configuration via YAML or JSON files, enabling complex engagements with multiple networks, VLANs, and mixed internal/external targets. The scope system provides: + +- **Network definitions** with CIDR ranges and VLAN info +- **Target metadata** including services, credentials, and focus areas +- **Exclusion rules** for hosts, URLs, ports, and paths +- **Domain boundaries** with wildcard support +- **Operational mode** control (recon-only, poc-only, full-pentest) + +## Quick Start + +```bash +# Run with scope file +strix --scope scope.yaml + +# Validate scope file without running +strix --scope scope.yaml --validate + +# Filter targets by criteria +strix --scope scope.yaml --filter "tags:critical" +strix --scope scope.yaml --filter "network:DMZ" + +# Combine scope with manual targets +strix --scope scope.yaml --target https://extra-target.com +``` + +--- + +## Scope File Formats + +### YAML (Recommended) + +```yaml +# scope.yaml +metadata: + engagement_name: "Acme Corp Pentest 2024" + engagement_type: "internal" # internal | external | hybrid + start_date: "2024-01-15" + end_date: "2024-01-30" + tester: "security-team" + +# Global settings +settings: + operational_mode: "poc-only" # recon-only | poc-only | full-pentest + max_agents: 20 + require_validation: true + generate_fixes: false + +# Network scope definitions +networks: + - name: "Corporate LAN" + type: "internal" + vlan: 10 + cidr: "10.0.10.0/24" + gateway: "10.0.10.1" + description: "Main corporate network" + + - name: "Server VLAN" + type: "internal" + vlan: 101 + cidr: "10.0.101.0/24" + gateway: "10.0.101.1" + description: "Production servers" + + - name: "DMZ" + type: "external" + cidr: "203.0.113.0/24" + description: "Public-facing servers" + +# Specific targets within scope +targets: + # Infrastructure targets + - host: "10.0.101.2" + name: "Proxmox Host 1" + type: "infrastructure" + network: "Server VLAN" + ports: [22, 8006] + services: + - port: 8006 + service: "proxmox-ve" + version: "8.1" + credentials: + - username: "root" + password_env: "PROXMOX_ROOT_PASS" # Reference env var + access_level: "admin" + tags: ["hypervisor", "critical"] + modules: ["proxmox_ve"] + + # Web application targets + - url: "https://app.acme.com" + name: "Customer Portal" + type: "web_application" + network: "DMZ" + technologies: ["Django", "PostgreSQL", "Redis"] + credentials: + - username: "testuser@acme.com" + password_env: "PORTAL_TEST_PASS" + access_level: "user" + focus_areas: ["authentication", "idor", "business_logic"] + tags: ["customer-facing", "pii"] + + # API targets + - url: "https://api.acme.com" + name: "REST API" + type: "api" + network: "DMZ" + auth_type: "bearer" + token_env: "API_BEARER_TOKEN" + openapi_spec: "./specs/api-v2.yaml" + tags: ["api", "critical"] + + # Code repositories + - repo: "https://github.com/acme/backend" + name: "Backend Codebase" + type: "repository" + branch: "main" + focus_areas: ["sql_injection", "authentication_jwt", "secrets"] + + - path: "./frontend" + name: "Frontend Codebase" + type: "local_code" + focus_areas: ["xss", "csrf"] + +# Exclusions - DO NOT TEST +exclusions: + hosts: + - "10.0.101.1" # Gateway + - "10.0.101.254" # Network monitoring + cidrs: + - "10.0.102.0/24" # Out of scope network + urls: + - "https://app.acme.com/health" + - "https://app.acme.com/metrics" + paths: + - "/api/v1/legacy/*" + - "/admin/dangerous/*" + ports: + - 161 # SNMP - don't touch + - 162 + +# Domain scope for web testing +domains: + in_scope: + - "*.acme.com" + - "*.acme-staging.com" + out_of_scope: + - "mail.acme.com" + - "vpn.acme.com" +``` + +### JSON + +```json +{ + "metadata": { + "engagement_name": "Acme Corp Pentest 2024", + "engagement_type": "internal" + }, + "settings": { + "operational_mode": "poc-only" + }, + "networks": [ + { + "name": "Server VLAN", + "type": "internal", + "vlan": 101, + "cidr": "10.0.101.0/24" + } + ], + "targets": [ + { + "host": "10.0.101.2", + "name": "Proxmox Host 1", + "type": "infrastructure", + "ports": [22, 8006] + } + ], + "exclusions": { + "hosts": ["10.0.101.1"], + "cidrs": ["10.0.102.0/24"] + } +} +``` + +--- + +## CLI Integration + +### Basic Usage + +```bash +# Load scope from YAML file +strix --scope scope.yaml + +# Load scope from JSON file +strix --scope scope.json + +# Validate scope file and exit +strix --scope scope.yaml --validate +``` + +### Filtering Targets + +Filter targets from the scope file by various criteria: + +```bash +# Filter by tags (comma-separated for multiple) +strix --scope scope.yaml --filter "tags:critical" +strix --scope scope.yaml --filter "tags:hypervisor,pii" + +# Filter by network name +strix --scope scope.yaml --filter "network:DMZ" +strix --scope scope.yaml --filter "network:Server VLAN" + +# Filter by target type +strix --scope scope.yaml --filter "type:infrastructure" +strix --scope scope.yaml --filter "type:web_application" + +# Combine multiple filters +strix --scope scope.yaml --filter "tags:critical" --filter "network:DMZ" +``` + +### Combining with Manual Targets + +Scope files can be combined with manual `--target` arguments: + +```bash +# Scope file + additional target +strix --scope scope.yaml --target https://extra-target.com + +# Scope file + custom instructions +strix --scope scope.yaml --instruction "Focus on authentication vulnerabilities" +``` + +--- + +## Validation + +The scope validator checks for: + +1. **Target validity** + - At least one identifier (host, url, repo, or path) + - Valid IP addresses for host fields + - Valid port ranges (1-65535) + +2. **Network validity** + - No duplicate network names + - Valid VLAN IDs (1-4094) + - Valid CIDR notation + - Gateway within CIDR range (warning if not) + +3. **Exclusion validity** + - Valid CIDR notation + - Valid port ranges + - Warnings for overlaps between targets and exclusions + +4. **Credential security** + - Warnings for missing environment variables + +5. **Reference integrity** + - Target network references must exist + - Module references validated against available modules + +### Validation Output + +```bash +$ strix --scope scope.yaml --validate + +Scope file is valid: scope.yaml + +Warnings: + - pve-node-01: Environment variable not set: PVE_ROOT_PASS + +Scope Summary: + Engagement: Proxmox Cluster Assessment + Type: internal + Mode: recon-only + Networks: 1 + Targets: 3 +``` + +--- + +## How Scope Affects Agent Behavior + +When a scope file is loaded, the root agent receives: + +1. **Scope Context** - Engagement metadata, settings, and network definitions +2. **Exclusion Rules** - Hosts, CIDRs, URLs, paths, and ports to avoid + +### Operational Mode + +The scope's `operational_mode` setting controls agent behavior: + +| Mode | Behavior | +|------|----------| +| `recon-only` | Reconnaissance only, no exploitation, PoCs generated but not executed | +| `poc-only` | Discovery and PoC validation in sandbox, no active exploitation | +| `full-pentest` | Full testing including exploitation within scope boundaries | + +### Scope Boundaries + +Agents automatically respect: + +- **In-scope**: Targets listed in scope, IPs within network CIDRs, domains matching `in_scope` patterns +- **Excluded**: Hosts, CIDRs, URLs, paths, ports, and domains in exclusion lists +- **Out-of-scope**: Domains matching `out_of_scope` patterns + +--- + +## Target Definition Reference + +### Common Fields + +| Field | Type | Description | +|-------|------|-------------| +| `name` | string | Human-readable name | +| `type` | string | `infrastructure`, `web_application`, `api`, `repository`, `local_code` | +| `network` | string | Reference to network definition | +| `tags` | list | Arbitrary tags for filtering | +| `focus_areas` | list | Vulnerability types to prioritize | +| `modules` | list | Prompt modules to load | + +### Target Identifiers (one required) + +| Field | Type | Description | +|-------|------|-------------| +| `host` | string | IP address | +| `url` | string | Web URL | +| `repo` | string | Git repository URL | +| `path` | string | Local filesystem path | + +### Infrastructure Fields + +| Field | Type | Description | +|-------|------|-------------| +| `ports` | list[int] | Open ports | +| `services` | list | Service definitions with port, service, version | +| `credentials` | list | Credential definitions | + +### Web/API Fields + +| Field | Type | Description | +|-------|------|-------------| +| `technologies` | list | Known tech stack | +| `auth_type` | string | `bearer`, `basic`, `api_key` | +| `token_env` | string | Environment variable for auth token | +| `openapi_spec` | string | Path to OpenAPI spec file | + +### Repository Fields + +| Field | Type | Description | +|-------|------|-------------| +| `branch` | string | Git branch to analyze | + +--- + +## Templates + +Pre-built scope templates are available in `templates/scope/`: + +- `scope.yaml` - Full YAML template with all options +- `scope.json` - JSON template +- `scope-simple.csv` - Simple CSV for target lists +- `proxmox-cluster.yaml` - Proxmox cluster assessment + +--- + +## Future: Database Integration + +The scope parser is designed with future SQLite/Redis integration in mind: + +```python +# Current: File-based +scope = ScopeConfig.from_file("scope.yaml") + +# Future: Database-backed +scope = ScopeConfig.from_dict(db.get_scope(engagement_id)) +db.save_scope(engagement_id, scope.to_dict()) + +# Future: Redis cache +scope = ScopeConfig.from_dict(redis.get_json("scope:12345")) +``` + +The `to_dict()` and `from_dict()` methods enable serialization to any backend. See `docs/CACHE_IMPLEMENTATION.md` for the planned caching architecture. + +--- + +## Module Structure + +``` +strix/scope/ +β”œβ”€β”€ __init__.py # Public exports +β”œβ”€β”€ models.py # Pydantic models for scope configuration +β”œβ”€β”€ parser.py # ScopeConfig class with parsing and conversion +└── validator.py # ScopeValidator with validation rules +``` + +### Key Classes + +- **ScopeConfig**: Main class for loading, validating, and querying scope +- **ScopeConfigModel**: Pydantic model for scope structure +- **ScopeValidator**: Validates scope for correctness and security +- **ValidationResult**: Contains errors and warnings from validation diff --git a/strix/agents/StrixAgent/strix_agent.py b/strix/agents/StrixAgent/strix_agent.py index 81f4886e8..d65542669 100644 --- a/strix/agents/StrixAgent/strix_agent.py +++ b/strix/agents/StrixAgent/strix_agent.py @@ -9,18 +9,26 @@ class StrixAgent(BaseAgent): def __init__(self, config: dict[str, Any]): default_modules = [] + agent_role = None state = config.get("state") if state is None or (hasattr(state, "parent_id") and state.parent_id is None): default_modules = ["root_agent"] + agent_role = "root" - self.default_llm_config = LLMConfig(prompt_modules=default_modules) + self.default_llm_config = LLMConfig(prompt_modules=default_modules, agent_role=agent_role) super().__init__(config) + # Set role on state for runtime enforcement (state may be created by BaseAgent) + if agent_role and hasattr(self.state, "agent_role"): + self.state.agent_role = agent_role + async def execute_scan(self, scan_config: dict[str, Any]) -> dict[str, Any]: # noqa: PLR0912 user_instructions = scan_config.get("user_instructions", "") targets = scan_config.get("targets", []) + scope_context = scan_config.get("scope_context") + exclusion_rules = scan_config.get("exclusion_rules") repositories = [] local_code = [] @@ -86,4 +94,41 @@ async def execute_scan(self, scan_config: dict[str, Any]) -> dict[str, Any]: # if user_instructions: task_description += f"\n\nSpecial instructions: {user_instructions}" + # Inject scope context for agent awareness + if scope_context: + task_description += "\n\n" + task_description += f"\nEngagement: {scope_context.get('engagement', {}).get('name', 'Unknown')}" + task_description += f"\nType: {scope_context.get('engagement', {}).get('type', 'Unknown')}" + task_description += f"\nMode: {scope_context.get('settings', {}).get('mode', 'poc-only')}" + task_description += f"\nTargets in scope: {scope_context.get('target_count', 0)}" + + networks = scope_context.get("networks", []) + if networks: + task_description += "\nNetworks:" + for net in networks: + task_description += f"\n - {net.get('name')}: {net.get('cidr', 'N/A')} ({net.get('type')})" + + in_scope_domains = scope_context.get("in_scope_domains", []) + if in_scope_domains: + task_description += f"\nIn-scope domains: {', '.join(in_scope_domains)}" + + task_description += "\n" + + # Inject exclusion rules for agents to respect + if exclusion_rules: + task_description += "\n\n" + if exclusion_rules.get("excluded_hosts"): + task_description += f"\nExcluded hosts: {', '.join(exclusion_rules['excluded_hosts'])}" + if exclusion_rules.get("excluded_cidrs"): + task_description += f"\nExcluded CIDRs: {', '.join(exclusion_rules['excluded_cidrs'])}" + if exclusion_rules.get("excluded_urls"): + task_description += f"\nExcluded URLs: {', '.join(exclusion_rules['excluded_urls'])}" + if exclusion_rules.get("excluded_paths"): + task_description += f"\nExcluded paths: {', '.join(exclusion_rules['excluded_paths'])}" + if exclusion_rules.get("excluded_ports"): + task_description += f"\nExcluded ports: {', '.join(map(str, exclusion_rules['excluded_ports']))}" + if exclusion_rules.get("out_of_scope_domains"): + task_description += f"\nOut-of-scope domains: {', '.join(exclusion_rules['out_of_scope_domains'])}" + task_description += "\n" + return await self.agent_loop(task=task_description) diff --git a/strix/agents/state.py b/strix/agents/state.py index 81ac6572e..8f2a1afcb 100644 --- a/strix/agents/state.py +++ b/strix/agents/state.py @@ -12,6 +12,7 @@ def _generate_agent_id() -> str: class AgentState(BaseModel): agent_id: str = Field(default_factory=_generate_agent_id) agent_name: str = "Strix Agent" + agent_role: str | None = None parent_id: str | None = None sandbox_id: str | None = None sandbox_token: str | None = None diff --git a/strix/interface/cli.py b/strix/interface/cli.py index c9bc78ffe..cee873bc5 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -70,6 +70,11 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 "run_name": args.run_name, } + # Add scope context if scope file was loaded + if hasattr(args, "scope_config") and args.scope_config is not None: + scan_config["scope_context"] = args.scope_config.get_agent_context() + scan_config["exclusion_rules"] = args.scope_config.get_exclusion_rules() + llm_config = LLMConfig() agent_config = { "llm_config": llm_config, diff --git a/strix/interface/main.py b/strix/interface/main.py index ef8e6f864..f632eb6b2 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -269,6 +269,11 @@ def parse_arguments() -> argparse.Namespace: # Custom instructions strix --target example.com --instruction "Focus on authentication vulnerabilities" + + # Using scope file + strix --scope scope.yaml + strix --scope scope.yaml --filter "tags:critical" + strix --scope scope.yaml --validate """, ) @@ -276,11 +281,29 @@ def parse_arguments() -> argparse.Namespace: "-t", "--target", type=str, - required=True, + required=False, action="append", help="Target to test (URL, repository, local directory path, domain name, or IP address). " "Can be specified multiple times for multi-target scans.", ) + parser.add_argument( + "--scope", + type=str, + help="Path to scope configuration file (YAML or JSON). " + "Defines targets, networks, exclusions, and settings.", + ) + parser.add_argument( + "--filter", + type=str, + action="append", + help="Filter scope targets by criteria. Format: 'key:value'. " + "Supported keys: tags, network, type. Can be specified multiple times.", + ) + parser.add_argument( + "--validate", + action="store_true", + help="Validate scope file and exit without running scan.", + ) parser.add_argument( "--instruction", type=str, @@ -310,27 +333,163 @@ def parse_arguments() -> argparse.Namespace: args = parser.parse_args() - args.targets_info = [] - for target in args.target: - try: - target_type, target_dict = infer_target_type(target) + # Require either --target or --scope + if not args.target and not args.scope: + parser.error("Either --target or --scope is required") - if target_type == "local_code": - display_target = target_dict.get("target_path", target) + args.targets_info = [] + args.scope_config = None + + # Handle scope file + if args.scope: + args.scope_config = _load_scope_file(args.scope, args.filter, args.validate, parser) + if args.validate: + return args # Early return for validation-only mode + + # Get targets from scope + scope_targets = args.scope_config.to_targets_info() + args.targets_info.extend(scope_targets) + + # Add scope-derived instruction context + scope_instruction = args.scope_config.get_instruction_context() + if scope_instruction: + if args.instruction: + args.instruction = f"{scope_instruction}. {args.instruction}" else: - display_target = target + args.instruction = scope_instruction - args.targets_info.append( - {"type": target_type, "details": target_dict, "original": display_target} - ) - except ValueError: - parser.error(f"Invalid target '{target}'") + # Handle individual --target arguments + if args.target: + for target in args.target: + try: + target_type, target_dict = infer_target_type(target) + + if target_type == "local_code": + display_target = target_dict.get("target_path", target) + else: + display_target = target + + args.targets_info.append( + {"type": target_type, "details": target_dict, "original": display_target} + ) + except ValueError: + parser.error(f"Invalid target '{target}'") + + if not args.targets_info: + parser.error("No valid targets found") assign_workspace_subdirs(args.targets_info) return args +def _load_scope_file( + scope_path: str, + filters: list[str] | None, + validate_only: bool, + parser: argparse.ArgumentParser, +) -> Any: + """Load and validate scope file.""" + from strix.scope import ScopeConfig, validate_scope + + console = Console() + + try: + scope_config = ScopeConfig.from_file(Path(scope_path)) + except FileNotFoundError: + parser.error(f"Scope file not found: {scope_path}") + except ValueError as e: + parser.error(f"Invalid scope file: {e}") + except Exception as e: + parser.error(f"Error loading scope file: {e}") + + # Validate scope + result = validate_scope(scope_config.model) + + if validate_only: + # Print validation results and exit + if result.valid: + console.print(f"\n[bold green]Scope file is valid:[/] {scope_path}\n") + else: + console.print(f"\n[bold red]Scope file has errors:[/] {scope_path}\n") + + if result.errors: + console.print("[bold red]Errors:[/]") + for error in result.errors: + console.print(f" [red]- {error}[/]") + + if result.warnings: + console.print("\n[bold yellow]Warnings:[/]") + for warning in result.warnings: + console.print(f" [yellow]- {warning}[/]") + + # Print summary + console.print(f"\n[bold cyan]Scope Summary:[/]") + console.print(f" Engagement: {scope_config.metadata.engagement_name}") + console.print(f" Type: {scope_config.metadata.engagement_type}") + console.print(f" Mode: {scope_config.settings.operational_mode}") + console.print(f" Networks: {len(scope_config.networks)}") + console.print(f" Targets: {len(scope_config.targets)}") + console.print() + + if not result.valid: + sys.exit(1) + sys.exit(0) + + # Show warnings but continue + if result.warnings: + for warning in result.warnings: + console.print(f"[yellow]Warning:[/] {warning}") + + if not result.valid: + console.print("\n[bold red]Scope file has errors:[/]") + for error in result.errors: + console.print(f" [red]- {error}[/]") + parser.error("Invalid scope file") + + # Apply filters if specified + if filters: + _apply_scope_filters(scope_config, filters, parser) + + return scope_config + + +def _apply_scope_filters( + scope_config: Any, filters: list[str], parser: argparse.ArgumentParser +) -> None: + """Apply filters to scope targets.""" + tags_filter: list[str] = [] + network_filter: str | None = None + type_filter: str | None = None + + for f in filters: + if ":" not in f: + parser.error(f"Invalid filter format: {f}. Use 'key:value'") + + key, value = f.split(":", 1) + key = key.lower().strip() + value = value.strip() + + if key == "tags": + tags_filter.extend(v.strip() for v in value.split(",")) + elif key == "network": + network_filter = value + elif key == "type": + type_filter = value + else: + parser.error(f"Unknown filter key: {key}. Supported: tags, network, type") + + # Filter targets in place + filtered = scope_config.filter_targets( + tags=tags_filter if tags_filter else None, + network=network_filter, + target_type=type_filter, + ) + + # Replace targets with filtered list + scope_config.model.targets = filtered + + def display_completion_message(args: argparse.Namespace, results_path: Path) -> None: console = Console() tracer = get_global_tracer() @@ -452,6 +611,11 @@ def main() -> None: args = parse_arguments() + # Validation-only mode exits in parse_arguments via _load_scope_file + # This is a safety check in case the flow changes + if args.validate: + return + check_docker_installed() pull_docker_image() diff --git a/strix/interface/tui.py b/strix/interface/tui.py index ff0a255ca..5f67e10cb 100644 --- a/strix/interface/tui.py +++ b/strix/interface/tui.py @@ -310,13 +310,20 @@ def __init__(self, args: argparse.Namespace): self._setup_cleanup_handlers() def _build_scan_config(self, args: argparse.Namespace) -> dict[str, Any]: - return { + config = { "scan_id": args.run_name, "targets": args.targets_info, "user_instructions": args.instruction or "", "run_name": args.run_name, } + # Add scope context if scope file was loaded + if hasattr(args, "scope_config") and args.scope_config is not None: + config["scope_context"] = args.scope_config.get_agent_context() + config["exclusion_rules"] = args.scope_config.get_exclusion_rules() + + return config + def _build_agent_config(self, args: argparse.Namespace) -> dict[str, Any]: llm_config = LLMConfig() diff --git a/strix/llm/config.py b/strix/llm/config.py index ea7163a1d..33e98898f 100644 --- a/strix/llm/config.py +++ b/strix/llm/config.py @@ -8,6 +8,7 @@ def __init__( enable_prompt_caching: bool = True, prompt_modules: list[str] | None = None, timeout: int | None = None, + agent_role: str | None = None, ): self.model_name = model_name or os.getenv("STRIX_LLM", "openai/gpt-5") @@ -16,5 +17,6 @@ def __init__( self.enable_prompt_caching = enable_prompt_caching self.prompt_modules = prompt_modules or [] + self.agent_role = agent_role self.timeout = timeout or int(os.getenv("LLM_TIMEOUT", "600")) diff --git a/strix/llm/llm.py b/strix/llm/llm.py index 99a566a49..351575476 100644 --- a/strix/llm/llm.py +++ b/strix/llm/llm.py @@ -169,8 +169,13 @@ def get_module(name: str) -> str: self.jinja_env.globals["get_module"] = get_module + agent_role = self.config.agent_role + + def get_tools_prompt_with_role() -> str: + return get_tools_prompt(role=agent_role) + self.system_prompt = self.jinja_env.get_template("system_prompt.jinja").render( - get_tools_prompt=get_tools_prompt, + get_tools_prompt=get_tools_prompt_with_role, loaded_module_names=list(prompt_module_content.keys()), **prompt_module_content, ) diff --git a/strix/prompts/coordination/root_agent.jinja b/strix/prompts/coordination/root_agent.jinja index e877bb913..a9cb473cf 100644 --- a/strix/prompts/coordination/root_agent.jinja +++ b/strix/prompts/coordination/root_agent.jinja @@ -2,40 +2,378 @@ 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 +1. Parse user instructions and determine operational mode +2. Create specialized agents for specific security tasks +3. Monitor agent progress and coordinate between them +4. Enforce safety controls and confirm destructive actions with user +5. Compile final scan reports from agent findings (JSON + Markdown) +6. Manage agent communication and dependencies CRITICAL RESTRICTIONS: -- NEVER perform vulnerability testing or security assessments +- NEVER perform vulnerability testing or security assessments directly - 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! +- ONLY use agent_graph, finish, and think tools for coordination +- ALWAYS confirm with user before any potentially destructive action +- You can create agents throughout the scan process based on findings, not just at the beginning - -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 + +Parse user instructions to determine the operational mode. Support both keywords and natural language: -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 +RECON-ONLY MODE (keywords: "recon only", "reconnaissance only", "no exploitation", "passive"): +- Create only reconnaissance and discovery agents +- Generate PoC scripts but DO NOT execute exploits +- Focus on: enumeration, version detection, CVE correlation, configuration analysis +- Validation agents MAY run PoCs in sandbox to confirm findings -COORDINATION GUIDELINES: -- Ensure clear task boundaries and success criteria -- Terminate redundant agents when objectives overlap -- Use message passing only when essential (requests/answers or critical handoffs); avoid routine status messages and prefer batched updates - +FULL PENTEST MODE (keywords: "full pentest", "full test", "exploitation allowed"): +- Create full agent tree: recon -> vuln testing -> validation -> reporting -> fixing +- Active exploitation permitted within scope +- Post-exploitation and pivoting allowed + +POC-ONLY MODE (keywords: "poc only", "proof of concept only"): +- Discovery and validation only +- Generate and validate PoCs in sandbox +- No active exploitation beyond PoC confirmation +- No fixing agents + +DEFAULT MODE (no specific keywords): +- Interpret user intent from natural language +- When uncertain, ASK the user for clarification +- Default to POC-ONLY if intent is unclear + +NOTE: If a block is provided, the mode may be overridden by scope settings. + + + +When and blocks are provided in the task, you MUST: + +RESPECT SCOPE BOUNDARIES: +1. Only test targets explicitly listed in scope or within defined network CIDRs +2. NEVER access, scan, or test hosts/URLs in exclusion lists +3. Respect out-of-scope domains - do not follow redirects or links to them +4. Skip excluded ports even if discovered on in-scope targets +5. Avoid testing URL paths that match exclusion patterns + +OPERATIONAL MODE FROM SCOPE: +- If scope_context specifies mode "recon-only": Follow RECON-ONLY restrictions +- If scope_context specifies mode "poc-only": Follow POC-ONLY restrictions +- If scope_context specifies mode "full-pentest": Allow full testing within scope + +NETWORK AWARENESS: +- Use network definitions to understand target relationships +- Internal networks may have different testing approaches than external +- VLAN information helps identify network segmentation boundaries +- Gateway IPs are typically infrastructure, test with extra caution + +SCOPE VALIDATION: +Before any child agent performs actions on a target: +1. Verify target is in scope (matches defined targets, CIDRs, or domain patterns) +2. Verify target is NOT in exclusion lists +3. Verify action respects operational mode constraints +4. If uncertain, escalate to coordinator (you) for decision + +WHEN SCOPE IS NOT PROVIDED: +- Fall back to targets specified in command line +- Apply default safety controls +- When discovering new targets, ASK user before expanding scope + + + +DESTRUCTIVE ACTION CATEGORIES (ALWAYS ASK USER FIRST): +- Any action that modifies target system state +- Exploitation that could cause service disruption +- Actions that write/delete/modify files on target +- Network attacks (DoS, flooding, etc.) +- Credential stuffing or brute force attacks +- Any action flagged by child agents as potentially destructive + +SAFE ACTIONS (proceed without asking): +- Passive reconnaissance (port scanning, version detection) +- Reading publicly accessible information +- Generating PoC scripts (not executing) +- Static code analysis +- Configuration review +- CVE correlation and research + +CONFIRMATION FORMAT: +When destructive action is proposed, message user with: +- Action description +- Potential impact +- Reversibility assessment +- Request explicit approval before proceeding + + + +STRUCTURE: Hybrid phase-based with target-specialized children +MAX DEPTH: 2 levels + fixing sub-agent when needed + +PHASE 1 - RECONNAISSANCE AGENTS (Top Level): +Create based on target type detected: + +For WEB TARGETS (URLs, domains, APIs): +- Web Recon Agent: subdomain enum, endpoint discovery, tech fingerprinting + prompt_modules: reconnaissance (when available) + +For INFRASTRUCTURE TARGETS (IPs, servers, hypervisors): +- Infrastructure Recon Agent: port scanning, service detection, version enumeration + prompt_modules: proxmox_ve, or relevant infra module + +For CLOUD TARGETS (AWS, Azure, GCP, K8s): +- Cloud Recon Agent: resource enumeration, IAM analysis, misconfiguration detection + prompt_modules: relevant cloud module + +For CODE TARGETS (repositories, local code): +- Code Analysis Agent: static analysis, dependency review, secret detection + prompt_modules: relevant framework module + +PHASE 2 - SPECIALIZED TESTING AGENTS (Children of Phase 1): +Created by Phase 1 agents based on discoveries: + +Web Vulnerability Agents: +- SQLi Agent (prompt_modules: sql_injection) +- XSS Agent (prompt_modules: xss) +- Auth Agent (prompt_modules: authentication_jwt) +- SSRF/XXE Agent (prompt_modules: ssrf, xxe) +- IDOR Agent (prompt_modules: idor) +- Business Logic Agent (prompt_modules: business_logic) + +Infrastructure Vulnerability Agents: +- API Security Agent (prompt_modules: relevant tech module) +- Container Escape Agent (prompt_modules: proxmox_ve or relevant) +- Network Segmentation Agent +- Privilege Escalation Agent + +Cloud Vulnerability Agents: +- IAM Misconfiguration Agent +- Storage Exposure Agent +- Network Security Agent + +VALIDATION AGENTS (Mandatory for ALL findings): +- Spawned by testing agents when potential vulnerability found +- Must validate with PoC execution in sandbox +- Reports back: confirmed/false-positive/needs-manual-review + +FIXING AGENTS (Only in white-box mode with code access): +- Spawned after validation confirms vulnerability +- Implements fix in code +- Creates verification test + + + +Automatically detect target type from input: + +WEB APPLICATION: +- Starts with http:// or https:// +- Contains domain name patterns +- Specified as URL + +INFRASTRUCTURE: +- IP address (IPv4 or IPv6) +- Hostname resolving to infrastructure service +- Known infrastructure ports (8006=Proxmox, 443=general, 22=SSH, etc.) +- User mentions: Proxmox, VMware, ESXi, hypervisor, server, firewall, router + +CLOUD: +- AWS ARNs, S3 buckets, EC2 references +- Azure resource IDs, blob storage +- GCP project references +- Kubernetes contexts, namespaces +- User mentions: AWS, Azure, GCP, cloud, k8s, kubernetes + +CODE/REPOSITORY: +- GitHub/GitLab/Bitbucket URLs +- Local path starting with / or ./ +- User mentions: repository, code, source, repo + +MOBILE: +- APK/IPA file references +- Mobile app store links +- User mentions: android, iOS, mobile app + + + +STRUCTURED MESSAGE TYPES: + +FINDING MESSAGE (from testing agent to coordinator): +{ + "type": "finding", + "severity": "critical|high|medium|low|info", + "category": "vulnerability|misconfiguration|exposure|weakness", + "title": "brief description", + "target": "affected component", + "confidence": "confirmed|likely|possible", + "needs_validation": true|false, + "poc_available": true|false +} + +REQUEST MESSAGE (agent requesting action/info): +{ + "type": "request", + "action": "create_agent|user_confirmation|info_needed", + "details": "what is needed", + "priority": "high|normal|low" +} + +STATUS MESSAGE (progress update - use sparingly): +{ + "type": "status", + "phase": "current phase", + "progress": "percentage or milestone", + "blockers": "any issues" | null +} + +GUIDELINES: +- Batch non-urgent messages +- Prefer completion reports over status updates +- Only use status messages for significant milestones or blockers + + + +PROMPT MODULE ASSIGNMENT: + +Infrastructure Targets: +- Proxmox VE: proxmox_ve +- General infrastructure: (create recon module if needed) +- Kubernetes/containers: (k8s module when available) + +Web Targets: +- Firebase apps: firebase_firestore +- Supabase apps: supabase +- FastAPI: fastapi +- Next.js: nextjs +- General web: relevant vulnerability modules + +Vulnerability-Specific: +- SQL Injection: sql_injection +- XSS: xss +- SSRF: ssrf +- XXE: xxe +- RCE: rce +- CSRF: csrf +- IDOR: idor +- Auth/JWT: authentication_jwt +- Business Logic: business_logic +- File Upload: insecure_file_uploads +- Path Traversal: path_traversal_lfi_rfi +- Race Conditions: race_conditions +- Mass Assignment: mass_assignment +- GraphQL: graphql + +RULES: +- Each agent: 1-3 modules preferred, max 5 +- Related vulns can share agent (SSRF+XXE, Auth+Business Logic) +- Infrastructure agents get tech-specific module + relevant vuln modules + + + +FINAL REPORT FORMAT: JSON + Markdown + +JSON STRUCTURE: +{ + "scan_metadata": { + "scan_id": "uuid", + "timestamp": "ISO8601", + "targets": ["list of targets"], + "mode": "recon-only|poc-only|full-pentest", + "duration_seconds": number + }, + "summary": { + "total_findings": number, + "by_severity": {"critical": n, "high": n, "medium": n, "low": n, "info": n}, + "validated": number, + "false_positives": number + }, + "findings": [ + { + "id": "finding-uuid", + "title": "string", + "severity": "critical|high|medium|low|info", + "category": "string", + "target": "affected component", + "description": "detailed description", + "poc": "proof of concept code/steps", + "remediation": "fix recommendation", + "references": ["CVE-xxx", "URL"], + "validated": true|false, + "validated_by": "agent-id" + } + ], + "agents": { + "total_created": number, + "by_type": {"recon": n, "testing": n, "validation": n, "reporting": n, "fixing": n} + } +} + +MARKDOWN REPORT: +Generated alongside JSON with human-readable format including: +- Executive summary +- Findings table sorted by severity +- Detailed finding descriptions +- PoC code blocks +- Remediation recommendations + + + +For critical decisions, complex targets, or high-confidence requirements, use MULTI-MODEL CONSENSUS: + +WHEN TO USE CONSENSUS: +- Severity classification for critical/high findings +- Complex vulnerability chains requiring correlation +- Uncertain findings needing multiple perspectives +- Final report generation for important assessments +- When user explicitly requests high-confidence validation + +HOW TO IMPLEMENT (using existing tools): + +1. CREATE ADVISOR AGENTS with different models: + + Analyze these findings and provide independent assessment: [findings] + Advisor-A + anthropic/claude-sonnet-4-20250514 + validation + + + + Analyze these findings and provide independent assessment: [findings] + Advisor-B + openai/gpt-4o + validation + + +2. WAIT for advisors to complete and report back + +3. SYNTHESIZE recommendations: + - Agreement: High confidence in shared conclusion + - Disagreement: Flag for manual review or create tie-breaker agent + - Partial agreement: Note confidence levels in report + +MODEL SELECTION GUIDELINES: +- Fast/cheap models (gemini-flash, gpt-4o-mini): Recon, simple validation +- Capable models (gpt-4o, claude-sonnet): Testing, complex analysis +- Reasoning models (o3, claude-sonnet-4-5): Critical decisions, consensus +- Model parameter cascades: explicit -> parent -> env default (STRIX_LLM) + +FALLBACK BEHAVIOR: +- If specified model fails, agent uses parent's model +- If parent has no model, falls back to STRIX_LLM environment variable +- Always maintain operation continuity over model preference + 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 +1. Collect findings from all agents via structured messages +2. Deduplicate and correlate related findings +3. For critical findings, consider consensus validation workflow +4. Generate JSON report structure +5. Generate Markdown report +6. Present summary to user +7. Use finish_scan tool to complete the assessment + +EARLY TERMINATION: +- If user requests stop, gracefully terminate all agents +- Compile partial report with findings collected so far +- Mark report as "incomplete" with reason -Your value is in orchestration, not execution. +Your value is in orchestration, safety enforcement, and clear reporting - not execution. diff --git a/strix/prompts/technologies/proxmox_ve.jinja b/strix/prompts/technologies/proxmox_ve.jinja new file mode 100644 index 000000000..dd7820f41 --- /dev/null +++ b/strix/prompts/technologies/proxmox_ve.jinja @@ -0,0 +1,274 @@ + +PROXMOX VE β€” ADVERSARIAL TESTING AND EXPLOITATION + +Proxmox VE exposes a REST API, web UI, SPICE/VNC consoles, SSH, and cluster services. Most impactful findings arise from API authentication bypasses, privilege escalation via misconfigured roles/permissions, exposed management interfaces, container/VM escape vectors, insecure cluster communication, and backup/storage misconfigurations. The API uses ticket-based authentication with CSRF tokens; always verify both are enforced across all endpoints. + + +- Web UI and API (port 8006, /api2/json/) +- Authentication: PAM, PVE, LDAP/AD, OpenID Connect +- Authorization: Users, Groups, Roles, ACLs, Pools, Realms +- VM/Container management: QEMU/KVM, LXC +- Storage backends: local, NFS, CIFS, Ceph, ZFS, iSCSI, PBS +- Networking: bridges, VLANs, SDN, firewall +- Cluster: Corosync (5405/udp), pmxcfs, HA +- Backup: vzdump, Proxmox Backup Server integration +- Console access: SPICE, noVNC, xterm.js + + + +1. Enumerate version and patch level via API or web UI; map to known CVEs. +2. Identify authentication realms and obtain credentials for multiple privilege levels: no auth, unprivileged user, VM admin, full admin. +3. Build Resource x Action x Principal matrix across API endpoints. Test CRUD operations on VMs, containers, storage, users, and cluster config. +4. Probe network segmentation: can VMs/containers reach management interfaces? Are cluster ports exposed externally? +5. Test console access controls, backup permissions, and storage ACLs for privilege boundaries. + + + +- API base: https://:8006/api2/json/ +- Authentication flow: POST /api2/json/access/ticket with username@realm + password; returns ticket (cookie) + CSRFPreventionToken (header). +- All state-changing requests require: Cookie: PVEAuthCookie= AND CSRFPreventionToken: +- Realms: pam (system users), pve (Proxmox native), ldap, ad, openid +- Permissions: Role-based ACLs on paths like /vms/, /storage/, /pool/, /access/groups/ +- Cluster: corosync for membership/quorum, pmxcfs for distributed config, pve-cluster service + + + +- Ticket format: PVE:@::: +- Tickets expire (default 2 hours); refresh via /api2/json/access/ticket with existing cookie. +- CSRFPreventionToken must match ticket; both required for POST/PUT/DELETE. +- API tokens: @!=; can have separate permissions from user. +- Pitfalls: + - CSRF token not validated on some endpoints or HTTP methods. + - Ticket replay across nodes in cluster if clocks are skewed. + - API tokens with overly broad permissions stored in automation scripts. + - TFA (TOTP/U2F/Recovery) bypass via API token authentication. +- Tests: + - Attempt state-changing requests without CSRFPreventionToken; verify rejection. + - Replay tickets across cluster nodes; check timestamp/signature validation. + - Enumerate API tokens via /api2/json/access/users//token; test token permissions. + - Bypass TFA: authenticate with API token instead of password+TFA. + + + +- Critical endpoints requiring elevated privileges: + - /nodes//qemu//agent/* - Guest agent commands (file read/write, exec) + - /nodes//execute - Arbitrary command execution (root only) + - /nodes//ceph/* - Ceph cluster management + - /cluster/config - Cluster-wide configuration + - /access/users, /access/groups, /access/roles, /access/acl - IAM management + - /nodes//storage//upload - File upload to storage + - /nodes//vzdump - Backup operations +- Common gaps: + - Unprivileged users with VM.Console can access guest agent if enabled. + - Pool administrators can escalate via storage access. + - Backup permissions allow reading VM disk contents. +- Tests: + - As low-priv user, enumerate all accessible /nodes//qemu and /nodes//lxc. + - Try guest agent endpoints: /agent/file-read, /agent/exec; verify ACL enforcement. + - Attempt /nodes//execute as non-root; should fail. + - Upload malicious content to storage; check for path traversal in filename. + + + +- Role hierarchy: PVEAdmin > PVEVMAdmin > PVEVMUser > PVEAuditor +- ACL inheritance: permissions cascade from / down through /vms, /storage, /pool paths. +- Escalation vectors: + - VM.Config.Options allows changing boot order, adding USB/PCI passthrough. + - VM.Config.Disk + storage access = mount host paths into VM. + - Container with nesting=1 or privileged=1 enables escape techniques. + - Datastore.AllocateTemplate on shared storage affects all VMs using it. + - Pool.Allocate allows adding VMs to pools, inheriting pool permissions. +- Tests: + - Enumerate effective permissions: GET /api2/json/access/permissions + - As VMAdmin, try to modify ACLs: PUT /api2/json/access/acl + - Check if unprivileged containers can enable nesting/features post-creation. + - Verify storage isolation: can user A's VM access user B's storage? + + + +- QEMU/KVM: + - Guest agent (qemu-ga): if enabled, host can execute commands in guest; verify ACLs. + - VNC/SPICE: websocket proxy at /api2/json/nodes//qemu//vncproxy. + - Passthrough (PCI/USB): can expose host devices; verify isolation. + - Live migration: transfers memory/disk; check for interception on network. +- LXC Containers: + - Privileged containers run as root with host capabilities; escape is trivial. + - Unprivileged containers use user namespaces; verify uid/gid mapping. + - Features: nesting, fuse, mknod, keyctl - each weakens isolation. + - Bind mounts: check for host path exposure via mp0, mp1, etc. + - AppArmor/seccomp profiles: verify enforcement, check for bypass. +- Tests: + - Inside container: cat /proc/1/cgroup, ls /dev, capsh --print to assess isolation. + - Check for CVE-2019-5736 (runc), CVE-2022-0185 (file_caps), container escapes. + - Verify AppArmor: cat /proc/self/attr/current should show profile. + - Attempt to access /dev/sda or mount host filesystems from within container. + + + +- VNC proxy: /api2/json/nodes//qemu//vncproxy returns ticket + port. +- SPICE proxy: /api2/json/nodes//qemu//spiceproxy returns connection file. +- xterm.js: /api2/json/nodes//lxc//termproxy for container shell. +- Websocket upgrade at wss://:8006/api2/json/nodes//qemu//vncwebsocket +- Risks: + - Console tickets may have longer validity than auth tickets. + - Websocket connections may not re-validate permissions after initial auth. + - VNC password (if set) often weak or shared across VMs. +- Tests: + - Obtain VNC ticket, revoke user access, verify ticket still works (ticket lifetime issue). + - Connect to websocket with expired/revoked auth cookie; check enforcement. + - Enumerate VMs with VNC vs SPICE vs serial console; check password requirements. + + + +- Corosync: UDP 5405 for cluster communication; uses pre-shared key. +- pmxcfs: distributed filesystem at /etc/pve; stores cluster config, VM configs, user database. +- Join token: pvecm expected -y creates token valid for limited time. +- Risks: + - Corosync key (/etc/corosync/authkey) disclosure allows cluster join. + - Exposed UDP 5405 allows cluster membership enumeration. + - pvecm add with stolen key adds malicious node. + - Split-brain scenarios during network partition. +- Tests: + - Scan for UDP 5405 from external/VM networks; should be firewalled. + - If key obtained, attempt cluster join simulation. + - Check /etc/pve permissions; should not be world-readable. + - Verify HA fencing configuration to prevent resource contention. + + + +- Storage types: local (dir/lvm/zfs), shared (NFS/CIFS/Ceph/iSCSI/GlusterFS), PBS +- Content types: images, rootdir, vztmpl, backup, iso, snippets +- Risks: + - NFS exports with no_root_squash allow host root escalation from VM. + - CIFS credentials stored in plaintext in /etc/pve/storage.cfg. + - Backup files (.vma, .tar) contain full disk images; sensitive data exposure. + - ISO storage allows uploading bootable images; malicious ISOs. + - Snippets (hookscripts, cloud-init) execute with elevated privileges. +- Tests: + - Enumerate storage: GET /api2/json/storage; check each for permissions. + - Read backup files if Datastore.Audit granted; extract sensitive data. + - Upload hookscript to snippets storage; verify execution context. + - Check NFS mount options; verify root_squash enforcement. + - Attempt path traversal in storage upload endpoints. + + + +- vzdump: creates backups to storage; options include compression, encryption, mode (snapshot/suspend/stop). +- PBS (Proxmox Backup Server): separate service for incremental, encrypted backups. +- Risks: + - Backups contain encryption keys, passwords, sensitive data. + - Backup permissions separate from VM permissions; user may backup but not access VM. + - Unencrypted backups on shared storage readable by storage admins. + - Restore to different VM bypasses original VM's access controls. +- Tests: + - With Datastore.AllocateSpace, attempt backup of VMs user doesn't own. + - Download backup file; extract and analyze contents. + - Restore backup to new VM under attacker's control. + - Check PBS encryption: are keys properly managed? Stored separately? + + + +- Bridges: vmbr0 typically bridges physical NIC; VMs share L2 domain. +- VLANs: 802.1q tagging; verify VLAN isolation between tenants. +- SDN: software-defined networking zones, vnets; check isolation. +- Firewall: datacenter, host, and VM/CT levels; nftables backend. +- Risks: + - VMs on same bridge can ARP spoof, sniff traffic. + - Management interface (8006) reachable from VM network. + - Firewall disabled by default; even when enabled, rules may be lax. + - VLAN hopping if trunking misconfigured. +- Tests: + - From VM, scan host IP on port 8006, 22, 3128 (spice), 5900-5999 (VNC). + - ARP scan from VM; identify other VMs and host on same bridge. + - Verify firewall rules: GET /api2/json/nodes//firewall/rules. + - Test VLAN isolation: can VM on VLAN 10 reach VLAN 20? + + + +- CVE-2022-35508: API authentication bypass via crafted ticket. +- CVE-2023-43320: XSS in web UI task viewer. +- CVE-2024-21545: Privilege escalation via vzdump. +- CVE-2024-21546: Information disclosure in cluster join. +- Always check: https://pve.proxmox.com/wiki/Roadmap#Security_Advisories +- Tests: + - Identify exact version: GET /api2/json/version + - Check pveversion -v output for package versions. + - Correlate with CVE databases; test applicable exploits in PoC mode. + + + +- Content-type switching: API accepts JSON; try form-encoded for parser differences. +- HTTP method override: X-HTTP-Method-Override header may bypass method restrictions. +- Path normalization: /api2/json/nodes/node1/qemu/100 vs /api2/json/nodes/node1/qemu/100/ trailing slash. +- Unicode/encoding: URL-encoded parameters may bypass input validation. +- Header injection: Host header manipulation for SSRF or cache poisoning. +- Race conditions: parallel requests during VM state changes (start/stop/migrate). + + + +- Task polling: /api2/json/nodes//tasks//status reveals operation success/failure. +- Error messages: detailed errors disclose internal paths, versions, configs. +- Timing: authentication failures vs. invalid user vs. wrong password timing differences. +- Cluster status: /api2/json/cluster/status reveals node names, IPs, quorum state. + + + +- API exploration: curl with cookie/CSRF headers; pve-api-viewer for endpoint discovery. +- Scanning: nmap for port enumeration; nuclei templates for Proxmox-specific checks. +- Authentication: script ticket acquisition and renewal for long-running tests. +- Container escape: use container escape toolkits (CDK, deepce) inside LXC. +- Traffic analysis: tcpdump on bridges; Wireshark for SPICE/VNC protocols. +- Backup analysis: qemu-img for .vma conversion; tar extraction for container backups. + + + +- Is the web UI (8006) restricted to management network only? +- Are all API state-changing requests validating CSRFPreventionToken? +- Do ACLs follow least-privilege; are default roles appropriately scoped? +- Are TFA requirements enforced consistently (API tokens bypass TFA by design)? +- Are containers unprivileged by default with appropriate AppArmor profiles? +- Is cluster communication (Corosync) firewalled from untrusted networks? +- Are backups encrypted and access-controlled separately from VM access? +- Is network segmentation enforced between management, VM, and storage networks? +- Are storage credentials (NFS/CIFS/iSCSI) properly secured? +- Is guest agent access restricted to appropriate administrators? + + + +1. Demonstrate authentication bypass or CSRF token bypass on state-changing endpoint. +2. Show privilege escalation from low-privilege user to VM access or admin functions. +3. Prove container escape or host access from within VM/container. +4. Document network segmentation failure (VM accessing management interface). +5. Provide minimal PoC scripts with exact API calls, headers, and expected vs. actual responses. + + + +- Management interface intentionally exposed (documented homelab setup). +- Privileged containers used for specific workloads with accepted risk. +- Guest agent enabled for authorized monitoring/management. +- Backup access granted to backup administrators by design. +- Cluster ports accessible within trusted management VLAN only. + + + +- Full infrastructure compromise via API authentication bypass. +- Cross-tenant VM access and data exfiltration. +- Container/VM escape leading to host root access. +- Cluster takeover via Corosync key disclosure. +- Sensitive data exposure via backup access or storage misconfiguration. +- Denial of service via resource exhaustion or cluster disruption. + + + +1. Start with version enumeration; older Proxmox installations often lag on patches. +2. Map the permission model completely before testing; ACLs are granular but complex. +3. Test both web UI and direct API; UI may have additional client-side restrictions. +4. Container escape is easier than VM escape; focus on LXC if present. +5. Backup access is often overlooked; it's equivalent to full disk read access. +6. Cluster join tokens and Corosync keys are crown jewels; trace their exposure. +7. Guest agent is powerful; VM.Console permission may grant more than expected. +8. Always verify network segmentation from inside VMs, not just from external scans. + + +Proxmox VE is a hypervisorβ€”compromise means full infrastructure access. Every permission boundary (API auth, ACLs, container isolation, network segmentation, cluster membership) must be validated independently. Focus on the gaps between intended isolation and actual enforcement. + diff --git a/strix/scope/__init__.py b/strix/scope/__init__.py new file mode 100644 index 000000000..3cb1af9d6 --- /dev/null +++ b/strix/scope/__init__.py @@ -0,0 +1,34 @@ +"""Scope configuration module for Strix.""" + +from .models import ( + CredentialDefinition, + DomainScope, + Exclusions, + NetworkDefinition, + ScopeConfigModel, + ScopeMetadata, + ScopeSettings, + ServiceDefinition, + TargetDefinition, +) +from .parser import ScopeConfig +from .validator import ScopeValidator, ValidationResult, validate_scope + +__all__ = [ + # Main class + "ScopeConfig", + # Models + "ScopeConfigModel", + "ScopeMetadata", + "ScopeSettings", + "NetworkDefinition", + "TargetDefinition", + "ServiceDefinition", + "CredentialDefinition", + "Exclusions", + "DomainScope", + # Validation + "ScopeValidator", + "ValidationResult", + "validate_scope", +] diff --git a/strix/scope/models.py b/strix/scope/models.py new file mode 100644 index 000000000..63143846d --- /dev/null +++ b/strix/scope/models.py @@ -0,0 +1,167 @@ +"""Pydantic models for scope configuration.""" + +from typing import Any + +from pydantic import BaseModel, Field, field_validator + + +class ScopeMetadata(BaseModel): + """Engagement metadata.""" + + engagement_name: str = "Unnamed Engagement" + engagement_type: str = "internal" # internal | external | hybrid + start_date: str | None = None + end_date: str | None = None + tester: str | None = None + notes: str | None = None + + @field_validator("engagement_type") + @classmethod + def validate_engagement_type(cls, v: str) -> str: + valid_types = {"internal", "external", "hybrid"} + if v.lower() not in valid_types: + raise ValueError(f"engagement_type must be one of: {valid_types}") + return v.lower() + + +class ScopeSettings(BaseModel): + """Global scan settings.""" + + operational_mode: str = "poc-only" # recon-only | poc-only | full-pentest + max_agents: int = 20 + require_validation: bool = True + generate_fixes: bool = False + + @field_validator("operational_mode") + @classmethod + def validate_operational_mode(cls, v: str) -> str: + valid_modes = {"recon-only", "poc-only", "full-pentest"} + if v.lower() not in valid_modes: + raise ValueError(f"operational_mode must be one of: {valid_modes}") + return v.lower() + + +class ServiceDefinition(BaseModel): + """Service running on a port.""" + + port: int + service: str + version: str | None = None + + +class CredentialDefinition(BaseModel): + """Credential for a target (password stored via env var reference).""" + + username: str + password_env: str | None = None # Environment variable name + token_env: str | None = None # For bearer tokens + access_level: str = "user" # user | admin | readonly + + +class NetworkDefinition(BaseModel): + """Network/VLAN definition.""" + + name: str + type: str = "internal" # internal | external + vlan: int | None = None + cidr: str | None = None + gateway: str | None = None + description: str | None = None + + @field_validator("type") + @classmethod + def validate_network_type(cls, v: str) -> str: + valid_types = {"internal", "external"} + if v.lower() not in valid_types: + raise ValueError(f"network type must be one of: {valid_types}") + return v.lower() + + +class TargetDefinition(BaseModel): + """Individual target definition.""" + + # One of these must be set + host: str | None = None # IP address + url: str | None = None # Web URL + repo: str | None = None # Git repository + path: str | None = None # Local path + + # Metadata + name: str | None = None + type: str | None = None # infrastructure | web_application | api | repository | local_code + network: str | None = None # Reference to NetworkDefinition.name + + # Details + ports: list[int] = Field(default_factory=list) + services: list[ServiceDefinition] = Field(default_factory=list) + technologies: list[str] = Field(default_factory=list) + credentials: list[CredentialDefinition] = Field(default_factory=list) + + # Testing configuration + focus_areas: list[str] = Field(default_factory=list) + modules: list[str] = Field(default_factory=list) + tags: list[str] = Field(default_factory=list) + + # API-specific + auth_type: str | None = None # bearer | basic | api_key + token_env: str | None = None + openapi_spec: str | None = None + + # Repository-specific + branch: str | None = None + + # Additional metadata + metadata: dict[str, Any] = Field(default_factory=dict) + + def get_target_value(self) -> str: + """Get the primary target identifier.""" + return self.host or self.url or self.repo or self.path or "" + + def get_target_type(self) -> str: + """Infer target type if not explicitly set.""" + if self.type: + return self.type + if self.host: + return "infrastructure" + if self.url: + return "web_application" + if self.repo: + return "repository" + if self.path: + return "local_code" + return "unknown" + + +class Exclusions(BaseModel): + """Targets and patterns to exclude from testing.""" + + hosts: list[str] = Field(default_factory=list) + cidrs: list[str] = Field(default_factory=list) + urls: list[str] = Field(default_factory=list) + paths: list[str] = Field(default_factory=list) + ports: list[int] = Field(default_factory=list) + + +class DomainScope(BaseModel): + """Domain boundaries for web testing.""" + + in_scope: list[str] = Field(default_factory=list) # Supports wildcards: *.example.com + out_of_scope: list[str] = Field(default_factory=list) + + +class ScopeConfigModel(BaseModel): + """Complete scope configuration model.""" + + metadata: ScopeMetadata = Field(default_factory=ScopeMetadata) + settings: ScopeSettings = Field(default_factory=ScopeSettings) + networks: list[NetworkDefinition] = Field(default_factory=list) + targets: list[TargetDefinition] = Field(default_factory=list) + exclusions: Exclusions = Field(default_factory=Exclusions) + domains: DomainScope = Field(default_factory=DomainScope) + + def get_network_by_name(self, name: str) -> NetworkDefinition | None: + """Look up network by name.""" + for network in self.networks: + if network.name == name: + return network + return None diff --git a/strix/scope/parser.py b/strix/scope/parser.py new file mode 100644 index 000000000..294b5c8ae --- /dev/null +++ b/strix/scope/parser.py @@ -0,0 +1,457 @@ +"""Scope configuration parser.""" + +import ipaddress +import json +import os +from fnmatch import fnmatch +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import yaml + +from .models import ( + CredentialDefinition, + DomainScope, + Exclusions, + NetworkDefinition, + ScopeConfigModel, + ScopeMetadata, + ScopeSettings, + TargetDefinition, +) + + +class ScopeConfig: + """ + Scope configuration manager. + + Parses scope files (YAML/JSON) and provides methods for: + - Converting to existing targets_info format + - Checking if targets are in scope + - Retrieving credentials + - Getting exclusion rules for agent context + + Designed for future SQLite/Redis integration via to_dict()/from_dict(). + """ + + def __init__(self, model: ScopeConfigModel): + self.model = model + self._network_cidrs: dict[str, ipaddress.IPv4Network | ipaddress.IPv6Network] = {} + self._exclusion_cidrs: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = [] + self._parse_cidrs() + + def _parse_cidrs(self) -> None: + """Pre-parse CIDR notations for efficient scope checking.""" + for network in self.model.networks: + if network.cidr: + try: + self._network_cidrs[network.name] = ipaddress.ip_network( + network.cidr, strict=False + ) + except ValueError: + pass # Invalid CIDR, will be caught by validator + + for cidr in self.model.exclusions.cidrs: + try: + self._exclusion_cidrs.append(ipaddress.ip_network(cidr, strict=False)) + except ValueError: + pass + + @classmethod + def from_file(cls, path: Path | str) -> "ScopeConfig": + """Load scope configuration from YAML or JSON file.""" + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Scope file not found: {path}") + + content = path.read_text(encoding="utf-8") + + if path.suffix in (".yaml", ".yml"): + config_dict = yaml.safe_load(content) or {} + elif path.suffix == ".json": + config_dict = json.loads(content) + else: + raise ValueError(f"Unsupported scope file format: {path.suffix}. Use .yaml or .json") + + return cls.from_dict(config_dict) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "ScopeConfig": + """Create ScopeConfig from dictionary (for future database loading).""" + model = ScopeConfigModel( + metadata=ScopeMetadata(**data.get("metadata", {})), + settings=ScopeSettings(**data.get("settings", {})), + networks=[NetworkDefinition(**n) for n in data.get("networks", [])], + targets=[TargetDefinition(**t) for t in data.get("targets", [])], + exclusions=Exclusions(**data.get("exclusions", {})), + domains=DomainScope(**data.get("domains", {})), + ) + return cls(model) + + def to_dict(self) -> dict[str, Any]: + """Serialize to dictionary (for future database storage).""" + return self.model.model_dump() + + # ------------------------------------------------------------------------- + # Conversion to existing targets_info format + # ------------------------------------------------------------------------- + + def to_targets_info(self) -> list[dict[str, Any]]: + """ + Convert scope targets to existing targets_info format for compatibility + with current Strix CLI/agent infrastructure. + """ + targets_info = [] + + for target in self.model.targets: + target_type = target.get_target_type() + target_value = target.get_target_value() + + if not target_value: + continue + + # Map to existing format + if target_type == "infrastructure": + info = { + "type": "ip_address", + "details": {"target_ip": target_value}, + "original": target_value, + } + elif target_type in ("web_application", "api"): + info = { + "type": "web_application", + "details": {"target_url": target_value}, + "original": target_value, + } + elif target_type == "repository": + info = { + "type": "repository", + "details": {"target_repo": target_value}, + "original": target_value, + } + elif target_type == "local_code": + resolved_path = str(Path(target_value).resolve()) + info = { + "type": "local_code", + "details": {"target_path": resolved_path}, + "original": target_value, + } + else: + continue + + # Add extended metadata for agent context + info["scope_metadata"] = { + "name": target.name, + "network": target.network, + "ports": target.ports, + "services": [s.model_dump() for s in target.services], + "technologies": target.technologies, + "focus_areas": target.focus_areas, + "modules": target.modules, + "tags": target.tags, + "auth_type": target.auth_type, + "openapi_spec": target.openapi_spec, + "branch": target.branch, + } + + targets_info.append(info) + + return targets_info + + def get_instruction_context(self) -> str: + """ + Generate instruction context string from scope settings. + This supplements the --instruction flag. + """ + parts = [] + + # Operational mode + mode = self.model.settings.operational_mode + mode_text = { + "recon-only": "Recon only, no exploitation, generate PoCs but do not execute", + "poc-only": "Discovery and PoC validation only, no active exploitation", + "full-pentest": "Full penetration test, exploitation allowed within scope", + } + parts.append(mode_text.get(mode, "")) + + # Engagement type + eng_type = self.model.metadata.engagement_type + if eng_type == "internal": + parts.append("Internal network assessment") + elif eng_type == "external": + parts.append("External/perimeter assessment") + + # Validation requirement + if self.model.settings.require_validation: + parts.append("All findings must be validated before reporting") + + return ". ".join(filter(None, parts)) + + # ------------------------------------------------------------------------- + # Scope checking + # ------------------------------------------------------------------------- + + def is_in_scope(self, target: str) -> bool: + """ + Check if a discovered target is within scope. + + Handles: + - IP addresses (against network CIDRs and explicit hosts) + - URLs (against domain patterns) + - Ports (against port exclusions) + """ + # Check exclusions first + if self._is_excluded(target): + return False + + # Check if explicitly in targets + for t in self.model.targets: + if t.get_target_value() == target: + return True + + # Check if IP is in any network CIDR + try: + ip = ipaddress.ip_address(target) + for _name, network in self._network_cidrs.items(): + if ip in network: + return True + except ValueError: + pass # Not an IP + + # Check domain patterns for URLs + try: + parsed = urlparse(target) + if parsed.netloc: + return self._is_domain_in_scope(parsed.netloc) + except ValueError: + pass + + return False + + def is_port_in_scope(self, port: int) -> bool: + """Check if a port is allowed (not in exclusions).""" + return port not in self.model.exclusions.ports + + def is_path_in_scope(self, path: str) -> bool: + """Check if a URL path is allowed (not matching exclusion patterns).""" + for pattern in self.model.exclusions.paths: + if fnmatch(path, pattern): + return False + return True + + def _is_excluded(self, target: str) -> bool: + """Check if target matches any exclusion rule.""" + # Check host exclusions + if target in self.model.exclusions.hosts: + return True + + # Check URL exclusions + if target in self.model.exclusions.urls: + return True + + # Check CIDR exclusions + try: + ip = ipaddress.ip_address(target) + for excluded_cidr in self._exclusion_cidrs: + if ip in excluded_cidr: + return True + except ValueError: + pass + + # Check domain out_of_scope + try: + parsed = urlparse(target) + if parsed.netloc: + for pattern in self.model.domains.out_of_scope: + if fnmatch(parsed.netloc, pattern): + return True + except ValueError: + pass + + return False + + def _is_domain_in_scope(self, domain: str) -> bool: + """Check if domain matches in_scope patterns.""" + # First check out_of_scope + for pattern in self.model.domains.out_of_scope: + if fnmatch(domain, pattern): + return False + + # Then check in_scope + for pattern in self.model.domains.in_scope: + if fnmatch(domain, pattern): + return True + + # If no in_scope patterns defined, allow anything not excluded + if not self.model.domains.in_scope: + return True + + return False + + # ------------------------------------------------------------------------- + # Credentials + # ------------------------------------------------------------------------- + + def get_credentials(self, target: str) -> list[dict[str, str]]: + """ + Get credentials for a target, resolving env var references. + + Returns list of dicts with resolved username/password/token. + """ + credentials = [] + + for t in self.model.targets: + if t.get_target_value() != target: + continue + + for cred in t.credentials: + resolved = self._resolve_credential(cred) + if resolved: + credentials.append(resolved) + + # Also check target-level token_env + if t.token_env: + token = os.environ.get(t.token_env) + if token: + credentials.append({ + "type": "bearer", + "token": token, + "access_level": "unknown", + }) + + return credentials + + def _resolve_credential(self, cred: CredentialDefinition) -> dict[str, str] | None: + """Resolve a credential definition to actual values.""" + result: dict[str, str] = { + "username": cred.username, + "access_level": cred.access_level, + } + + if cred.password_env: + password = os.environ.get(cred.password_env) + if password: + result["password"] = password + result["type"] = "password" + else: + return None # Env var not set + + if cred.token_env: + token = os.environ.get(cred.token_env) + if token: + result["token"] = token + result["type"] = "token" + elif not cred.password_env: + return None # Neither password nor token available + + return result + + # ------------------------------------------------------------------------- + # Agent context + # ------------------------------------------------------------------------- + + def get_exclusion_rules(self) -> dict[str, Any]: + """Get exclusion rules formatted for agent context.""" + return { + "excluded_hosts": self.model.exclusions.hosts, + "excluded_cidrs": self.model.exclusions.cidrs, + "excluded_urls": self.model.exclusions.urls, + "excluded_paths": self.model.exclusions.paths, + "excluded_ports": self.model.exclusions.ports, + "out_of_scope_domains": self.model.domains.out_of_scope, + } + + def get_agent_context(self) -> dict[str, Any]: + """ + Get complete scope context for injection into agent prompts. + """ + return { + "engagement": { + "name": self.model.metadata.engagement_name, + "type": self.model.metadata.engagement_type, + }, + "settings": { + "mode": self.model.settings.operational_mode, + "require_validation": self.model.settings.require_validation, + "max_agents": self.model.settings.max_agents, + }, + "networks": [ + { + "name": n.name, + "type": n.type, + "cidr": n.cidr, + "vlan": n.vlan, + } + for n in self.model.networks + ], + "target_count": len(self.model.targets), + "exclusions": self.get_exclusion_rules(), + "in_scope_domains": self.model.domains.in_scope, + } + + def get_target_by_value(self, value: str) -> TargetDefinition | None: + """Look up target definition by its value (IP, URL, etc.).""" + for target in self.model.targets: + if target.get_target_value() == value: + return target + return None + + # ------------------------------------------------------------------------- + # Filtering + # ------------------------------------------------------------------------- + + def filter_targets( + self, + tags: list[str] | None = None, + network: str | None = None, + target_type: str | None = None, + ) -> list[TargetDefinition]: + """Filter targets by criteria.""" + results = [] + + for target in self.model.targets: + # Filter by tags (any match) + if tags: + if not any(t in target.tags for t in tags): + continue + + # Filter by network + if network and target.network != network: + continue + + # Filter by type + if target_type and target.get_target_type() != target_type: + continue + + results.append(target) + + return results + + # ------------------------------------------------------------------------- + # Properties + # ------------------------------------------------------------------------- + + @property + def metadata(self) -> ScopeMetadata: + return self.model.metadata + + @property + def settings(self) -> ScopeSettings: + return self.model.settings + + @property + def networks(self) -> list[NetworkDefinition]: + return self.model.networks + + @property + def targets(self) -> list[TargetDefinition]: + return self.model.targets + + @property + def exclusions(self) -> Exclusions: + return self.model.exclusions + + @property + def domains(self) -> DomainScope: + return self.model.domains diff --git a/strix/scope/validator.py b/strix/scope/validator.py new file mode 100644 index 000000000..2509cc7c6 --- /dev/null +++ b/strix/scope/validator.py @@ -0,0 +1,214 @@ +"""Scope configuration validation.""" + +import ipaddress +import os +from dataclasses import dataclass, field +from pathlib import Path + +from .models import ScopeConfigModel + + +@dataclass +class ValidationResult: + """Result of scope validation.""" + + valid: bool = True + errors: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + def add_error(self, message: str) -> None: + self.errors.append(message) + self.valid = False + + def add_warning(self, message: str) -> None: + self.warnings.append(message) + + def __bool__(self) -> bool: + return self.valid + + +class ScopeValidator: + """Validates scope configuration for correctness and security.""" + + def __init__(self, model: ScopeConfigModel): + self.model = model + self.result = ValidationResult() + + def validate(self) -> ValidationResult: + """Run all validation checks.""" + self._validate_targets() + self._validate_networks() + self._validate_cidrs() + self._validate_exclusions() + self._validate_credentials() + self._validate_network_references() + self._validate_modules() + return self.result + + def _validate_targets(self) -> None: + """Validate target definitions.""" + for i, target in enumerate(self.model.targets): + target_id = target.name or f"targets[{i}]" + + # Must have at least one identifier + if not any([target.host, target.url, target.repo, target.path]): + self.result.add_error( + f"{target_id}: Must specify at least one of: host, url, repo, path" + ) + + # Validate host is valid IP + if target.host: + try: + ipaddress.ip_address(target.host) + except ValueError: + self.result.add_error(f"{target_id}: Invalid IP address: {target.host}") + + # Validate URL format + if target.url: + if not target.url.startswith(("http://", "https://")): + self.result.add_warning( + f"{target_id}: URL should start with http:// or https://: {target.url}" + ) + + # Validate local path exists + if target.path: + path = Path(target.path).expanduser() + if not path.exists(): + self.result.add_warning(f"{target_id}: Local path does not exist: {target.path}") + elif not path.is_dir(): + self.result.add_error(f"{target_id}: Local path is not a directory: {target.path}") + + # Validate port ranges + for port in target.ports: + if not 1 <= port <= 65535: + self.result.add_error(f"{target_id}: Invalid port number: {port}") + + def _validate_networks(self) -> None: + """Validate network definitions.""" + names = set() + for i, network in enumerate(self.model.networks): + network_id = network.name or f"networks[{i}]" + + # Check for duplicate names + if network.name in names: + self.result.add_error(f"Duplicate network name: {network.name}") + names.add(network.name) + + # Validate VLAN range + if network.vlan is not None: + if not 1 <= network.vlan <= 4094: + self.result.add_error(f"{network_id}: Invalid VLAN ID: {network.vlan}") + + # Validate gateway is within CIDR + if network.cidr and network.gateway: + try: + net = ipaddress.ip_network(network.cidr, strict=False) + gw = ipaddress.ip_address(network.gateway) + if gw not in net: + self.result.add_warning( + f"{network_id}: Gateway {network.gateway} is not within CIDR {network.cidr}" + ) + except ValueError: + pass # CIDR validation handled separately + + def _validate_cidrs(self) -> None: + """Validate all CIDR notations.""" + # Network CIDRs + for network in self.model.networks: + if network.cidr: + try: + ipaddress.ip_network(network.cidr, strict=False) + except ValueError as e: + self.result.add_error(f"Invalid CIDR in network '{network.name}': {e}") + + # Exclusion CIDRs + for cidr in self.model.exclusions.cidrs: + try: + ipaddress.ip_network(cidr, strict=False) + except ValueError as e: + self.result.add_error(f"Invalid CIDR in exclusions: {cidr} - {e}") + + def _validate_exclusions(self) -> None: + """Validate exclusion rules.""" + # Validate excluded hosts are valid IPs + for host in self.model.exclusions.hosts: + try: + ipaddress.ip_address(host) + except ValueError: + self.result.add_warning( + f"Excluded host is not a valid IP: {host}. Will be treated as hostname." + ) + + # Validate excluded ports + for port in self.model.exclusions.ports: + if not 1 <= port <= 65535: + self.result.add_error(f"Invalid excluded port number: {port}") + + # Check for overlap between targets and exclusions + excluded_ips = set(self.model.exclusions.hosts) + for target in self.model.targets: + if target.host and target.host in excluded_ips: + self.result.add_warning( + f"Target '{target.name or target.host}' is also in exclusions list" + ) + + def _validate_credentials(self) -> None: + """Validate credential security and availability.""" + for target in self.model.targets: + target_id = target.name or target.get_target_value() + + for cred in target.credentials: + # Check password_env is set and available + if cred.password_env: + if not os.environ.get(cred.password_env): + self.result.add_warning( + f"{target_id}: Environment variable not set: {cred.password_env}" + ) + + # Check token_env is set and available + if cred.token_env: + if not os.environ.get(cred.token_env): + self.result.add_warning( + f"{target_id}: Environment variable not set: {cred.token_env}" + ) + + # Check target-level token_env + if target.token_env: + if not os.environ.get(target.token_env): + self.result.add_warning( + f"{target_id}: Environment variable not set: {target.token_env}" + ) + + def _validate_network_references(self) -> None: + """Validate that target network references exist.""" + network_names = {n.name for n in self.model.networks} + + for target in self.model.targets: + if target.network and target.network not in network_names: + self.result.add_error( + f"Target '{target.name or target.get_target_value()}' references " + f"undefined network: {target.network}" + ) + + def _validate_modules(self) -> None: + """Validate prompt module references.""" + try: + from strix.prompts import get_all_module_names + + available_modules = set(get_all_module_names()) + + for target in self.model.targets: + for module in target.modules: + if module not in available_modules: + self.result.add_warning( + f"Target '{target.name or target.get_target_value()}' " + f"references unknown module: {module}" + ) + except ImportError: + pass # Can't validate modules without strix.prompts + + +def validate_scope(model: ScopeConfigModel) -> ValidationResult: + """Convenience function to validate a scope configuration.""" + validator = ScopeValidator(model) + return validator.validate() diff --git a/strix/tools/__init__.py b/strix/tools/__init__.py index 8d5f896b4..01ac4f011 100644 --- a/strix/tools/__init__.py +++ b/strix/tools/__init__.py @@ -10,10 +10,12 @@ validate_tool_availability, ) from .registry import ( + TOOL_PROFILES, ImplementedInClientSideOnlyError, get_tool_by_name, get_tool_names, get_tools_prompt, + is_tool_allowed_for_role, needs_agent_state, register_tool, tools, @@ -48,6 +50,7 @@ __all__ = [ "ImplementedInClientSideOnlyError", + "TOOL_PROFILES", "execute_tool", "execute_tool_invocation", "execute_tool_with_validation", @@ -55,6 +58,7 @@ "get_tool_by_name", "get_tool_names", "get_tools_prompt", + "is_tool_allowed_for_role", "needs_agent_state", "process_tool_invocations", "register_tool", diff --git a/strix/tools/agents_graph/agents_graph_actions.py b/strix/tools/agents_graph/agents_graph_actions.py index 2e384c01c..f2fb9a8b2 100644 --- a/strix/tools/agents_graph/agents_graph_actions.py +++ b/strix/tools/agents_graph/agents_graph_actions.py @@ -191,6 +191,8 @@ def create_agent( name: str, inherit_context: bool = True, prompt_modules: str | None = None, + agent_role: str | None = None, + model: str | None = None, ) -> dict[str, Any]: try: parent_id = agent_state.agent_id @@ -228,19 +230,29 @@ def create_agent( from strix.agents.state import AgentState from strix.llm.config import LLMConfig - state = AgentState(task=task, agent_name=name, parent_id=parent_id, max_iterations=300) + state = AgentState( + task=task, agent_name=name, parent_id=parent_id, max_iterations=300, agent_role=agent_role + ) parent_agent = _agent_instances.get(parent_id) timeout = None - if ( - parent_agent - and hasattr(parent_agent, "llm_config") - and hasattr(parent_agent.llm_config, "timeout") - ): - timeout = parent_agent.llm_config.timeout - - llm_config = LLMConfig(prompt_modules=module_list, timeout=timeout) + parent_model = None + if parent_agent and hasattr(parent_agent, "llm_config"): + if hasattr(parent_agent.llm_config, "timeout"): + timeout = parent_agent.llm_config.timeout + if hasattr(parent_agent.llm_config, "model_name"): + parent_model = parent_agent.llm_config.model_name + + # Model fallback chain: explicit model -> parent model -> env default (in LLMConfig) + effective_model = model or parent_model + + llm_config = LLMConfig( + model_name=effective_model, + prompt_modules=module_list, + timeout=timeout, + agent_role=agent_role, + ) agent_config = { "llm_config": llm_config, diff --git a/strix/tools/agents_graph/agents_graph_actions_schema.xml b/strix/tools/agents_graph/agents_graph_actions_schema.xml index 3a36c269c..d1eec1490 100644 --- a/strix/tools/agents_graph/agents_graph_actions_schema.xml +++ b/strix/tools/agents_graph/agents_graph_actions_schema.xml @@ -82,40 +82,47 @@ Only create a new agent if no existing agent is handling the specific task. Comma-separated list of prompt modules to use for the agent (MAXIMUM 5 modules allowed). Most agents should have at least one module in order to be useful. Agents should be highly specialized - use 1-3 related modules; up to 5 for complex contexts. {{DYNAMIC_MODULES_DESCRIPTION}} + + Role determining which tools the agent can access. Available roles: recon (reconnaissance agents), testing (vulnerability testing), validation (PoC validation), reporting (vulnerability reporting), fixing (code fixes). Each role has a specialized toolset. + + + LLM model to use for this agent (e.g., "gemini/gemini-1.5-flash", "openai/gpt-4o"). If not specified, inherits from parent agent. If parent has no model, falls back to STRIX_LLM environment variable. Use faster/cheaper models for recon, more capable models for complex testing. + Response containing: - agent_id: Unique identifier for the created agent - success: Whether the agent was created successfully - message: Status message - agent_info: Details about the created agent - # After confirming no SQL testing agent exists, create agent for vulnerability validation + # Create a recon agent for infrastructure scanning - Validate and exploit the suspected SQL injection vulnerability found in - the login form. Confirm exploitability and document proof of concept. - SQLi Validator - sql_injection + Perform port scanning, service detection, and version enumeration on target 10.0.101.2 + Infrastructure Recon + proxmox_ve + recon + # Create a validation agent for SQL injection - Test authentication mechanisms, JWT implementation, and session management - for security vulnerabilities and bypass techniques. - Auth Specialist - authentication_jwt, business_logic + Validate and create PoC for the suspected SQL injection in login form + SQLi Validator + sql_injection + validation - # Example of single-module specialization (most focused) + # Create a reporting agent - Perform comprehensive XSS testing including reflected, stored, and DOM-based - variants across all identified input points. - XSS Specialist - xss + Document the confirmed SQL injection vulnerability with full details and remediation + SQLi Reporter + sql_injection + reporting - # Example of up to 5 related modules (borderline acceptable) + # Create a testing agent for authentication - Test for server-side vulnerabilities including SSRF, XXE, and potential - RCE vectors in file upload and XML processing endpoints. - Server-Side Attack Specialist - ssrf, xxe, rce + Test authentication mechanisms, JWT implementation, and session management + Auth Tester + authentication_jwt, business_logic + testing diff --git a/strix/tools/executor.py b/strix/tools/executor.py index 6dd1b04d7..1eefb255c 100644 --- a/strix/tools/executor.py +++ b/strix/tools/executor.py @@ -12,6 +12,7 @@ from .registry import ( get_tool_by_name, get_tool_names, + is_tool_allowed_for_role, needs_agent_state, should_execute_in_sandbox, ) @@ -97,20 +98,31 @@ async def _execute_tool_locally(tool_name: str, agent_state: Any | None, **kwarg return await result if inspect.isawaitable(result) else result -def validate_tool_availability(tool_name: str | None) -> tuple[bool, str]: +def validate_tool_availability( + tool_name: str | None, agent_state: Any | None = None +) -> tuple[bool, str]: if tool_name is None: return False, "Tool name is missing" if tool_name not in get_tool_names(): return False, f"Tool '{tool_name}' is not available" + # Runtime role enforcement + agent_role = None + if agent_state is not None and hasattr(agent_state, "agent_role"): + agent_role = agent_state.agent_role + + is_allowed, role_error = is_tool_allowed_for_role(tool_name, agent_role) + if not is_allowed: + return False, role_error + return True, "" async def execute_tool_with_validation( tool_name: str | None, agent_state: Any | None = None, **kwargs: Any ) -> Any: - is_valid, error_msg = validate_tool_availability(tool_name) + is_valid, error_msg = validate_tool_availability(tool_name, agent_state) if not is_valid: return f"Error: {error_msg}" diff --git a/strix/tools/registry.py b/strix/tools/registry.py index a12ae2b0f..e43b1bbc5 100644 --- a/strix/tools/registry.py +++ b/strix/tools/registry.py @@ -12,6 +12,76 @@ _tools_by_name: dict[str, Callable[..., Any]] = {} logger = logging.getLogger(__name__) +# Role-based tool profiles - None means all tools available +TOOL_PROFILES: dict[str, list[str] | None] = { + "root": [ + "create_agent", + "view_agent_graph", + "finish_scan", + "send_message_to_agent", + "wait_for_message", + "think", + ], + "recon": [ + "terminal", + "python", + "browser", + "proxy", + "think", + "agent_finish", + "create_agent", + "view_agent_graph", + "send_message_to_agent", + "wait_for_message", + "read_file", + "write_file", + "list_directory", + "web_search", + ], + "testing": [ + "terminal", + "python", + "browser", + "proxy", + "think", + "agent_finish", + "create_agent", + "view_agent_graph", + "send_message_to_agent", + "wait_for_message", + "read_file", + "write_file", + "web_search", + ], + "validation": [ + "terminal", + "python", + "browser", + "proxy", + "think", + "agent_finish", + "read_file", + "send_message_to_agent", + ], + "reporting": [ + "create_vulnerability_report", + "read_file", + "write_file", + "think", + "agent_finish", + "send_message_to_agent", + ], + "fixing": [ + "read_file", + "write_file", + "terminal", + "python", + "think", + "agent_finish", + "send_message_to_agent", + ], +} + class ImplementedInClientSideOnlyError(Exception): def __init__( @@ -168,9 +238,52 @@ def should_execute_in_sandbox(tool_name: str) -> bool: return True -def get_tools_prompt() -> str: +def is_tool_allowed_for_role(tool_name: str, role: str | None) -> tuple[bool, str]: + """Check if a tool is allowed for a given agent role. + + Args: + tool_name: Name of the tool to check + role: Agent role (e.g., 'root', 'recon', 'testing', 'validation', 'reporting', 'fixing') + + Returns: + Tuple of (is_allowed, error_message) + - If allowed: (True, "") + - If denied: (False, "error description") + """ + # No role specified = all tools allowed (backward compatibility) + if role is None: + return True, "" + + # Unknown role = all tools allowed (permissive for custom roles) + if role not in TOOL_PROFILES: + logger.warning(f"Unknown agent role '{role}', allowing all tools") + return True, "" + + allowed_tools = TOOL_PROFILES[role] + + # None in profile means all tools allowed + if allowed_tools is None: + return True, "" + + if tool_name in allowed_tools: + return True, "" + + return False, ( + f"Tool '{tool_name}' is not permitted for role '{role}'. " + f"Allowed tools: {', '.join(sorted(allowed_tools))}" + ) + + +def get_tools_prompt(role: str | None = None) -> str: + allowed_tools = None + if role and role in TOOL_PROFILES: + allowed_tools = TOOL_PROFILES[role] + tools_by_module: dict[str, list[dict[str, Any]]] = {} for tool in tools: + tool_name = tool.get("name", "") + if allowed_tools is not None and tool_name not in allowed_tools: + continue module = tool.get("module", "unknown") if module not in tools_by_module: tools_by_module[module] = [] diff --git a/templates/scope/proxmox-cluster.yaml b/templates/scope/proxmox-cluster.yaml new file mode 100644 index 000000000..41b5e9785 --- /dev/null +++ b/templates/scope/proxmox-cluster.yaml @@ -0,0 +1,102 @@ +# Proxmox Cluster Assessment Scope Template + +metadata: + engagement_name: "Proxmox Cluster Security Assessment" + engagement_type: "internal" + tester: "your-name" + +settings: + operational_mode: "recon-only" # Start with recon + require_validation: true + +networks: + - name: "Proxmox Management" + type: "internal" + vlan: 101 + cidr: "10.0.101.0/24" + gateway: "10.0.101.1" + description: "Proxmox cluster management network" + + - name: "VM Network" + type: "internal" + vlan: 102 + cidr: "10.0.102.0/24" + description: "Virtual machine network" + + - name: "Storage Network" + type: "internal" + vlan: 103 + cidr: "10.0.103.0/24" + description: "Ceph/storage cluster network" + +targets: + # Proxmox nodes + - host: "10.0.101.2" + name: "pve-node-01" + type: "infrastructure" + network: "Proxmox Management" + ports: [22, 8006, 3128, 5900] # VNC starts at 5900 + services: + - port: 8006 + service: "proxmox-ve" + version: "8.x" + - port: 3128 + service: "spiceproxy" + credentials: + - username: "root@pam" + password_env: "PVE_ROOT_PASS" + access_level: "admin" + tags: ["hypervisor", "cluster-master", "critical"] + modules: ["proxmox_ve"] + + - host: "10.0.101.3" + name: "pve-node-02" + type: "infrastructure" + network: "Proxmox Management" + ports: [22, 8006, 3128] + tags: ["hypervisor", "cluster-member"] + modules: ["proxmox_ve"] + + - host: "10.0.101.4" + name: "pve-node-03" + type: "infrastructure" + network: "Proxmox Management" + ports: [22, 8006, 3128] + tags: ["hypervisor", "cluster-member"] + modules: ["proxmox_ve"] + + # PBS (Proxmox Backup Server) if present + - host: "10.0.101.10" + name: "pbs-backup-01" + type: "infrastructure" + network: "Proxmox Management" + ports: [22, 8007] + services: + - port: 8007 + service: "proxmox-backup" + tags: ["backup", "critical"] + modules: ["proxmox_ve"] + +exclusions: + hosts: + - "10.0.101.1" # Gateway + - "10.0.101.254" # Monitoring/IPMI + cidrs: + - "10.0.103.0/24" # Storage network - don't scan (performance impact) + ports: + - 111 # rpcbind + - 2049 # NFS + - 6789 # Ceph monitor + - 6800 # Ceph OSD range start (6800-7300) + +# Focus areas for Proxmox testing +test_focus: + - "API authentication bypass" + - "Privilege escalation (user to admin)" + - "VM escape vectors" + - "Container (LXC) escape" + - "Cluster authentication (corosync)" + - "Default credentials" + - "Known CVEs (CVE-2022-35508, etc.)" + - "Backup access controls" + - "Storage permission misconfigurations" diff --git a/templates/scope/scope-simple.csv b/templates/scope/scope-simple.csv new file mode 100644 index 000000000..74989b03c --- /dev/null +++ b/templates/scope/scope-simple.csv @@ -0,0 +1,8 @@ +type,target,name,network,ports,tags,focus_areas,credentials_env +infrastructure,10.0.10.5,Web Server,Production,22;80;443,web;production,, +infrastructure,10.0.10.6,Database Server,Production,22;3306;5432,database;critical,, +web_application,https://app.example.com,Main Application,DMZ,,customer-facing,authentication;idor,APP_TEST_USER:APP_TEST_PASS +web_application,https://admin.example.com,Admin Panel,DMZ,,admin;critical,authentication;business_logic,ADMIN_USER:ADMIN_PASS +api,https://api.example.com,REST API,DMZ,,api,sql_injection;idor,API_TOKEN +repository,https://github.com/example/backend,Backend Code,,,code,sql_injection;secrets, +local_code,./frontend,Frontend Code,,,code,xss;csrf, diff --git a/templates/scope/scope.json b/templates/scope/scope.json new file mode 100644 index 000000000..ce9c45a31 --- /dev/null +++ b/templates/scope/scope.json @@ -0,0 +1,86 @@ +{ + "metadata": { + "engagement_name": "Example Engagement", + "engagement_type": "internal", + "start_date": "2024-01-01", + "end_date": "2024-01-31", + "tester": "your-name" + }, + "settings": { + "operational_mode": "poc-only", + "max_agents": 20, + "require_validation": true, + "generate_fixes": false + }, + "networks": [ + { + "name": "Production Network", + "type": "internal", + "vlan": 10, + "cidr": "10.0.10.0/24", + "gateway": "10.0.10.1", + "description": "Production servers" + }, + { + "name": "DMZ", + "type": "external", + "cidr": "203.0.113.0/24", + "description": "Public-facing servers" + } + ], + "targets": [ + { + "host": "10.0.10.5", + "name": "Web Server", + "type": "infrastructure", + "network": "Production Network", + "ports": [22, 80, 443], + "tags": ["web", "production"] + }, + { + "url": "https://app.example.com", + "name": "Main Application", + "type": "web_application", + "network": "DMZ", + "technologies": ["Django", "PostgreSQL"], + "credentials": [ + { + "username": "testuser", + "password_env": "APP_TEST_PASSWORD", + "access_level": "user" + } + ], + "focus_areas": ["authentication", "idor"], + "tags": ["customer-facing"] + }, + { + "url": "https://api.example.com", + "name": "REST API", + "type": "api", + "auth_type": "bearer", + "token_env": "API_TOKEN", + "tags": ["api"] + }, + { + "repo": "https://github.com/example/app", + "name": "Application Source", + "type": "repository", + "branch": "main", + "focus_areas": ["sql_injection", "secrets"] + } + ], + "exclusions": { + "hosts": ["10.0.10.1", "10.0.10.254"], + "cidrs": ["10.0.20.0/24"], + "urls": [ + "https://app.example.com/health", + "https://app.example.com/metrics" + ], + "paths": ["/api/v1/internal/*"], + "ports": [161, 162] + }, + "domains": { + "in_scope": ["*.example.com", "*.example-staging.com"], + "out_of_scope": ["mail.example.com", "vpn.example.com"] + } +} diff --git a/templates/scope/scope.yaml b/templates/scope/scope.yaml new file mode 100644 index 000000000..17e544f3e --- /dev/null +++ b/templates/scope/scope.yaml @@ -0,0 +1,100 @@ +# Strix Scope Configuration Template +# Copy this file and customize for your engagement + +metadata: + engagement_name: "Example Engagement" + engagement_type: "internal" # internal | external | hybrid + start_date: "2024-01-01" + end_date: "2024-01-31" + tester: "your-name" + notes: "Additional engagement notes" + +settings: + operational_mode: "poc-only" # recon-only | poc-only | full-pentest + max_agents: 20 + require_validation: true + generate_fixes: false + +# Define networks/VLANs in scope +networks: + - name: "Production Network" + type: "internal" + vlan: 10 + cidr: "10.0.10.0/24" + gateway: "10.0.10.1" + description: "Production servers" + + - name: "DMZ" + type: "external" + cidr: "203.0.113.0/24" + description: "Public-facing servers" + +# Specific targets to test +targets: + # Infrastructure example + - host: "10.0.10.5" + name: "Web Server" + type: "infrastructure" + network: "Production Network" + ports: [22, 80, 443] + tags: ["web", "production"] + + # Web application example + - url: "https://app.example.com" + name: "Main Application" + type: "web_application" + network: "DMZ" + technologies: ["Django", "PostgreSQL"] + credentials: + - username: "testuser" + password_env: "APP_TEST_PASSWORD" # Set this env var + access_level: "user" + focus_areas: ["authentication", "idor"] + tags: ["customer-facing"] + + # API example + - url: "https://api.example.com" + name: "REST API" + type: "api" + auth_type: "bearer" + token_env: "API_TOKEN" + openapi_spec: "./api-spec.yaml" # Optional + tags: ["api"] + + # Repository example + - repo: "https://github.com/example/app" + name: "Application Source" + type: "repository" + branch: "main" + focus_areas: ["sql_injection", "secrets"] + + # Local code example + - path: "./src" + name: "Local Codebase" + type: "local_code" + focus_areas: ["xss", "csrf"] + +# What NOT to test +exclusions: + hosts: + - "10.0.10.1" # Gateway - do not test + - "10.0.10.254" # Monitoring - do not test + cidrs: + - "10.0.20.0/24" # Out of scope network + urls: + - "https://app.example.com/health" + - "https://app.example.com/metrics" + paths: + - "/api/v1/internal/*" + ports: + - 161 # SNMP + - 162 + +# Domain boundaries for web testing +domains: + in_scope: + - "*.example.com" + - "*.example-staging.com" + out_of_scope: + - "mail.example.com" + - "vpn.example.com" From 7b95699dd55d1137b91c79fd352bd20e5c4ee5c4 Mon Sep 17 00:00:00 2001 From: yokoszn Date: Sun, 23 Nov 2025 08:49:22 +1100 Subject: [PATCH 2/3] refactor: simplify state management and improve robustness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit State Management: - Add AgentStatus enum replacing boolean flags (completed, stop_requested, waiting_for_input, llm_failed) for single source of truth - Add always-on timeout (max_wait_seconds=300) to prevent indefinite stalls - Add consecutive_empty_responses tracking with MAX_EMPTY_RESPONSES=3 limit - Add failure_reason field for explicit failure tracking - Add backward-compatible properties for existing code LLM Improvements: - Add retry with exponential backoff for transient errors (rate limits, timeouts, service unavailable) - Add RETRYABLE_ERRORS tuple and MAX_RETRIES=3 configuration - Add real-time event emission for LLM requests/responses Code Quality Fixes: - Add thread locks (_graph_lock) to global mutable state in agents_graph_actions - Fix ProxyManager singleton caching bug (was creating new instance each call) - Add queue timeouts (QUEUE_TIMEOUT=120) in tool_server worker - Convert request_queue from threading to asyncio primitives - Add container cleanup on failure in docker_runtime - Re-enable file-based logging in tool_server worker processes New Features: - Add real-time event streaming with verbose mode (-v/--verbose flag) - Add --debug flag for full LiteLLM debug output - Add progress tracking tools (save_progress, load_progress, list_progress) - Add crash-proof persistence via events.jsonl πŸ€– Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- docker-compose.dev.yml | 4 + strix/agents/__init__.py | 3 +- strix/agents/base_agent.py | 193 ++++++---- strix/agents/state.py | 88 +++-- strix/interface/cli.py | 141 +++++++- strix/interface/main.py | 35 ++ strix/llm/__init__.py | 7 +- strix/llm/llm.py | 255 +++++++++---- strix/llm/request_queue.py | 26 +- strix/runtime/docker_runtime.py | 43 ++- strix/runtime/tool_server.py | 26 +- strix/telemetry/tracer.py | 336 ++++++++++++++++++ strix/tools/__init__.py | 1 + .../agents_graph/agents_graph_actions.py | 320 +++++++++-------- strix/tools/executor.py | 100 +++++- strix/tools/notes/notes_actions.py | 90 +++++ strix/tools/progress/__init__.py | 5 + strix/tools/progress/progress_actions.py | 243 +++++++++++++ strix/tools/proxy/proxy_manager.py | 3 +- strix/tools/registry.py | 120 +++++-- strix/tools/reporting/reporting_actions.py | 72 +++- templates/scope/proxmox-cluster.yaml | 47 ++- 22 files changed, 1775 insertions(+), 383 deletions(-) create mode 100644 strix/tools/progress/__init__.py create mode 100644 strix/tools/progress/progress_actions.py diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index ada40c2df..85f8bd50f 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -10,6 +10,10 @@ services: - ./strix:/app/strix:ro # Mount target directory (change as needed) - ./target:/workspace:ro + # Mount templates for scope files + - ./templates:/app/templates:ro + # Mount agent_runs for persistent results + - ./agent_runs:/app/agent_runs # Docker socket for spawning sandbox containers - /var/run/docker.sock:/var/run/docker.sock environment: diff --git a/strix/agents/__init__.py b/strix/agents/__init__.py index c7e542e46..4a8639f2b 100644 --- a/strix/agents/__init__.py +++ b/strix/agents/__init__.py @@ -1,10 +1,11 @@ from .base_agent import BaseAgent -from .state import AgentState +from .state import AgentState, AgentStatus from .StrixAgent import StrixAgent __all__ = [ "AgentState", + "AgentStatus", "BaseAgent", "StrixAgent", ] diff --git a/strix/agents/base_agent.py b/strix/agents/base_agent.py index 67aeb3838..f2afee163 100644 --- a/strix/agents/base_agent.py +++ b/strix/agents/base_agent.py @@ -18,7 +18,11 @@ from strix.llm.utils import clean_content from strix.tools import process_tool_invocations -from .state import AgentState +from .state import AgentState, AgentStatus + + +# Maximum consecutive empty responses before failing +MAX_EMPTY_RESPONSES = 3 logger = logging.getLogger(__name__) @@ -131,8 +135,8 @@ def _add_to_agents_graph(self) -> None: } agents_graph_actions._agent_graph["nodes"][self.state.agent_id] = node + # Store agent instance only - state is accessible via agent.state agents_graph_actions._agent_instances[self.state.agent_id] = self - agents_graph_actions._agent_states[self.state.agent_id] = self.state if self.state.parent_id: agents_graph_actions._agent_graph["edges"].append( @@ -151,6 +155,7 @@ def cancel_current_execution(self) -> None: self._current_task = None async def agent_loop(self, task: str) -> dict[str, Any]: # noqa: PLR0912, PLR0915 + """Main agent loop using status-based dispatch.""" await self._initialize_sandbox_and_state(task) from strix.telemetry.tracer import get_global_tracer @@ -158,50 +163,48 @@ async def agent_loop(self, task: str) -> dict[str, Any]: # noqa: PLR0912, PLR09 tracer = get_global_tracer() while True: + # Check for inter-agent messages first self._check_agent_messages(self.state) - if self.state.is_waiting_for_input(): + # Status-based dispatch + status = self.state.status + + if status == AgentStatus.COMPLETED: + return self._finalize_agent(tracer, success=True) + + if status == AgentStatus.FAILED: + return self._finalize_agent(tracer, success=False) + + if status == AgentStatus.STOPPED: + if self.non_interactive: + return self.state.final_result or {} + await self._enter_waiting_state(tracer, was_cancelled=True) + continue + + if status in (AgentStatus.WAITING_FOR_MESSAGE, AgentStatus.WAITING_FOR_RECOVERY): await self._wait_for_input() continue + # RUNNING status - check if we should stop if self.state.should_stop(): if self.non_interactive: return self.state.final_result or {} await self._enter_waiting_state(tracer) continue - if self.state.llm_failed: - await self._wait_for_input() - continue - + # Execute an iteration self.state.increment_iteration() - if ( - self.state.is_approaching_max_iterations() - and not self.state.max_iterations_warning_sent - ): - self.state.max_iterations_warning_sent = True - remaining = self.state.max_iterations - self.state.iteration - warning_msg = ( - f"URGENT: You are approaching the maximum iteration limit. " - f"Current: {self.state.iteration}/{self.state.max_iterations} " - f"({remaining} iterations remaining). " - f"Please prioritize completing your required task(s) and calling " - f"the appropriate finish tool (finish_scan for root agent, " - f"agent_finish for sub-agents) as soon as possible." - ) - self.state.add_message("user", warning_msg) - - if self.state.iteration == self.state.max_iterations - 3: - final_warning_msg = ( - "CRITICAL: You have only 3 iterations left! " - "Your next message MUST be the tool call to the appropriate " - "finish tool: finish_scan if you are the root agent, or " - "agent_finish if you are a sub-agent. " - "No other actions should be taken except finishing your work " - "immediately." + # Emit iteration event + if tracer: + tracer.log_agent_iteration( + agent_id=self.state.agent_id, + iteration=self.state.iteration, + max_iterations=self.state.max_iterations, ) - self.state.add_message("user", final_warning_msg) + + # Send warnings as we approach max iterations + self._check_iteration_warnings() try: should_finish = await self._process_iteration(tracer) @@ -221,37 +224,9 @@ async def agent_loop(self, task: str) -> dict[str, Any]: # noqa: PLR0912, PLR09 continue except LLMRequestFailedError as e: - error_msg = str(e) - error_details = getattr(e, "details", None) - self.state.add_error(error_msg) - + self._handle_llm_error(e, tracer) if self.non_interactive: - self.state.set_completed({"success": False, "error": error_msg}) - if tracer: - tracer.update_agent_status(self.state.agent_id, "failed", error_msg) - if error_details: - tracer.log_tool_execution_start( - self.state.agent_id, - "llm_error_details", - {"error": error_msg, "details": error_details}, - ) - tracer.update_tool_execution( - tracer._next_execution_id - 1, "failed", error_details - ) - return {"success": False, "error": error_msg} - - self.state.enter_waiting_state(llm_failed=True) - if tracer: - tracer.update_agent_status(self.state.agent_id, "llm_failed", error_msg) - if error_details: - tracer.log_tool_execution_start( - self.state.agent_id, - "llm_error_details", - {"error": error_msg, "details": error_details}, - ) - tracer.update_tool_execution( - tracer._next_execution_id - 1, "failed", error_details - ) + return {"success": False, "error": str(e)} continue except (RuntimeError, ValueError, TypeError) as e: @@ -264,6 +239,65 @@ async def agent_loop(self, task: str) -> dict[str, Any]: # noqa: PLR0912, PLR09 await self._enter_waiting_state(tracer, error_occurred=True) continue + def _finalize_agent(self, tracer: Optional["Tracer"], success: bool) -> dict[str, Any]: + """Finalize agent and return result.""" + if tracer: + status_str = "completed" if success else "failed" + tracer.update_agent_status(self.state.agent_id, status_str, self.state.failure_reason) + return self.state.final_result or {"success": success} + + def _check_iteration_warnings(self) -> None: + """Send warnings as agent approaches max iterations.""" + if ( + self.state.is_approaching_max_iterations() + and not self.state.max_iterations_warning_sent + ): + self.state.max_iterations_warning_sent = True + remaining = self.state.max_iterations - self.state.iteration + warning_msg = ( + f"URGENT: You are approaching the maximum iteration limit. " + f"Current: {self.state.iteration}/{self.state.max_iterations} " + f"({remaining} iterations remaining). " + f"Please prioritize completing your required task(s) and calling " + f"the appropriate finish tool (finish_scan for root agent, " + f"agent_finish for sub-agents) as soon as possible." + ) + self.state.add_message("user", warning_msg) + + if self.state.iteration == self.state.max_iterations - 3: + final_warning_msg = ( + "CRITICAL: You have only 3 iterations left! " + "Your next message MUST be the tool call to the appropriate " + "finish tool: finish_scan if you are the root agent, or " + "agent_finish if you are a sub-agent. " + "No other actions should be taken except finishing your work " + "immediately." + ) + self.state.add_message("user", final_warning_msg) + + def _handle_llm_error(self, e: LLMRequestFailedError, tracer: Optional["Tracer"]) -> None: + """Handle LLM request failure.""" + error_msg = str(e) + error_details = getattr(e, "details", None) + self.state.add_error(error_msg) + + if self.non_interactive: + self.state.set_failed(error_msg) + else: + self.state.enter_waiting_state(llm_failed=True) + + if tracer: + tracer.update_agent_status(self.state.agent_id, "llm_failed", error_msg) + if error_details: + tracer.log_tool_execution_start( + self.state.agent_id, + "llm_error_details", + {"error": error_msg, "details": error_details}, + ) + tracer.update_tool_execution( + tracer._next_execution_id - 1, "failed", error_details + ) + async def _wait_for_input(self) -> None: import asyncio @@ -296,18 +330,35 @@ async def _enter_waiting_state( error_occurred: bool = False, was_cancelled: bool = False, ) -> None: + old_state = "running" self.state.enter_waiting_state() if tracer: if task_completed: + new_state = "completed" + reason = "Task completed successfully" tracer.update_agent_status(self.state.agent_id, "completed") elif error_occurred: + new_state = "error" + reason = "An error occurred" tracer.update_agent_status(self.state.agent_id, "error") elif was_cancelled: + new_state = "stopped" + reason = "Execution cancelled by user" tracer.update_agent_status(self.state.agent_id, "stopped") else: + new_state = "waiting" + reason = "Execution paused" tracer.update_agent_status(self.state.agent_id, "stopped") + # Emit state transition event + tracer.log_agent_state_transition( + agent_id=self.state.agent_id, + from_state=old_state, + to_state=new_state, + reason=reason, + ) + if task_completed: self.state.add_message( "assistant", @@ -356,7 +407,23 @@ async def _process_iteration(self, tracer: Optional["Tracer"]) -> bool: content_stripped = (response.content or "").strip() if not content_stripped: + # Track consecutive empty responses + self.state.consecutive_empty_responses += 1 + + if self.state.consecutive_empty_responses >= MAX_EMPTY_RESPONSES: + # Fail after too many empty responses + error_msg = ( + f"Agent failed: {MAX_EMPTY_RESPONSES} consecutive empty responses. " + "The LLM is not responding properly." + ) + logger.error(error_msg) + self.state.set_failed(error_msg) + if tracer: + tracer.update_agent_status(self.state.agent_id, "failed", error_msg) + return True # Signal to exit loop + corrective_message = ( + f"WARNING: Empty response ({self.state.consecutive_empty_responses}/{MAX_EMPTY_RESPONSES}). " "You MUST NOT respond with empty messages. " "If you currently have nothing to do or say, use an appropriate tool instead:\n" "- Use agents_graph_actions.wait_for_message to wait for messages " @@ -369,6 +436,8 @@ async def _process_iteration(self, tracer: Optional["Tracer"]) -> bool: self.state.add_message("user", corrective_message) return False + # Reset empty response counter on successful response + self.state.consecutive_empty_responses = 0 self.state.add_message("assistant", response.content) if tracer: tracer.log_chat_message( diff --git a/strix/agents/state.py b/strix/agents/state.py index 8f2a1afcb..ef0ecbe69 100644 --- a/strix/agents/state.py +++ b/strix/agents/state.py @@ -1,5 +1,6 @@ import uuid from datetime import UTC, datetime +from enum import Enum from typing import Any from pydantic import BaseModel, Field @@ -9,6 +10,17 @@ def _generate_agent_id() -> str: return f"agent_{uuid.uuid4().hex[:8]}" +class AgentStatus(Enum): + """Explicit agent status - single source of truth for agent state.""" + + RUNNING = "running" + WAITING_FOR_MESSAGE = "waiting_for_message" + WAITING_FOR_RECOVERY = "waiting_for_recovery" # LLM retry pending + COMPLETED = "completed" + FAILED = "failed" + STOPPED = "stopped" + + class AgentState(BaseModel): agent_id: str = Field(default_factory=_generate_agent_id) agent_name: str = "Strix Agent" @@ -21,13 +33,15 @@ class AgentState(BaseModel): task: str = "" iteration: int = 0 max_iterations: int = 300 - completed: bool = False - stop_requested: bool = False - waiting_for_input: bool = False - llm_failed: bool = False + max_wait_seconds: int = 300 # 5 minutes default, always-on timeout + + # Single status field replaces: completed, stop_requested, waiting_for_input, llm_failed + status: AgentStatus = AgentStatus.RUNNING waiting_start_time: datetime | None = None final_result: dict[str, Any] | None = None + failure_reason: str | None = None # Reason for FAILED status max_iterations_warning_sent: bool = False + consecutive_empty_responses: int = 0 # Track empty LLM responses messages: list[dict[str, Any]] = Field(default_factory=list) context: dict[str, Any] = Field(default_factory=dict) @@ -75,32 +89,46 @@ def update_context(self, key: str, value: Any) -> None: self.last_updated = datetime.now(UTC).isoformat() def set_completed(self, final_result: dict[str, Any] | None = None) -> None: - self.completed = True + self.status = AgentStatus.COMPLETED self.final_result = final_result + self.waiting_start_time = None + self.last_updated = datetime.now(UTC).isoformat() + + def set_failed(self, reason: str) -> None: + """Mark agent as failed with a reason.""" + self.status = AgentStatus.FAILED + self.failure_reason = reason + self.waiting_start_time = None self.last_updated = datetime.now(UTC).isoformat() def request_stop(self) -> None: - self.stop_requested = True + self.status = AgentStatus.STOPPED + self.waiting_start_time = None self.last_updated = datetime.now(UTC).isoformat() def should_stop(self) -> bool: - return self.stop_requested or self.completed or self.has_reached_max_iterations() + """Check if agent should exit its loop.""" + return ( + self.status in (AgentStatus.COMPLETED, AgentStatus.FAILED, AgentStatus.STOPPED) + or self.has_reached_max_iterations() + ) def is_waiting_for_input(self) -> bool: - return self.waiting_for_input + return self.status in (AgentStatus.WAITING_FOR_MESSAGE, AgentStatus.WAITING_FOR_RECOVERY) def enter_waiting_state(self, llm_failed: bool = False) -> None: - self.waiting_for_input = True + if llm_failed: + self.status = AgentStatus.WAITING_FOR_RECOVERY + else: + self.status = AgentStatus.WAITING_FOR_MESSAGE self.waiting_start_time = datetime.now(UTC) - self.llm_failed = llm_failed self.last_updated = datetime.now(UTC).isoformat() def resume_from_waiting(self, new_task: str | None = None) -> None: - self.waiting_for_input = False + self.status = AgentStatus.RUNNING self.waiting_start_time = None - self.stop_requested = False - self.completed = False - self.llm_failed = False + self.failure_reason = None + self.consecutive_empty_responses = 0 if new_task: self.task = new_task self.last_updated = datetime.now(UTC).isoformat() @@ -112,19 +140,29 @@ def is_approaching_max_iterations(self, threshold: float = 0.85) -> bool: return self.iteration >= int(self.max_iterations * threshold) def has_waiting_timeout(self) -> bool: - if not self.waiting_for_input or not self.waiting_start_time: - return False - - if ( - self.stop_requested - or self.llm_failed - or self.completed - or self.has_reached_max_iterations() - ): + """Check if waiting state has exceeded max_wait_seconds. Always-on timeout.""" + if not self.is_waiting_for_input() or not self.waiting_start_time: return False elapsed = (datetime.now(UTC) - self.waiting_start_time).total_seconds() - return elapsed > 600 + return elapsed > self.max_wait_seconds + + # Backward compatibility properties + @property + def completed(self) -> bool: + return self.status == AgentStatus.COMPLETED + + @property + def stop_requested(self) -> bool: + return self.status == AgentStatus.STOPPED + + @property + def waiting_for_input(self) -> bool: + return self.is_waiting_for_input() + + @property + def llm_failed(self) -> bool: + return self.status == AgentStatus.WAITING_FOR_RECOVERY def has_empty_last_messages(self, count: int = 3) -> bool: if len(self.messages) < count: @@ -152,7 +190,9 @@ def get_execution_summary(self) -> dict[str, Any]: "task": self.task, "iteration": self.iteration, "max_iterations": self.max_iterations, + "status": self.status.value, "completed": self.completed, + "failure_reason": self.failure_reason, "final_result": self.final_result, "start_time": self.start_time, "last_updated": self.last_updated, diff --git a/strix/interface/cli.py b/strix/interface/cli.py index cee873bc5..9a47928b9 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -1,7 +1,8 @@ import atexit +import os import signal import sys -from typing import Any +from typing import TYPE_CHECKING, Any from rich.console import Console from rich.panel import Panel @@ -9,11 +10,137 @@ from strix.agents.StrixAgent import StrixAgent from strix.llm.config import LLMConfig -from strix.telemetry.tracer import Tracer, set_global_tracer +from strix.telemetry.tracer import EventType, Tracer, TracerEvent, set_global_tracer from .utils import get_severity_color +if TYPE_CHECKING: + pass + + +def _format_event_for_cli(event: TracerEvent, console: Console) -> None: + """Format and print a tracer event to the CLI.""" + agent_id = event.agent_id or "system" + # Truncate agent_id for display + agent_display = agent_id[:12] if len(agent_id) > 12 else agent_id + data = event.data + + if event.event_type == EventType.AGENT_ITERATION: + iteration = data.get("iteration", 0) + max_iter = data.get("max_iterations", 0) + progress = data.get("progress_pct", 0) + console.print( + f"[dim cyan][{agent_display}][/] " + f"[bold]Iteration {iteration}/{max_iter}[/] ({progress}%)" + ) + + elif event.event_type == EventType.LLM_REQUEST: + model = data.get("model", "unknown") + console.print( + f"[dim cyan][{agent_display}][/] " + f"[yellow]β†’ LLM request[/] ({model})" + ) + + elif event.event_type == EventType.LLM_RESPONSE: + input_tok = data.get("input_tokens", 0) + output_tok = data.get("output_tokens", 0) + duration = data.get("duration_ms", 0) + cost = data.get("cost") + cached = data.get("cached_tokens", 0) + + cost_str = f", ${cost:.4f}" if cost else "" + cache_str = f", {cached} cached" if cached > 0 else "" + + console.print( + f"[dim cyan][{agent_display}][/] " + f"[green]← LLM response[/] " + f"({input_tok}+{output_tok} tokens, {duration:.0f}ms{cost_str}{cache_str})" + ) + + elif event.event_type == EventType.LLM_ERROR: + error = data.get("error", "unknown error") + duration = data.get("duration_ms") + duration_str = f" ({duration:.0f}ms)" if duration else "" + # Show full error with proper formatting + console.print( + f"[dim cyan][{agent_display}][/] " + f"[bold red]βœ— LLM error{duration_str}:[/]" + ) + # Print error details on separate lines for readability + for line in str(error).split("\n"): + if line.strip(): + console.print(f" [red]{line}[/]") + + elif event.event_type == EventType.TOOL_START: + tool_name = data.get("tool_name", "unknown") + args = data.get("args", data.get("args_preview", {})) + # Truncate args for display + args_str = str(args) + if len(args_str) > 100: + args_str = args_str[:100] + "..." + console.print( + f"[dim cyan][{agent_display}][/] " + f"[bold magenta]β†’ {tool_name}[/] {args_str}" + ) + + elif event.event_type == EventType.TOOL_COMPLETE: + tool_name = data.get("tool_name", "unknown") + duration = data.get("duration_ms") + duration_str = f" ({duration:.0f}ms)" if duration else "" + console.print( + f"[dim cyan][{agent_display}][/] " + f"[green]βœ“ {tool_name} completed[/]{duration_str}" + ) + + elif event.event_type == EventType.TOOL_ERROR: + tool_name = data.get("tool_name", "unknown") + error = data.get("error", "unknown error") + duration = data.get("duration_ms") + duration_str = f" ({duration:.0f}ms)" if duration else "" + # Show full error with proper formatting + console.print( + f"[dim cyan][{agent_display}][/] " + f"[bold red]βœ— {tool_name} error{duration_str}:[/]" + ) + # Print error details on separate line for readability + for line in str(error).split("\n"): + console.print(f" [red]{line}[/]") + + elif event.event_type == EventType.AGENT_STATE_TRANSITION: + from_state = data.get("from_state", "?") + to_state = data.get("to_state", "?") + reason = data.get("reason", "") + reason_str = f" - {reason}" if reason else "" + console.print( + f"[dim cyan][{agent_display}][/] " + f"[blue]State: {from_state} β†’ {to_state}[/]{reason_str}" + ) + + elif event.event_type == EventType.AGENT_MESSAGE_SENT: + to_agent = data.get("to_agent_id", "?")[:12] + msg_preview = data.get("message_preview", "")[:50] + console.print( + f"[dim cyan][{agent_display}][/] " + f"[yellow]πŸ“€ β†’ {to_agent}:[/] {msg_preview}..." + ) + + elif event.event_type == EventType.AGENT_MESSAGE_RECEIVED: + from_agent = data.get("from_agent_id", "?")[:12] + msg_preview = data.get("message_preview", "")[:50] + console.print( + f"[dim cyan][{agent_display}][/] " + f"[green]πŸ“₯ ← {from_agent}:[/] {msg_preview}..." + ) + + elif event.event_type == EventType.AGENT_CREATED: + name = data.get("name", "unknown") + console.print( + f"[dim cyan][{agent_display}][/] " + f"[bold green]+ Agent created:[/] {name}" + ) + + async def run_cli(args: Any) -> None: # noqa: PLR0915 console = Console() @@ -120,6 +247,16 @@ def display_vulnerability(report_id: str, title: str, content: str, severity: st tracer.vulnerability_found_callback = display_vulnerability + # Enable real-time event streaming in verbose mode + verbose_mode = getattr(args, "verbose", False) or os.getenv("STRIX_VERBOSE", "").lower() == "true" + if verbose_mode: + console.print("[dim]Verbose mode enabled - showing all agent events[/]\n") + + def event_callback(event: TracerEvent) -> None: + _format_event_for_cli(event, console) + + tracer.event_callback = event_callback + def cleanup_on_exit() -> None: tracer.cleanup() diff --git a/strix/interface/main.py b/strix/interface/main.py index f632eb6b2..3ab4007f8 100644 --- a/strix/interface/main.py +++ b/strix/interface/main.py @@ -10,6 +10,7 @@ import shutil import sys from pathlib import Path +from typing import Any import litellm from docker.errors import DockerException @@ -331,6 +332,26 @@ def parse_arguments() -> argparse.Namespace: ), ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help=( + "Enable verbose output showing real-time agent events " + "(iterations, LLM requests/responses, tool executions, etc.). " + "Best used with -n/--non-interactive mode." + ), + ) + + parser.add_argument( + "--debug", + action="store_true", + help=( + "Enable debug mode with full LiteLLM debug logging. " + "Shows detailed API request/response information." + ), + ) + args = parser.parse_args() # Require either --target or --scope @@ -616,6 +637,20 @@ def main() -> None: if args.validate: return + # Enable debug mode if requested + if getattr(args, "debug", False): + console = Console() + console.print("[bold yellow]Debug mode enabled - showing full LiteLLM debug output[/]\n") + os.environ["STRIX_DEBUG"] = "true" + litellm._turn_on_debug() + logging.getLogger().setLevel(logging.DEBUG) + + # Enable verbose mode via environment for CLI to pick up + if getattr(args, "verbose", False): + os.environ["STRIX_VERBOSE"] = "true" + # Verbose implies we want to see more logging too + logging.getLogger().setLevel(logging.INFO) + check_docker_installed() pull_docker_image() diff --git a/strix/llm/__init__.py b/strix/llm/__init__.py index 6dde52543..33ce1b6df 100644 --- a/strix/llm/__init__.py +++ b/strix/llm/__init__.py @@ -1,3 +1,5 @@ +import os + import litellm from .config import LLMConfig @@ -10,6 +12,9 @@ "LLMRequestFailedError", ] -litellm._logging._disable_debugging() +# Only disable debugging if not in debug mode +# Debug mode is enabled via --debug flag or STRIX_DEBUG env var +if os.getenv("STRIX_DEBUG", "").lower() != "true": + litellm._logging._disable_debugging() litellm.drop_params = True diff --git a/strix/llm/llm.py b/strix/llm/llm.py index 351575476..a28a1d602 100644 --- a/strix/llm/llm.py +++ b/strix/llm/llm.py @@ -1,5 +1,7 @@ +import asyncio import logging import os +import time from dataclasses import dataclass from enum import Enum from fnmatch import fnmatch @@ -25,6 +27,16 @@ logger = logging.getLogger(__name__) + +def _get_tracer() -> Any: + """Get global tracer if available.""" + try: + from strix.telemetry.tracer import get_global_tracer + + return get_global_tracer() + except ImportError: + return None + api_key = os.getenv("LLM_API_KEY") if api_key: litellm.api_key = api_key @@ -70,6 +82,20 @@ def __init__(self, message: str, details: str | None = None): "claude-haiku-4-5*", ] +# Retry configuration for transient errors +MAX_RETRIES = 3 +INITIAL_RETRY_DELAY = 2.0 # seconds +MAX_RETRY_DELAY = 16.0 # seconds + +# Errors that are transient and should be retried +RETRYABLE_ERRORS = ( + litellm.RateLimitError, + litellm.Timeout, + litellm.ServiceUnavailableError, + litellm.InternalServerError, + litellm.APIConnectionError, +) + def normalize_model_name(model: str) -> str: raw = (model or "").strip().lower() @@ -289,86 +315,173 @@ async def generate( # noqa: PLR0912, PLR0915 if identity_message: messages.append(identity_message) + # Compress history (creates new list, doesn't mutate input) compressed_history = list(self.memory_compressor.compress_history(conversation_history)) - conversation_history.clear() - conversation_history.extend(compressed_history) + # Update caller's history in-place with compressed version + # Note: This is intentional - it keeps agent's state.messages compressed + if len(compressed_history) < len(conversation_history): + conversation_history.clear() + conversation_history.extend(compressed_history) + messages.extend(compressed_history) cached_messages = self._prepare_cached_messages(messages) - try: - response = await self._make_request(cached_messages) - self._update_usage_stats(response) - - content = "" - if ( - response.choices - and hasattr(response.choices[0], "message") - and response.choices[0].message - ): - content = getattr(response.choices[0].message, "content", "") or "" - - content = _truncate_to_first_function(content) - - if "" in content: - function_end_index = content.find("") + len("") - content = content[:function_end_index] - - tool_invocations = parse_tool_invocations(content) - - return LLMResponse( - scan_id=scan_id, - step_number=step_number, - role=StepRole.AGENT, - content=content, - tool_invocations=tool_invocations if tool_invocations else None, + # Emit LLM request event + tracer = _get_tracer() + request_id = None + start_time = time.time() + + if tracer and self.agent_id: + request_id = tracer.log_llm_request( + agent_id=self.agent_id, + model=self.config.model_name, ) - except litellm.RateLimitError as e: - raise LLMRequestFailedError("LLM request failed: Rate limit exceeded", str(e)) from e - except litellm.AuthenticationError as e: - raise LLMRequestFailedError("LLM request failed: Invalid API key", str(e)) from e - except litellm.NotFoundError as e: - raise LLMRequestFailedError("LLM request failed: Model not found", str(e)) from e - except litellm.ContextWindowExceededError as e: - raise LLMRequestFailedError("LLM request failed: Context too long", str(e)) from e - except litellm.ContentPolicyViolationError as e: - raise LLMRequestFailedError( - "LLM request failed: Content policy violation", str(e) - ) from e - except litellm.ServiceUnavailableError as e: - raise LLMRequestFailedError("LLM request failed: Service unavailable", str(e)) from e - except litellm.Timeout as e: - raise LLMRequestFailedError("LLM request failed: Request timed out", str(e)) from e - except litellm.UnprocessableEntityError as e: - raise LLMRequestFailedError("LLM request failed: Unprocessable entity", str(e)) from e - except litellm.InternalServerError as e: - raise LLMRequestFailedError("LLM request failed: Internal server error", str(e)) from e - except litellm.APIConnectionError as e: - raise LLMRequestFailedError("LLM request failed: Connection error", str(e)) from e - except litellm.UnsupportedParamsError as e: - raise LLMRequestFailedError("LLM request failed: Unsupported parameters", str(e)) from e - except litellm.BudgetExceededError as e: - raise LLMRequestFailedError("LLM request failed: Budget exceeded", str(e)) from e - except litellm.APIResponseValidationError as e: - raise LLMRequestFailedError( - "LLM request failed: Response validation error", str(e) - ) from e - except litellm.JSONSchemaValidationError as e: + # Retry loop with exponential backoff for transient errors + last_error: Exception | None = None + retry_delay = INITIAL_RETRY_DELAY + + for attempt in range(MAX_RETRIES + 1): + try: + response = await self._make_request(cached_messages) + self._update_usage_stats(response) + + # Emit LLM response event + duration_ms = (time.time() - start_time) * 1000 + if tracer and self.agent_id and request_id: + tracer.log_llm_response( + agent_id=self.agent_id, + request_id=request_id, + input_tokens=self._last_request_stats.input_tokens, + output_tokens=self._last_request_stats.output_tokens, + duration_ms=duration_ms, + cost=self._last_request_stats.cost, + cached_tokens=self._last_request_stats.cached_tokens, + ) + + content = "" + if ( + response.choices + and hasattr(response.choices[0], "message") + and response.choices[0].message + ): + content = getattr(response.choices[0].message, "content", "") or "" + + content = _truncate_to_first_function(content) + + if "" in content: + function_end_index = content.find("") + len("") + content = content[:function_end_index] + + tool_invocations = parse_tool_invocations(content) + + return LLMResponse( + scan_id=scan_id, + step_number=step_number, + role=StepRole.AGENT, + content=content, + tool_invocations=tool_invocations if tool_invocations else None, + ) + + except RETRYABLE_ERRORS as e: + last_error = e + error_type = type(e).__name__ + + if attempt < MAX_RETRIES: + # Log retry attempt + logger.warning( + f"LLM request failed ({error_type}), retrying in {retry_delay}s " + f"(attempt {attempt + 1}/{MAX_RETRIES}): {e}" + ) + await asyncio.sleep(retry_delay) + retry_delay = min(retry_delay * 2, MAX_RETRY_DELAY) + continue + + # All retries exhausted + self._emit_llm_error( + tracer, request_id, start_time, + f"{error_type} (after {MAX_RETRIES} retries)", + error_type, str(e) + ) + raise LLMRequestFailedError( + f"LLM request failed: {error_type} (after {MAX_RETRIES} retries)", + str(e) + ) from e + + except ( + litellm.AuthenticationError, + litellm.NotFoundError, + litellm.ContextWindowExceededError, + litellm.ContentPolicyViolationError, + litellm.BudgetExceededError, + litellm.UnsupportedParamsError, + litellm.InvalidRequestError, + litellm.BadRequestError, + ) as e: + # Non-retryable errors - fail immediately + error_type = type(e).__name__ + error_msg = self._get_error_message(e) + self._emit_llm_error(tracer, request_id, start_time, error_msg, error_type, str(e)) + raise LLMRequestFailedError(f"LLM request failed: {error_msg}", str(e)) from e + + except Exception as e: + # Unknown errors - fail immediately + error_type = type(e).__name__ + self._emit_llm_error(tracer, request_id, start_time, str(e), error_type, str(e)) + raise LLMRequestFailedError(f"LLM request failed: {error_type}", str(e)) from e + + # Should not reach here, but handle edge case + if last_error: raise LLMRequestFailedError( - "LLM request failed: JSON schema validation error", str(e) - ) from e - except litellm.InvalidRequestError as e: - raise LLMRequestFailedError("LLM request failed: Invalid request", str(e)) from e - except litellm.BadRequestError as e: - raise LLMRequestFailedError("LLM request failed: Bad request", str(e)) from e - except litellm.APIError as e: - raise LLMRequestFailedError("LLM request failed: API error", str(e)) from e - except litellm.OpenAIError as e: - raise LLMRequestFailedError("LLM request failed: OpenAI error", str(e)) from e - except Exception as e: - raise LLMRequestFailedError(f"LLM request failed: {type(e).__name__}", str(e)) from e + f"LLM request failed after {MAX_RETRIES} retries", + str(last_error) + ) from last_error + raise LLMRequestFailedError("LLM request failed: Unknown error") + + def _emit_llm_error( + self, tracer: Any, request_id: str | None, start_time: float, error: str, + error_type: str | None = None, details: str | None = None + ) -> None: + """Emit LLM error event to tracer with full details.""" + if tracer and self.agent_id and request_id: + duration_ms = (time.time() - start_time) * 1000 + # Build detailed error message + error_parts = [error] + if error_type: + error_parts.insert(0, f"[{error_type}]") + if details: + # Truncate very long details but keep useful info + if len(details) > 500: + details = details[:500] + "..." + error_parts.append(f"\nDetails: {details}") + + full_error = " ".join(error_parts) if error_type else error + if details: + full_error = f"{full_error}\nDetails: {details}" + + tracer.log_llm_error( + agent_id=self.agent_id, + request_id=request_id, + error=full_error, + duration_ms=duration_ms, + ) + + def _get_error_message(self, error: Exception) -> str: + """Get human-readable error message for common LiteLLM errors.""" + error_messages = { + "AuthenticationError": "Invalid API key", + "NotFoundError": "Model not found", + "ContextWindowExceededError": "Context too long", + "ContentPolicyViolationError": "Content policy violation", + "BudgetExceededError": "Budget exceeded", + "UnsupportedParamsError": "Unsupported parameters", + "InvalidRequestError": "Invalid request", + "BadRequestError": "Bad request", + } + error_type = type(error).__name__ + return error_messages.get(error_type, error_type) @property def usage_stats(self) -> dict[str, dict[str, int | float]]: diff --git a/strix/llm/request_queue.py b/strix/llm/request_queue.py index cd99bcfbb..770639a67 100644 --- a/strix/llm/request_queue.py +++ b/strix/llm/request_queue.py @@ -1,6 +1,5 @@ import asyncio import logging -import threading import time from typing import Any @@ -29,16 +28,25 @@ 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._semaphore: asyncio.Semaphore | None = None self._last_request_time = 0.0 - self._lock = threading.Lock() + self._lock: asyncio.Lock | None = None - 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) + def _get_semaphore(self) -> asyncio.Semaphore: + """Lazily create semaphore in the current event loop.""" + if self._semaphore is None: + self._semaphore = asyncio.Semaphore(self.max_concurrent) + return self._semaphore + + def _get_lock(self) -> asyncio.Lock: + """Lazily create lock in the current event loop.""" + if self._lock is None: + self._lock = asyncio.Lock() + return self._lock - with self._lock: + async def make_request(self, completion_args: dict[str, Any]) -> ModelResponse: + async with self._get_semaphore(): + async with self._get_lock(): now = time.time() time_since_last = now - self._last_request_time sleep_needed = max(0, self.delay_between_requests - time_since_last) @@ -48,8 +56,6 @@ async def make_request(self, completion_args: dict[str, Any]) -> ModelResponse: await asyncio.sleep(sleep_needed) return await self._reliable_request(completion_args) - finally: - self._semaphore.release() @retry( # type: ignore[misc] stop=stop_after_attempt(5), diff --git a/strix/runtime/docker_runtime.py b/strix/runtime/docker_runtime.py index 63cb7a033..4bd516dd1 100644 --- a/strix/runtime/docker_runtime.py +++ b/strix/runtime/docker_runtime.py @@ -77,25 +77,31 @@ def _validate_image(image: docker.models.images.Image) -> None: logger.debug(f"Image {image_name} verified as available") return + def _cleanup_failed_container(self, container_name: str) -> None: + """Clean up a partially created or failed container.""" + try: + container = self.client.containers.get(container_name) + logger.warning(f"Cleaning up failed container {container_name}") + with contextlib.suppress(Exception): + container.stop(timeout=5) + container.remove(force=True) + except NotFound: + pass + except DockerException as e: + logger.warning(f"Error during container cleanup: {e}") + def _create_container_with_retry(self, scan_id: str, max_retries: int = 3) -> Container: last_exception = None container_name = f"strix-scan-{scan_id}" + container: Container | None = None for attempt in range(max_retries): try: self._verify_image_available(STRIX_IMAGE) - try: - existing_container = self.client.containers.get(container_name) - logger.warning(f"Container {container_name} already exists, removing it") - with contextlib.suppress(Exception): - existing_container.stop(timeout=5) - existing_container.remove(force=True) - time.sleep(1) - except NotFound: - pass - except DockerException as e: - logger.warning(f"Error checking/removing existing container: {e}") + # Clean up any existing container before creation + self._cleanup_failed_container(container_name) + time.sleep(1) caido_port = self._find_available_port() tool_server_port = self._find_available_port() @@ -131,17 +137,20 @@ def _create_container_with_retry(self, scan_id: str, max_retries: int = 3) -> Co self._initialize_container( container, caido_port, tool_server_port, tool_server_token ) - except DockerException as e: + except (DockerException, RuntimeError, OSError) as e: last_exception = e - if attempt == max_retries - 1: - logger.exception(f"Failed to create container after {max_retries} attempts") - break - - logger.warning(f"Container creation attempt {attempt + 1}/{max_retries} failed") + logger.warning(f"Container creation attempt {attempt + 1}/{max_retries} failed: {e}") + # Clean up failed container to avoid orphans + self._cleanup_failed_container(container_name) + self._scan_container = None self._tool_server_port = None self._tool_server_token = None + if attempt == max_retries - 1: + logger.exception(f"Failed to create container after {max_retries} attempts") + break + sleep_time = (2**attempt) + (0.1 * attempt) time.sleep(sleep_time) else: diff --git a/strix/runtime/tool_server.py b/strix/runtime/tool_server.py index 6461f8c79..aeb61719c 100644 --- a/strix/runtime/tool_server.py +++ b/strix/runtime/tool_server.py @@ -65,19 +65,37 @@ class ToolExecutionResponse(BaseModel): error: str | None = None +QUEUE_TIMEOUT = 120 # 2 minutes timeout for queue operations + + def agent_worker(_agent_id: str, request_queue: Queue[Any], response_queue: Queue[Any]) -> None: - null_handler = logging.NullHandler() + import os + from pathlib import Path + from queue import Empty + + # Configure file-based logging for worker process (instead of suppressing) + log_dir = Path("/tmp/strix_workers") + log_dir.mkdir(exist_ok=True) + log_file = log_dir / f"worker_{os.getpid()}.log" + + file_handler = logging.FileHandler(log_file) + file_handler.setLevel(logging.WARNING) + file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")) root_logger = logging.getLogger() - root_logger.handlers = [null_handler] - root_logger.setLevel(logging.CRITICAL) + root_logger.handlers = [file_handler] + root_logger.setLevel(logging.WARNING) from strix.tools.argument_parser import ArgumentConversionError, convert_arguments from strix.tools.registry import get_tool_by_name while True: try: - request = request_queue.get() + # Use timeout to prevent infinite blocking + try: + request = request_queue.get(timeout=QUEUE_TIMEOUT) + except Empty: + continue # Keep worker alive, just no request received if request is None: break diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index 15a4b4238..f39d8a8a0 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -1,5 +1,8 @@ +import json import logging +from dataclasses import dataclass, field from datetime import UTC, datetime +from enum import Enum from pathlib import Path from typing import TYPE_CHECKING, Any, Optional from uuid import uuid4 @@ -11,6 +14,55 @@ logger = logging.getLogger(__name__) + +class EventType(Enum): + """Types of events tracked by the tracer.""" + + # Agent lifecycle events + AGENT_CREATED = "agent_created" + AGENT_STATUS_CHANGED = "agent_status_changed" + AGENT_ITERATION = "agent_iteration" + AGENT_STATE_TRANSITION = "agent_state_transition" + + # Tool events + TOOL_START = "tool_start" + TOOL_COMPLETE = "tool_complete" + TOOL_ERROR = "tool_error" + + # LLM events + LLM_REQUEST = "llm_request" + LLM_RESPONSE = "llm_response" + LLM_ERROR = "llm_error" + + # Inter-agent communication + AGENT_MESSAGE_SENT = "agent_message_sent" + AGENT_MESSAGE_RECEIVED = "agent_message_received" + + # Vulnerability events + VULNERABILITY_FOUND = "vulnerability_found" + + # Scan events + SCAN_START = "scan_start" + SCAN_COMPLETE = "scan_complete" + + +@dataclass +class TracerEvent: + """A single event in the tracer event stream.""" + + event_type: EventType + timestamp: str + agent_id: str | None = None + data: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + return { + "event_type": self.event_type.value, + "timestamp": self.timestamp, + "agent_id": self.agent_id, + "data": self.data, + } + _global_tracer: Optional["Tracer"] = None @@ -50,8 +102,15 @@ def __init__(self, run_name: str | None = None): self._run_dir: Path | None = None self._next_execution_id = 1 self._next_message_id = 1 + self._next_event_id = 1 + # Unified event stream for real-time visibility + self.events: list[TracerEvent] = [] + self._event_cursor: int = 0 # For consumers to track their position + + # Callbacks for real-time notifications self.vulnerability_found_callback: Callable[[str, str, str, str], None] | None = None + self.event_callback: Callable[[TracerEvent], None] | None = None def set_run_name(self, run_name: str) -> None: self.run_name = run_name @@ -68,6 +127,19 @@ def get_run_dir(self) -> Path: return self._run_dir + def _save_metadata(self) -> None: + """Save metadata.json to run directory (called on scan start and updates).""" + try: + metadata_file = self.get_run_dir() / "metadata.json" + # Use atomic write pattern + temp_file = metadata_file.with_suffix(".json.tmp") + with temp_file.open("w", encoding="utf-8") as f: + json.dump(self.run_metadata, f, indent=2, ensure_ascii=False, default=str) + temp_file.replace(metadata_file) + logger.debug(f"Metadata saved to {metadata_file}") + except (OSError, IOError) as e: + logger.warning(f"Failed to save metadata: {e}") + def add_vulnerability_report( self, title: str, @@ -198,11 +270,50 @@ def set_scan_config(self, config: dict[str, Any]) -> None: } ) + # Immediately persist metadata on scan start + self._save_metadata() + + # Emit scan start event + event = TracerEvent( + event_type=EventType.SCAN_START, + timestamp=datetime.now(UTC).isoformat(), + agent_id=None, + data={ + "run_id": self.run_id, + "run_name": self.run_name, + "targets": config.get("targets", []), + "max_iterations": config.get("max_iterations", 200), + }, + ) + self._emit_event(event) + def save_run_data(self) -> None: try: run_dir = self.get_run_dir() self.end_time = datetime.now(UTC).isoformat() + # Update and persist final metadata + self.run_metadata["end_time"] = self.end_time + self.run_metadata["status"] = "completed" + self.run_metadata["vulnerability_count"] = len(self.vulnerability_reports) + self.run_metadata["duration_seconds"] = self._calculate_duration() + self._save_metadata() + + # Emit scan complete event + event = TracerEvent( + event_type=EventType.SCAN_COMPLETE, + timestamp=self.end_time, + agent_id=None, + data={ + "run_id": self.run_id, + "duration_seconds": self._calculate_duration(), + "vulnerability_count": len(self.vulnerability_reports), + "agent_count": len(self.agents), + "tool_execution_count": len(self.tool_executions), + }, + ) + self._emit_event(event) + if self.final_scan_result: penetration_test_report_file = run_dir / "penetration_test_report.md" with penetration_test_report_file.open("w", encoding="utf-8") as f: @@ -321,3 +432,228 @@ def get_total_llm_stats(self) -> dict[str, Any]: def cleanup(self) -> None: self.save_run_data() + + # ========================================================================= + # New Event Stream Methods for Real-Time Visibility + # ========================================================================= + + def _emit_event(self, event: TracerEvent) -> None: + """Add event to stream, persist to JSONL, and notify callback if registered.""" + self.events.append(event) + + # Append to JSONL file (crash-safe, append-only) + try: + events_file = self.get_run_dir() / "events.jsonl" + with events_file.open("a", encoding="utf-8") as f: + f.write(json.dumps(event.to_dict(), default=str) + "\n") + except (OSError, IOError) as e: + logger.warning(f"Failed to append event to JSONL: {e}") + + if self.event_callback: + try: + self.event_callback(event) + except Exception: # noqa: BLE001 + logger.exception("Error in event callback") + + def log_agent_iteration( + self, + agent_id: str, + iteration: int, + max_iterations: int, + ) -> None: + """Log an agent iteration event.""" + event = TracerEvent( + event_type=EventType.AGENT_ITERATION, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data={ + "iteration": iteration, + "max_iterations": max_iterations, + "progress_pct": round((iteration / max_iterations) * 100, 1), + }, + ) + self._emit_event(event) + + def log_agent_state_transition( + self, + agent_id: str, + from_state: str, + to_state: str, + reason: str | None = None, + ) -> None: + """Log an agent state transition event.""" + event = TracerEvent( + event_type=EventType.AGENT_STATE_TRANSITION, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data={ + "from_state": from_state, + "to_state": to_state, + "reason": reason, + }, + ) + self._emit_event(event) + + def log_llm_request( + self, + agent_id: str, + model: str, + prompt_tokens: int | None = None, + request_id: str | None = None, + ) -> str: + """Log an LLM request event. Returns request_id for correlation.""" + if request_id is None: + request_id = f"llm-{self._next_event_id}" + self._next_event_id += 1 + + event = TracerEvent( + event_type=EventType.LLM_REQUEST, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data={ + "request_id": request_id, + "model": model, + "prompt_tokens": prompt_tokens, + }, + ) + self._emit_event(event) + return request_id + + def log_llm_response( + self, + agent_id: str, + request_id: str, + input_tokens: int, + output_tokens: int, + duration_ms: float, + cost: float | None = None, + cached_tokens: int = 0, + ) -> None: + """Log an LLM response event.""" + event = TracerEvent( + event_type=EventType.LLM_RESPONSE, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data={ + "request_id": request_id, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cached_tokens": cached_tokens, + "total_tokens": input_tokens + output_tokens, + "duration_ms": round(duration_ms, 1), + "cost": round(cost, 6) if cost else None, + }, + ) + self._emit_event(event) + + def log_llm_error( + self, + agent_id: str, + request_id: str, + error: str, + duration_ms: float | None = None, + ) -> None: + """Log an LLM error event.""" + event = TracerEvent( + event_type=EventType.LLM_ERROR, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data={ + "request_id": request_id, + "error": error, + "duration_ms": round(duration_ms, 1) if duration_ms else None, + }, + ) + self._emit_event(event) + + def log_agent_message_sent( + self, + from_agent_id: str, + to_agent_id: str, + message: str, + ) -> None: + """Log an inter-agent message sent event.""" + event = TracerEvent( + event_type=EventType.AGENT_MESSAGE_SENT, + timestamp=datetime.now(UTC).isoformat(), + agent_id=from_agent_id, + data={ + "to_agent_id": to_agent_id, + "message_preview": message[:200] + "..." if len(message) > 200 else message, + "message_length": len(message), + }, + ) + self._emit_event(event) + + def log_agent_message_received( + self, + agent_id: str, + from_agent_id: str, + message: str, + ) -> None: + """Log an inter-agent message received event.""" + event = TracerEvent( + event_type=EventType.AGENT_MESSAGE_RECEIVED, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data={ + "from_agent_id": from_agent_id, + "message_preview": message[:200] + "..." if len(message) > 200 else message, + "message_length": len(message), + }, + ) + self._emit_event(event) + + def log_tool_event( + self, + agent_id: str, + tool_name: str, + event_type: EventType, + args: dict[str, Any] | None = None, + result: Any | None = None, + error: str | None = None, + duration_ms: float | None = None, + ) -> None: + """Log a tool execution event (start, complete, or error).""" + data: dict[str, Any] = {"tool_name": tool_name} + + if args is not None: + # Truncate large args for display + args_str = str(args) + if len(args_str) > 500: + data["args_preview"] = args_str[:500] + "..." + else: + data["args"] = args + + if result is not None: + result_str = str(result) + if len(result_str) > 500: + data["result_preview"] = result_str[:500] + "..." + else: + data["result"] = result + + if error is not None: + data["error"] = error + + if duration_ms is not None: + data["duration_ms"] = round(duration_ms, 1) + + event = TracerEvent( + event_type=event_type, + timestamp=datetime.now(UTC).isoformat(), + agent_id=agent_id, + data=data, + ) + self._emit_event(event) + + def get_events_since(self, cursor: int = 0) -> tuple[list[TracerEvent], int]: + """Get all events since the given cursor position. + + Returns (events, new_cursor) tuple. + """ + new_events = self.events[cursor:] + return new_events, len(self.events) + + def get_recent_events(self, count: int = 50) -> list[TracerEvent]: + """Get the most recent N events.""" + return self.events[-count:] if self.events else [] diff --git a/strix/tools/__init__.py b/strix/tools/__init__.py index 01ac4f011..9f39d1983 100644 --- a/strix/tools/__init__.py +++ b/strix/tools/__init__.py @@ -32,6 +32,7 @@ from .file_edit import * # noqa: F403 from .finish import * # noqa: F403 from .notes import * # noqa: F403 + from .progress import * # noqa: F403 from .proxy import * # noqa: F403 from .python import * # noqa: F403 from .reporting import * # noqa: F403 diff --git a/strix/tools/agents_graph/agents_graph_actions.py b/strix/tools/agents_graph/agents_graph_actions.py index f2fb9a8b2..fe19284fa 100644 --- a/strix/tools/agents_graph/agents_graph_actions.py +++ b/strix/tools/agents_graph/agents_graph_actions.py @@ -5,6 +5,9 @@ from strix.tools.registry import register_tool +# Thread lock for all global mutable state - use RLock to allow re-entrant locking +_graph_lock = threading.RLock() + _agent_graph: dict[str, Any] = { "nodes": {}, "edges": [], @@ -18,8 +21,6 @@ _agent_instances: dict[str, Any] = {} -_agent_states: dict[str, Any] = {} - def _run_agent_in_thread( agent: Any, state: Any, inherited_messages: list[dict[str, Any]] @@ -69,9 +70,8 @@ def _run_agent_in_thread( state.add_message("user", task_xml) - _agent_states[state.agent_id] = state - - _agent_graph["nodes"][state.agent_id]["state"] = state.model_dump() + with _graph_lock: + _agent_graph["nodes"][state.agent_id]["state"] = state.model_dump() import asyncio @@ -83,105 +83,108 @@ def _run_agent_in_thread( loop.close() except Exception as e: - _agent_graph["nodes"][state.agent_id]["status"] = "error" - _agent_graph["nodes"][state.agent_id]["finished_at"] = datetime.now(UTC).isoformat() - _agent_graph["nodes"][state.agent_id]["result"] = {"error": str(e)} - _running_agents.pop(state.agent_id, None) - _agent_instances.pop(state.agent_id, None) + with _graph_lock: + _agent_graph["nodes"][state.agent_id]["status"] = "error" + _agent_graph["nodes"][state.agent_id]["finished_at"] = datetime.now(UTC).isoformat() + _agent_graph["nodes"][state.agent_id]["result"] = {"error": str(e)} + _running_agents.pop(state.agent_id, None) + _agent_instances.pop(state.agent_id, None) raise else: - if state.stop_requested: - _agent_graph["nodes"][state.agent_id]["status"] = "stopped" - else: - _agent_graph["nodes"][state.agent_id]["status"] = "completed" - _agent_graph["nodes"][state.agent_id]["finished_at"] = datetime.now(UTC).isoformat() - _agent_graph["nodes"][state.agent_id]["result"] = result - _running_agents.pop(state.agent_id, None) - _agent_instances.pop(state.agent_id, None) + with _graph_lock: + if state.stop_requested: + _agent_graph["nodes"][state.agent_id]["status"] = "stopped" + else: + _agent_graph["nodes"][state.agent_id]["status"] = "completed" + _agent_graph["nodes"][state.agent_id]["finished_at"] = datetime.now(UTC).isoformat() + _agent_graph["nodes"][state.agent_id]["result"] = result + _running_agents.pop(state.agent_id, None) + _agent_instances.pop(state.agent_id, None) return {"result": result} @register_tool(sandbox_execution=False) def view_agent_graph(agent_state: Any) -> dict[str, Any]: - try: - structure_lines = ["=== AGENT GRAPH STRUCTURE ==="] - - def _build_tree(agent_id: str, depth: int = 0) -> None: - node = _agent_graph["nodes"][agent_id] - indent = " " * depth - - you_indicator = " ← This is you" if agent_id == agent_state.agent_id else "" - - structure_lines.append(f"{indent}* {node['name']} ({agent_id}){you_indicator}") - structure_lines.append(f"{indent} Task: {node['task']}") - structure_lines.append(f"{indent} Status: {node['status']}") - - children = [ - edge["to"] - for edge in _agent_graph["edges"] - if edge["from"] == agent_id and edge["type"] == "delegation" - ] - - if children: - structure_lines.append(f"{indent} Children:") - for child_id in children: - _build_tree(child_id, depth + 2) - - root_agent_id = _root_agent_id - if not root_agent_id and _agent_graph["nodes"]: - for agent_id, node in _agent_graph["nodes"].items(): - if node.get("parent_id") is None: - root_agent_id = agent_id - break - if not root_agent_id: - root_agent_id = next(iter(_agent_graph["nodes"].keys())) - - if root_agent_id and root_agent_id in _agent_graph["nodes"]: - _build_tree(root_agent_id) + with _graph_lock: + try: + structure_lines = ["=== AGENT GRAPH STRUCTURE ==="] + + def _build_tree(agent_id: str, depth: int = 0) -> None: + node = _agent_graph["nodes"][agent_id] + indent = " " * depth + + you_indicator = " ← This is you" if agent_id == agent_state.agent_id else "" + + structure_lines.append(f"{indent}* {node['name']} ({agent_id}){you_indicator}") + structure_lines.append(f"{indent} Task: {node['task']}") + structure_lines.append(f"{indent} Status: {node['status']}") + + children = [ + edge["to"] + for edge in _agent_graph["edges"] + if edge["from"] == agent_id and edge["type"] == "delegation" + ] + + if children: + structure_lines.append(f"{indent} Children:") + for child_id in children: + _build_tree(child_id, depth + 2) + + root_agent_id = _root_agent_id + if not root_agent_id and _agent_graph["nodes"]: + for agent_id, node in _agent_graph["nodes"].items(): + if node.get("parent_id") is None: + root_agent_id = agent_id + break + if not root_agent_id: + root_agent_id = next(iter(_agent_graph["nodes"].keys())) + + if root_agent_id and root_agent_id in _agent_graph["nodes"]: + _build_tree(root_agent_id) + else: + structure_lines.append("No agents in the graph yet") + + graph_structure = "\n".join(structure_lines) + + total_nodes = len(_agent_graph["nodes"]) + running_count = sum( + 1 for node in _agent_graph["nodes"].values() if node["status"] == "running" + ) + waiting_count = sum( + 1 for node in _agent_graph["nodes"].values() if node["status"] == "waiting" + ) + stopping_count = sum( + 1 for node in _agent_graph["nodes"].values() if node["status"] == "stopping" + ) + completed_count = sum( + 1 for node in _agent_graph["nodes"].values() if node["status"] == "completed" + ) + stopped_count = sum( + 1 for node in _agent_graph["nodes"].values() if node["status"] == "stopped" + ) + failed_count = sum( + 1 for node in _agent_graph["nodes"].values() if node["status"] in ["failed", "error"] + ) + + except Exception as e: # noqa: BLE001 + return { + "error": f"Failed to view agent graph: {e}", + "graph_structure": "Error retrieving graph structure", + } else: - structure_lines.append("No agents in the graph yet") - - graph_structure = "\n".join(structure_lines) - - total_nodes = len(_agent_graph["nodes"]) - running_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "running" - ) - waiting_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "waiting" - ) - stopping_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "stopping" - ) - completed_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "completed" - ) - stopped_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] == "stopped" - ) - failed_count = sum( - 1 for node in _agent_graph["nodes"].values() if node["status"] in ["failed", "error"] - ) - - except Exception as e: # noqa: BLE001 - return { - "error": f"Failed to view agent graph: {e}", - "graph_structure": "Error retrieving graph structure", - } - else: - return { - "graph_structure": graph_structure, - "summary": { - "total_agents": total_nodes, - "running": running_count, - "waiting": waiting_count, - "stopping": stopping_count, - "completed": completed_count, - "stopped": stopped_count, - "failed": failed_count, - }, - } + return { + "graph_structure": graph_structure, + "summary": { + "total_agents": total_nodes, + "running": running_count, + "waiting": waiting_count, + "stopping": stopping_count, + "completed": completed_count, + "stopped": stopped_count, + "failed": failed_count, + }, + } @register_tool(sandbox_execution=False) @@ -267,16 +270,18 @@ def create_agent( if inherit_context: inherited_messages = agent_state.get_conversation_history() - _agent_instances[state.agent_id] = agent + # Atomic registration of agent instance and thread + with _graph_lock: + _agent_instances[state.agent_id] = agent - thread = threading.Thread( - target=_run_agent_in_thread, - args=(agent, state, inherited_messages), - daemon=True, - name=f"Agent-{name}-{state.agent_id}", - ) - thread.start() - _running_agents[state.agent_id] = thread + thread = threading.Thread( + target=_run_agent_in_thread, + args=(agent, state, inherited_messages), + daemon=True, + name=f"Agent-{name}-{state.agent_id}", + ) + thread.start() + _running_agents[state.agent_id] = thread except Exception as e: # noqa: BLE001 return {"success": False, "error": f"Failed to create agent: {e}", "agent_id": None} @@ -303,51 +308,67 @@ def send_message_to_agent( priority: Literal["low", "normal", "high", "urgent"] = "normal", ) -> dict[str, Any]: try: - if target_agent_id not in _agent_graph["nodes"]: - return { - "success": False, - "error": f"Target agent '{target_agent_id}' not found in graph", - "message_id": None, - } - - sender_id = agent_state.agent_id - - from uuid import uuid4 - - message_id = f"msg_{uuid4().hex[:8]}" - message_data = { - "id": message_id, - "from": sender_id, - "to": target_agent_id, - "content": message, - "message_type": message_type, - "priority": priority, - "timestamp": datetime.now(UTC).isoformat(), - "delivered": False, - "read": False, - } + with _graph_lock: + if target_agent_id not in _agent_graph["nodes"]: + return { + "success": False, + "error": f"Target agent '{target_agent_id}' not found in graph", + "message_id": None, + } - if target_agent_id not in _agent_messages: - _agent_messages[target_agent_id] = [] + sender_id = agent_state.agent_id - _agent_messages[target_agent_id].append(message_data) + from uuid import uuid4 - _agent_graph["edges"].append( - { + message_id = f"msg_{uuid4().hex[:8]}" + message_data = { + "id": message_id, "from": sender_id, "to": target_agent_id, - "type": "message", - "message_id": message_id, + "content": message, "message_type": message_type, "priority": priority, - "created_at": datetime.now(UTC).isoformat(), + "timestamp": datetime.now(UTC).isoformat(), + "delivered": False, + "read": False, } - ) - message_data["delivered"] = True + if target_agent_id not in _agent_messages: + _agent_messages[target_agent_id] = [] + + _agent_messages[target_agent_id].append(message_data) - target_name = _agent_graph["nodes"][target_agent_id]["name"] - sender_name = _agent_graph["nodes"][sender_id]["name"] + _agent_graph["edges"].append( + { + "from": sender_id, + "to": target_agent_id, + "type": "message", + "message_id": message_id, + "message_type": message_type, + "priority": priority, + "created_at": datetime.now(UTC).isoformat(), + } + ) + + message_data["delivered"] = True + + target_name = _agent_graph["nodes"][target_agent_id]["name"] + sender_name = _agent_graph["nodes"][sender_id]["name"] + target_status = _agent_graph["nodes"][target_agent_id]["status"] + + # Emit inter-agent message event (outside lock) + try: + from strix.telemetry.tracer import get_global_tracer + + tracer = get_global_tracer() + if tracer: + tracer.log_agent_message_sent( + from_agent_id=sender_id, + to_agent_id=target_agent_id, + message=message, + ) + except (ImportError, AttributeError): + pass return { "success": True, @@ -357,7 +378,7 @@ def send_message_to_agent( "target_agent": { "id": target_agent_id, "name": target_name, - "status": _agent_graph["nodes"][target_agent_id]["status"], + "status": target_status, }, } @@ -453,6 +474,20 @@ def agent_finish( } ) + # Emit completion report message event + try: + from strix.telemetry.tracer import get_global_tracer + + tracer = get_global_tracer() + if tracer: + tracer.log_agent_message_sent( + from_agent_id=agent_id, + to_agent_id=parent_id, + message=f"[Completion Report] {result_summary}", + ) + except (ImportError, AttributeError): + pass + parent_notified = True _running_agents.pop(agent_id, None) @@ -498,10 +533,7 @@ def stop_agent(agent_id: str) -> dict[str, Any]: "previous_status": agent_node["status"], } - if agent_id in _agent_states: - agent_state = _agent_states[agent_id] - agent_state.request_stop() - + # Use _agent_instances to access agent state (single source of truth) if agent_id in _agent_instances: agent_instance = _agent_instances[agent_id] if hasattr(agent_instance, "state"): diff --git a/strix/tools/executor.py b/strix/tools/executor.py index 1eefb255c..15c656c32 100644 --- a/strix/tools/executor.py +++ b/strix/tools/executor.py @@ -1,10 +1,16 @@ import inspect import os +import time +import traceback from typing import Any import httpx +# Sandbox tool execution timeout (2 minutes) +SANDBOX_TIMEOUT = httpx.Timeout(120.0, connect=10.0) + + if os.getenv("STRIX_SANDBOX_MODE", "false").lower() == "false": from strix.runtime import get_runtime @@ -63,16 +69,20 @@ async def _execute_tool_in_sandbox(tool_name: str, agent_state: Any, **kwargs: A "Content-Type": "application/json", } - async with httpx.AsyncClient(trust_env=False) as client: + async with httpx.AsyncClient(trust_env=False, timeout=SANDBOX_TIMEOUT) as client: try: response = await client.post( - request_url, json=request_data, headers=headers, timeout=None + request_url, json=request_data, headers=headers ) response.raise_for_status() response_data = response.json() if response_data.get("error"): raise RuntimeError(f"Sandbox execution error: {response_data['error']}") return response_data.get("result") + except httpx.TimeoutException as e: + raise RuntimeError( + f"Tool execution timed out after {SANDBOX_TIMEOUT.read}s: {tool_name}" + ) from e except httpx.HTTPStatusError as e: if e.response.status_code == 401: raise RuntimeError("Authentication failed: Invalid or missing sandbox token") from e @@ -219,12 +229,16 @@ async def _execute_single_tool( args = tool_inv.get("args", {}) execution_id = None should_agent_finish = False + start_time = time.time() if tracer: execution_id = tracer.log_tool_execution_start(agent_id, tool_name, args) + # Emit TOOL_START event for real-time visibility + _emit_tool_start_event(tracer, agent_id, tool_name, args) try: result = await execute_tool_invocation(tool_inv, agent_state) + duration_ms = (time.time() - start_time) * 1000 is_error, error_payload = _check_error_result(result) @@ -240,16 +254,98 @@ async def _execute_single_tool( _update_tracer_with_result(tracer, execution_id, is_error, result, error_payload) + # Emit completion or error event for real-time visibility + if tracer: + if is_error: + _emit_tool_error_event(tracer, agent_id, tool_name, error_payload, duration_ms) + else: + _emit_tool_complete_event(tracer, agent_id, tool_name, result, duration_ms) + except (ConnectionError, RuntimeError, ValueError, TypeError, OSError) as e: + duration_ms = (time.time() - start_time) * 1000 error_msg = str(e) + error_type = type(e).__name__ + stack_trace = traceback.format_exc() + if tracer and execution_id: tracer.update_tool_execution(execution_id, "error", error_msg) + + # Emit detailed error event for real-time visibility + if tracer: + _emit_tool_error_event( + tracer, agent_id, tool_name, + {"error": error_msg, "error_type": error_type, "stack_trace": stack_trace}, + duration_ms + ) raise observation_xml, images = _format_tool_result(tool_name, result) return observation_xml, images, should_agent_finish +def _emit_tool_start_event( + tracer: Any, agent_id: str, tool_name: str, args: dict[str, Any] +) -> None: + """Emit TOOL_START event to the tracer event stream.""" + try: + from strix.telemetry.tracer import EventType + + tracer.log_tool_event( + agent_id=agent_id, + tool_name=tool_name, + event_type=EventType.TOOL_START, + args=args, + ) + except (ImportError, AttributeError): + pass + + +def _emit_tool_complete_event( + tracer: Any, agent_id: str, tool_name: str, result: Any, duration_ms: float +) -> None: + """Emit TOOL_COMPLETE event to the tracer event stream.""" + try: + from strix.telemetry.tracer import EventType + + tracer.log_tool_event( + agent_id=agent_id, + tool_name=tool_name, + event_type=EventType.TOOL_COMPLETE, + result=result, + duration_ms=duration_ms, + ) + except (ImportError, AttributeError): + pass + + +def _emit_tool_error_event( + tracer: Any, agent_id: str, tool_name: str, error: Any, duration_ms: float +) -> None: + """Emit TOOL_ERROR event to the tracer event stream.""" + try: + from strix.telemetry.tracer import EventType + + # Extract detailed error info + if isinstance(error, dict): + error_str = error.get("error", str(error)) + error_type = error.get("error_type", "UnknownError") + stack_trace = error.get("stack_trace") + else: + error_str = str(error) + error_type = "Error" + stack_trace = None + + tracer.log_tool_event( + agent_id=agent_id, + tool_name=tool_name, + event_type=EventType.TOOL_ERROR, + error=f"[{error_type}] {error_str}", + duration_ms=duration_ms, + ) + except (ImportError, AttributeError): + pass + + def _get_tracer_and_agent_id(agent_state: Any | None) -> tuple[Any | None, str]: try: from strix.telemetry.tracer import get_global_tracer diff --git a/strix/tools/notes/notes_actions.py b/strix/tools/notes/notes_actions.py index 0f91ecd4b..454ccdedb 100644 --- a/strix/tools/notes/notes_actions.py +++ b/strix/tools/notes/notes_actions.py @@ -1,11 +1,73 @@ +import json +import logging import uuid from datetime import UTC, datetime +from pathlib import Path from typing import Any from strix.tools.registry import register_tool +logger = logging.getLogger(__name__) + _notes_storage: dict[str, dict[str, Any]] = {} +_notes_file_path: Path | None = None + + +def _get_notes_file() -> Path | None: + """Get the path to the notes.json file in the run directory.""" + global _notes_file_path # noqa: PLW0603 + if _notes_file_path is not None: + return _notes_file_path + + try: + from strix.telemetry.tracer import get_global_tracer + + tracer = get_global_tracer() + if tracer: + run_dir = tracer.get_run_dir() + _notes_file_path = run_dir / "notes.json" + return _notes_file_path + except (ImportError, AttributeError): + pass + return None + + +def _load_notes_from_disk() -> None: + """Load notes from disk if file exists.""" + global _notes_storage # noqa: PLW0603 + notes_file = _get_notes_file() + if notes_file and notes_file.exists(): + try: + with notes_file.open("r", encoding="utf-8") as f: + _notes_storage = json.load(f) + logger.info(f"Loaded {len(_notes_storage)} notes from {notes_file}") + except (json.JSONDecodeError, OSError) as e: + logger.warning(f"Failed to load notes from disk: {e}") + + +def _save_notes_to_disk() -> bool: + """Save all notes to disk (crash-proof).""" + notes_file = _get_notes_file() + if not notes_file: + return False + + try: + # Ensure parent directory exists + notes_file.parent.mkdir(parents=True, exist_ok=True) + + # Write atomically using temp file + temp_file = notes_file.with_suffix(".json.tmp") + with temp_file.open("w", encoding="utf-8") as f: + json.dump(_notes_storage, f, indent=2, ensure_ascii=False) + + # Atomic rename + temp_file.replace(notes_file) + return True + + except (OSError, IOError) as e: + logger.error(f"Failed to save notes to disk: {e}") + return False def _filter_notes( @@ -52,6 +114,10 @@ def create_note( priority: str = "normal", ) -> dict[str, Any]: try: + # Load existing notes from disk on first access + if not _notes_storage: + _load_notes_from_disk() + if not title or not title.strip(): return {"success": False, "error": "Title cannot be empty", "note_id": None} @@ -89,6 +155,9 @@ def create_note( _notes_storage[note_id] = note + # Immediately persist to disk + persisted = _save_notes_to_disk() + except (ValueError, TypeError) as e: return {"success": False, "error": f"Failed to create note: {e}", "note_id": None} else: @@ -96,6 +165,7 @@ def create_note( "success": True, "note_id": note_id, "message": f"Note '{title}' created successfully", + "persisted_to_disk": persisted, } @@ -107,6 +177,10 @@ def list_notes( search: str | None = None, ) -> dict[str, Any]: try: + # Load existing notes from disk on first access + if not _notes_storage: + _load_notes_from_disk() + filtered_notes = _filter_notes( category=category, tags=tags, priority=priority, search_query=search ) @@ -135,6 +209,10 @@ def update_note( priority: str | None = None, ) -> dict[str, Any]: try: + # Load existing notes from disk on first access + if not _notes_storage: + _load_notes_from_disk() + if note_id not in _notes_storage: return {"success": False, "error": f"Note with ID '{note_id}' not found"} @@ -164,9 +242,13 @@ def update_note( note["updated_at"] = datetime.now(UTC).isoformat() + # Immediately persist to disk + persisted = _save_notes_to_disk() + return { "success": True, "message": f"Note '{note['title']}' updated successfully", + "persisted_to_disk": persisted, } except (ValueError, TypeError) as e: @@ -176,16 +258,24 @@ def update_note( @register_tool def delete_note(note_id: str) -> dict[str, Any]: try: + # Load existing notes from disk on first access + if not _notes_storage: + _load_notes_from_disk() + if note_id not in _notes_storage: return {"success": False, "error": f"Note with ID '{note_id}' not found"} note_title = _notes_storage[note_id]["title"] del _notes_storage[note_id] + # Immediately persist to disk + persisted = _save_notes_to_disk() + except (ValueError, TypeError) as e: return {"success": False, "error": f"Failed to delete note: {e}"} else: return { "success": True, "message": f"Note '{note_title}' deleted successfully", + "persisted_to_disk": persisted, } diff --git a/strix/tools/progress/__init__.py b/strix/tools/progress/__init__.py new file mode 100644 index 000000000..74269c3d6 --- /dev/null +++ b/strix/tools/progress/__init__.py @@ -0,0 +1,5 @@ +"""Progress tracking tools for agents to persist and retrieve scan progress.""" + +from .progress_actions import list_progress, load_progress, save_progress + +__all__ = ["save_progress", "load_progress", "list_progress"] diff --git a/strix/tools/progress/progress_actions.py b/strix/tools/progress/progress_actions.py new file mode 100644 index 000000000..69f61af34 --- /dev/null +++ b/strix/tools/progress/progress_actions.py @@ -0,0 +1,243 @@ +"""Progress tracking tools for agents to offload context and track scan progress. + +These tools allow agents to: +- Save structured data (findings, scanned endpoints, etc.) to disk +- Load previously saved progress data +- List all available progress keys +- Persist state across crashes/restarts +""" + +import json +import logging +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from strix.tools.registry import register_tool + + +logger = logging.getLogger(__name__) + +_progress_cache: dict[str, Any] = {} +_progress_file_path: Path | None = None + + +def _get_progress_file() -> Path | None: + """Get the path to the progress.json file in the run directory.""" + global _progress_file_path # noqa: PLW0603 + if _progress_file_path is not None: + return _progress_file_path + + try: + from strix.telemetry.tracer import get_global_tracer + + tracer = get_global_tracer() + if tracer: + run_dir = tracer.get_run_dir() + _progress_file_path = run_dir / "progress.json" + return _progress_file_path + except (ImportError, AttributeError): + pass + return None + + +def _load_progress_from_disk() -> dict[str, Any]: + """Load progress data from disk if file exists.""" + global _progress_cache # noqa: PLW0603 + progress_file = _get_progress_file() + + if progress_file and progress_file.exists(): + try: + with progress_file.open("r", encoding="utf-8") as f: + _progress_cache = json.load(f) + logger.info(f"Loaded progress data with {len(_progress_cache)} keys from {progress_file}") + except (json.JSONDecodeError, OSError) as e: + logger.warning(f"Failed to load progress from disk: {e}") + _progress_cache = {} + + return _progress_cache + + +def _save_progress_to_disk() -> bool: + """Save all progress data to disk atomically.""" + progress_file = _get_progress_file() + if not progress_file: + return False + + try: + # Ensure parent directory exists + progress_file.parent.mkdir(parents=True, exist_ok=True) + + # Write atomically using temp file + temp_file = progress_file.with_suffix(".json.tmp") + with temp_file.open("w", encoding="utf-8") as f: + json.dump(_progress_cache, f, indent=2, ensure_ascii=False, default=str) + + # Atomic rename + temp_file.replace(progress_file) + logger.debug(f"Progress saved to {progress_file}") + return True + + except (OSError, IOError) as e: + logger.error(f"Failed to save progress to disk: {e}") + return False + + +@register_tool(sandbox_execution=False) +def save_progress( + key: str, + data: dict[str, Any], + append: bool = False, +) -> dict[str, Any]: + """Save structured progress data to disk for persistence across crashes. + + Use this to track: + - Scanned endpoints/targets + - Discovered services + - Intermediate findings + - Task completion status + - Any data you want to persist and reference later + + Args: + key: Unique identifier for this progress data (e.g., "scanned_ports", "discovered_services") + data: The data to save (must be JSON-serializable) + append: If True and key exists with a list value, append data to the list instead of replacing + + Returns: + Success status and file path + """ + try: + if not key or not key.strip(): + return {"success": False, "error": "Key cannot be empty"} + + if not isinstance(data, dict): + return {"success": False, "error": "Data must be a dictionary"} + + # Load existing progress on first access + if not _progress_cache: + _load_progress_from_disk() + + key = key.strip() + timestamp = datetime.now(UTC).isoformat() + + if append and key in _progress_cache: + existing = _progress_cache[key].get("data") + if isinstance(existing, list) and isinstance(data.get("items"), list): + # Append items to existing list + existing.extend(data["items"]) + _progress_cache[key]["updated_at"] = timestamp + else: + # Can't append, just update + _progress_cache[key] = { + "data": data, + "created_at": _progress_cache[key].get("created_at", timestamp), + "updated_at": timestamp, + } + else: + _progress_cache[key] = { + "data": data, + "created_at": timestamp, + "updated_at": timestamp, + } + + # Immediately persist to disk + persisted = _save_progress_to_disk() + + return { + "success": True, + "message": f"Progress '{key}' saved successfully", + "key": key, + "persisted_to_disk": persisted, + } + + except (ValueError, TypeError) as e: + return {"success": False, "error": f"Failed to save progress: {e}"} + + +@register_tool(sandbox_execution=False) +def load_progress(key: str) -> dict[str, Any]: + """Load previously saved progress data. + + Args: + key: The key used when saving the progress data + + Returns: + The saved data or error if not found + """ + try: + if not key or not key.strip(): + return {"success": False, "error": "Key cannot be empty", "data": None} + + # Load from disk on first access + if not _progress_cache: + _load_progress_from_disk() + + key = key.strip() + + if key not in _progress_cache: + return { + "success": False, + "error": f"Progress key '{key}' not found", + "data": None, + "available_keys": list(_progress_cache.keys()), + } + + entry = _progress_cache[key] + return { + "success": True, + "key": key, + "data": entry.get("data"), + "created_at": entry.get("created_at"), + "updated_at": entry.get("updated_at"), + } + + except (ValueError, TypeError) as e: + return {"success": False, "error": f"Failed to load progress: {e}", "data": None} + + +@register_tool(sandbox_execution=False) +def list_progress() -> dict[str, Any]: + """List all available progress keys and their metadata. + + Returns: + List of all progress keys with creation/update timestamps + """ + try: + # Load from disk on first access + if not _progress_cache: + _load_progress_from_disk() + + progress_list = [] + for key, entry in _progress_cache.items(): + data = entry.get("data", {}) + # Calculate size hint + if isinstance(data, dict): + size_hint = f"{len(data)} keys" + elif isinstance(data, list): + size_hint = f"{len(data)} items" + else: + size_hint = "unknown" + + progress_list.append({ + "key": key, + "created_at": entry.get("created_at"), + "updated_at": entry.get("updated_at"), + "size_hint": size_hint, + }) + + # Sort by updated_at descending + progress_list.sort(key=lambda x: x.get("updated_at", ""), reverse=True) + + return { + "success": True, + "progress": progress_list, + "total_count": len(progress_list), + } + + except (ValueError, TypeError) as e: + return { + "success": False, + "error": f"Failed to list progress: {e}", + "progress": [], + "total_count": 0, + } diff --git a/strix/tools/proxy/proxy_manager.py b/strix/tools/proxy/proxy_manager.py index e02d85b7d..59a2f5758 100644 --- a/strix/tools/proxy/proxy_manager.py +++ b/strix/tools/proxy/proxy_manager.py @@ -780,6 +780,7 @@ def close(self) -> None: def get_proxy_manager() -> ProxyManager: + global _PROXY_MANAGER if _PROXY_MANAGER is None: - return ProxyManager() + _PROXY_MANAGER = ProxyManager() return _PROXY_MANAGER diff --git a/strix/tools/registry.py b/strix/tools/registry.py index e43b1bbc5..2d92d7ea1 100644 --- a/strix/tools/registry.py +++ b/strix/tools/registry.py @@ -21,64 +21,136 @@ "send_message_to_agent", "wait_for_message", "think", + # Progress tracking (available to all roles) + "save_progress", + "load_progress", + "list_progress", + # Notes + "create_note", + "list_notes", ], "recon": [ - "terminal", - "python", - "browser", - "proxy", + "terminal_execute", + "python_action", + "browser_action", + "list_requests", + "view_request", + "send_request", + "repeat_request", + "scope_rules", + "list_sitemap", + "view_sitemap_entry", "think", "agent_finish", "create_agent", "view_agent_graph", "send_message_to_agent", "wait_for_message", - "read_file", - "write_file", - "list_directory", + "str_replace_editor", + "list_files", + "search_files", "web_search", + # Progress tracking + "save_progress", + "load_progress", + "list_progress", + # Notes + "create_note", + "list_notes", + "update_note", ], "testing": [ - "terminal", - "python", - "browser", - "proxy", + "terminal_execute", + "python_action", + "browser_action", + "list_requests", + "view_request", + "send_request", + "repeat_request", + "scope_rules", + "list_sitemap", + "view_sitemap_entry", "think", "agent_finish", "create_agent", "view_agent_graph", "send_message_to_agent", "wait_for_message", - "read_file", - "write_file", + "str_replace_editor", + "list_files", + "search_files", "web_search", + # Progress tracking + "save_progress", + "load_progress", + "list_progress", + # Notes + "create_note", + "list_notes", + "update_note", + # Vulnerability reporting + "create_vulnerability_report", ], "validation": [ - "terminal", - "python", - "browser", - "proxy", + "terminal_execute", + "python_action", + "browser_action", + "list_requests", + "view_request", + "send_request", + "repeat_request", + "scope_rules", + "list_sitemap", + "view_sitemap_entry", "think", "agent_finish", - "read_file", + "str_replace_editor", + "list_files", + "search_files", "send_message_to_agent", + # Progress tracking + "save_progress", + "load_progress", + "list_progress", + # Notes + "create_note", + "list_notes", + # Vulnerability reporting + "create_vulnerability_report", ], "reporting": [ "create_vulnerability_report", - "read_file", - "write_file", + "str_replace_editor", + "list_files", + "search_files", "think", "agent_finish", "send_message_to_agent", + # Progress tracking + "save_progress", + "load_progress", + "list_progress", + # Notes + "create_note", + "list_notes", + "update_note", ], "fixing": [ - "read_file", - "write_file", - "terminal", - "python", + "str_replace_editor", + "list_files", + "search_files", + "terminal_execute", + "python_action", "think", "agent_finish", "send_message_to_agent", + # Progress tracking + "save_progress", + "load_progress", + "list_progress", + # Notes + "create_note", + "list_notes", ], } diff --git a/strix/tools/reporting/reporting_actions.py b/strix/tools/reporting/reporting_actions.py index dd98d6dbc..ffac75191 100644 --- a/strix/tools/reporting/reporting_actions.py +++ b/strix/tools/reporting/reporting_actions.py @@ -1,8 +1,65 @@ +import csv +import logging +from datetime import UTC, datetime +from pathlib import Path from typing import Any from strix.tools.registry import register_tool +logger = logging.getLogger(__name__) + + +def _save_vulnerability_to_disk( + run_dir: Path, + report_id: str, + title: str, + content: str, + severity: str, + timestamp: str, +) -> bool: + """Immediately save vulnerability report to disk (crash-proof).""" + try: + # Ensure directories exist + vuln_dir = run_dir / "vulnerabilities" + vuln_dir.mkdir(parents=True, exist_ok=True) + + # Write individual vulnerability markdown file + vuln_file = vuln_dir / f"{report_id}.md" + with vuln_file.open("w", encoding="utf-8") as f: + f.write(f"# {title}\n\n") + f.write(f"**ID:** {report_id}\n") + f.write(f"**Severity:** {severity.upper()}\n") + f.write(f"**Found:** {timestamp}\n\n") + f.write("## Description\n\n") + f.write(f"{content}\n") + + # Append to CSV index (create header if new file) + csv_file = run_dir / "vulnerabilities.csv" + write_header = not csv_file.exists() + + with csv_file.open("a", encoding="utf-8", newline="") as f: + writer = csv.DictWriter( + f, fieldnames=["id", "title", "severity", "timestamp", "file"] + ) + if write_header: + writer.writeheader() + writer.writerow({ + "id": report_id, + "title": title, + "severity": severity.upper(), + "timestamp": timestamp, + "file": f"vulnerabilities/{report_id}.md", + }) + + logger.info(f"Vulnerability {report_id} saved to disk: {vuln_file}") + return True + + except (OSError, IOError) as e: + logger.error(f"Failed to save vulnerability {report_id} to disk: {e}") + return False + + @register_tool(sandbox_execution=False) def create_vulnerability_report( title: str, @@ -37,13 +94,26 @@ def create_vulnerability_report( severity=severity, ) + # IMMEDIATELY save to disk (crash-proof) + timestamp = datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S UTC") + run_dir = tracer.get_run_dir() + disk_saved = _save_vulnerability_to_disk( + run_dir=run_dir, + report_id=report_id, + title=title.strip(), + content=content.strip(), + severity=severity.lower().strip(), + timestamp=timestamp, + ) + return { "success": True, "message": f"Vulnerability report '{title}' created successfully", "report_id": report_id, "severity": severity.lower(), + "persisted_to_disk": disk_saved, + "file_path": f"vulnerabilities/{report_id}.md" if disk_saved else None, } - import logging logging.warning("Global tracer not available - vulnerability report not stored") diff --git a/templates/scope/proxmox-cluster.yaml b/templates/scope/proxmox-cluster.yaml index 41b5e9785..1e747c333 100644 --- a/templates/scope/proxmox-cluster.yaml +++ b/templates/scope/proxmox-cluster.yaml @@ -30,12 +30,12 @@ networks: description: "Ceph/storage cluster network" targets: - # Proxmox nodes + # Proxmox node on management network - host: "10.0.101.2" - name: "pve-node-01" + name: "pve-mgmt" type: "infrastructure" network: "Proxmox Management" - ports: [22, 8006, 3128, 5900] # VNC starts at 5900 + ports: [22, 8006, 3128, 5900] services: - port: 8006 service: "proxmox-ve" @@ -49,32 +49,41 @@ targets: tags: ["hypervisor", "cluster-master", "critical"] modules: ["proxmox_ve"] - - host: "10.0.101.3" - name: "pve-node-02" + # Proxmox cluster nodes on VM network + - host: "10.0.102.2" + name: "pve-node-01" type: "infrastructure" - network: "Proxmox Management" + network: "VM Network" ports: [22, 8006, 3128] - tags: ["hypervisor", "cluster-member"] + services: + - port: 8006 + service: "proxmox-ve" + version: "8.x" + tags: ["hypervisor", "cluster-member", "critical"] modules: ["proxmox_ve"] - - host: "10.0.101.4" - name: "pve-node-03" + - host: "10.0.102.3" + name: "pve-node-02" type: "infrastructure" - network: "Proxmox Management" + network: "VM Network" ports: [22, 8006, 3128] - tags: ["hypervisor", "cluster-member"] + services: + - port: 8006 + service: "proxmox-ve" + version: "8.x" + tags: ["hypervisor", "cluster-member", "critical"] modules: ["proxmox_ve"] - # PBS (Proxmox Backup Server) if present - - host: "10.0.101.10" - name: "pbs-backup-01" + - host: "10.0.102.4" + name: "pve-node-03" type: "infrastructure" - network: "Proxmox Management" - ports: [22, 8007] + network: "VM Network" + ports: [22, 8006, 3128] services: - - port: 8007 - service: "proxmox-backup" - tags: ["backup", "critical"] + - port: 8006 + service: "proxmox-ve" + version: "8.x" + tags: ["hypervisor", "cluster-member", "critical"] modules: ["proxmox_ve"] exclusions: From f462b97b58b27c9fca5a96ed2a0f7d533284de90 Mon Sep 17 00:00:00 2001 From: yokoszn Date: Sun, 23 Nov 2025 21:22:32 +1100 Subject: [PATCH 3/3] feat: add performance optimizations and stability improvements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Performance: - Tool parallelization: read-only tools now execute concurrently - Streaming LLM responses with early termination on - LLM response caching with configurable TTL and LRU eviction Stability: - Circuit breaker pattern for LLM API (prevents cascading failures) - Worker health checks with automatic restart on death - Graceful degradation for tool parsing failures - Individual error isolation in parallel tool execution - Response queue timeout to prevent indefinite blocking Bug fixes: - Fix asyncio Semaphore event loop binding issue (revert to threading) New environment variables: - LLM_STREAMING: enable/disable streaming (default: true) - LLM_CACHE_ENABLED: enable response caching (default: true) - LLM_CACHE_MAX_SIZE: max cached responses (default: 100) - LLM_CACHE_TTL: cache TTL in seconds (default: 3600) - LLM_CIRCUIT_FAILURE_THRESHOLD: failures before circuit opens (default: 5) - LLM_CIRCUIT_RECOVERY_TIMEOUT: seconds before retry (default: 60) Docs: - Add scope configuration section to README πŸ€– Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- README.md | 38 +++++++ strix/llm/__init__.py | 7 ++ strix/llm/circuit_breaker.py | 185 +++++++++++++++++++++++++++++++++++ strix/llm/config.py | 7 ++ strix/llm/llm.py | 124 +++++++++++++++++++++-- strix/llm/request_queue.py | 130 +++++++++++++++++++----- strix/llm/response_cache.py | 160 ++++++++++++++++++++++++++++++ strix/runtime/tool_server.py | 89 ++++++++++++++--- strix/tools/executor.py | 81 ++++++++++++++- 9 files changed, 774 insertions(+), 47 deletions(-) create mode 100644 strix/llm/circuit_breaker.py create mode 100644 strix/llm/response_cache.py diff --git a/README.md b/README.md index 2efa9d2e3..0ebe498ab 100644 --- a/README.md +++ b/README.md @@ -196,6 +196,44 @@ jobs: run: strix -n -t ./ ``` +### πŸ“‹ Scope Configuration + +For complex assessments with multiple targets, networks, and credentials, use a scope configuration file: + +```bash +# Run with a scope file +strix --scope scope.yaml + +# Validate scope file before running +strix --scope scope.yaml --validate + +# Filter to specific targets +strix --scope scope.yaml --filter "tags:critical" +strix --scope scope.yaml --filter "network:Management" +``` + +**Scope templates** are available in `templates/scope/` for common scenarios: + +```bash +# Proxmox cluster assessment +strix --scope templates/scope/proxmox-cluster.yaml + +# Filter to management network only +strix --scope templates/scope/proxmox-cluster.yaml --filter "network:Proxmox Management" +``` + +Technology modules provide specialized testing guidance. Enable them per-target in your scope file: + +```yaml +targets: + - host: "10.0.101.2" + name: "pve-node" + type: "infrastructure" + modules: ["proxmox_ve"] # Loads Proxmox VE testing guidance +``` + +See [Scope Configuration Guide](docs/SCOPE_CONFIGURATION.md) for full documentation. + ### βš™οΈ Configuration ```bash diff --git a/strix/llm/__init__.py b/strix/llm/__init__.py index 33ce1b6df..96274d70b 100644 --- a/strix/llm/__init__.py +++ b/strix/llm/__init__.py @@ -2,14 +2,21 @@ import litellm +from .circuit_breaker import CircuitBreaker, CircuitBreakerError, get_llm_circuit_breaker from .config import LLMConfig from .llm import LLM, LLMRequestFailedError +from .response_cache import ResponseCache, get_global_cache __all__ = [ "LLM", "LLMConfig", "LLMRequestFailedError", + "CircuitBreaker", + "CircuitBreakerError", + "get_llm_circuit_breaker", + "ResponseCache", + "get_global_cache", ] # Only disable debugging if not in debug mode diff --git a/strix/llm/circuit_breaker.py b/strix/llm/circuit_breaker.py new file mode 100644 index 000000000..f4d3b0a88 --- /dev/null +++ b/strix/llm/circuit_breaker.py @@ -0,0 +1,185 @@ +"""Circuit breaker pattern for LLM API calls. + +Prevents cascading failures when LLM API is unavailable by +failing fast after repeated failures. +""" + +import logging +import os +import threading +import time +from enum import Enum +from typing import Any + + +logger = logging.getLogger(__name__) + + +class CircuitState(str, Enum): + CLOSED = "closed" # Normal operation + OPEN = "open" # Failing fast + HALF_OPEN = "half_open" # Testing recovery + + +class CircuitBreakerError(Exception): + """Raised when circuit breaker is open.""" + + def __init__(self, message: str, time_until_retry: float): + super().__init__(message) + self.time_until_retry = time_until_retry + + +class CircuitBreaker: + """Thread-safe circuit breaker for protecting external service calls.""" + + def __init__( + self, + failure_threshold: int | None = None, + recovery_timeout: float | None = None, + half_open_max_calls: int = 1, + name: str = "default", + ): + self.failure_threshold = failure_threshold or int( + os.environ.get("LLM_CIRCUIT_FAILURE_THRESHOLD", "5") + ) + self.recovery_timeout = recovery_timeout or float( + os.environ.get("LLM_CIRCUIT_RECOVERY_TIMEOUT", "60") + ) + self.half_open_max_calls = half_open_max_calls + self.name = name + + self._state = CircuitState.CLOSED + self._failure_count = 0 + self._last_failure_time: float | None = None + self._half_open_calls = 0 + self._lock = threading.Lock() + + # Stats + self._total_calls = 0 + self._total_failures = 0 + self._total_circuit_breaks = 0 + + @property + def state(self) -> CircuitState: + """Get current circuit state, updating if recovery timeout has passed.""" + with self._lock: + if self._state == CircuitState.OPEN: + if self._should_attempt_recovery(): + self._state = CircuitState.HALF_OPEN + self._half_open_calls = 0 + logger.info(f"Circuit breaker '{self.name}' entering half-open state") + return self._state + + def _should_attempt_recovery(self) -> bool: + """Check if enough time has passed to attempt recovery.""" + if self._last_failure_time is None: + return True + return time.time() - self._last_failure_time >= self.recovery_timeout + + def _time_until_recovery(self) -> float: + """Get seconds until circuit breaker will attempt recovery.""" + if self._last_failure_time is None: + return 0 + elapsed = time.time() - self._last_failure_time + return max(0, self.recovery_timeout - elapsed) + + def can_execute(self) -> bool: + """Check if a call can be executed.""" + current_state = self.state + + if current_state == CircuitState.CLOSED: + return True + + if current_state == CircuitState.HALF_OPEN: + with self._lock: + return self._half_open_calls < self.half_open_max_calls + + return False + + def record_success(self) -> None: + """Record a successful call.""" + with self._lock: + self._total_calls += 1 + + if self._state == CircuitState.HALF_OPEN: + # Recovery successful, close circuit + self._state = CircuitState.CLOSED + self._failure_count = 0 + self._half_open_calls = 0 + logger.info(f"Circuit breaker '{self.name}' recovered, closing circuit") + elif self._state == CircuitState.CLOSED: + # Reset failure count on success + self._failure_count = 0 + + def record_failure(self, error: Exception | None = None) -> None: + """Record a failed call.""" + with self._lock: + self._total_calls += 1 + self._total_failures += 1 + self._failure_count += 1 + self._last_failure_time = time.time() + + if self._state == CircuitState.HALF_OPEN: + # Recovery failed, reopen circuit + self._state = CircuitState.OPEN + self._total_circuit_breaks += 1 + logger.warning( + f"Circuit breaker '{self.name}' recovery failed, reopening circuit" + ) + elif self._state == CircuitState.CLOSED: + if self._failure_count >= self.failure_threshold: + self._state = CircuitState.OPEN + self._total_circuit_breaks += 1 + logger.warning( + f"Circuit breaker '{self.name}' opened after {self._failure_count} failures" + ) + + def raise_if_open(self) -> None: + """Raise CircuitBreakerError if circuit is open.""" + if not self.can_execute(): + time_until_retry = self._time_until_recovery() + raise CircuitBreakerError( + f"Circuit breaker '{self.name}' is open. Service unavailable. " + f"Retry in {time_until_retry:.1f}s", + time_until_retry=time_until_retry, + ) + + # Track half-open calls + with self._lock: + if self._state == CircuitState.HALF_OPEN: + self._half_open_calls += 1 + + def reset(self) -> None: + """Manually reset the circuit breaker.""" + with self._lock: + self._state = CircuitState.CLOSED + self._failure_count = 0 + self._half_open_calls = 0 + logger.info(f"Circuit breaker '{self.name}' manually reset") + + @property + def stats(self) -> dict[str, Any]: + """Get circuit breaker statistics.""" + with self._lock: + return { + "name": self.name, + "state": self._state.value, + "failure_count": self._failure_count, + "failure_threshold": self.failure_threshold, + "total_calls": self._total_calls, + "total_failures": self._total_failures, + "total_circuit_breaks": self._total_circuit_breaks, + "time_until_recovery": self._time_until_recovery() if self._state == CircuitState.OPEN else 0, + } + + +# Global circuit breaker for LLM API +_llm_circuit_breaker: CircuitBreaker | None = None + + +def get_llm_circuit_breaker() -> CircuitBreaker: + """Get the global LLM circuit breaker instance.""" + global _llm_circuit_breaker # noqa: PLW0603 + if _llm_circuit_breaker is None: + _llm_circuit_breaker = CircuitBreaker(name="llm_api") + return _llm_circuit_breaker diff --git a/strix/llm/config.py b/strix/llm/config.py index 33e98898f..91d4a6d2e 100644 --- a/strix/llm/config.py +++ b/strix/llm/config.py @@ -9,6 +9,7 @@ def __init__( prompt_modules: list[str] | None = None, timeout: int | None = None, agent_role: str | None = None, + enable_streaming: bool | None = None, ): self.model_name = model_name or os.getenv("STRIX_LLM", "openai/gpt-5") @@ -20,3 +21,9 @@ def __init__( self.agent_role = agent_role self.timeout = timeout or int(os.getenv("LLM_TIMEOUT", "600")) + + # Streaming enabled by default, can be disabled via env var + if enable_streaming is not None: + self.enable_streaming = enable_streaming + else: + self.enable_streaming = os.getenv("LLM_STREAMING", "true").lower() == "true" diff --git a/strix/llm/llm.py b/strix/llm/llm.py index a28a1d602..cb041ba11 100644 --- a/strix/llm/llm.py +++ b/strix/llm/llm.py @@ -17,9 +17,11 @@ from litellm import ModelResponse, completion_cost from litellm.utils import supports_prompt_caching +from strix.llm.circuit_breaker import CircuitBreakerError, get_llm_circuit_breaker from strix.llm.config import LLMConfig from strix.llm.memory_compressor import MemoryCompressor from strix.llm.request_queue import get_global_queue +from strix.llm.response_cache import get_global_cache from strix.llm.utils import _truncate_to_first_function, parse_tool_invocations from strix.prompts import load_prompt_modules from strix.tools import get_tools_prompt @@ -339,6 +341,14 @@ async def generate( # noqa: PLR0912, PLR0915 model=self.config.model_name, ) + # Check circuit breaker before attempting request + circuit_breaker = get_llm_circuit_breaker() + try: + circuit_breaker.raise_if_open() + except CircuitBreakerError as e: + self._emit_llm_error(tracer, request_id, start_time, str(e), "CircuitBreakerError") + raise LLMRequestFailedError(str(e), f"Retry in {e.time_until_retry:.1f}s") from e + # Retry loop with exponential backoff for transient errors last_error: Exception | None = None retry_delay = INITIAL_RETRY_DELAY @@ -348,6 +358,9 @@ async def generate( # noqa: PLR0912, PLR0915 response = await self._make_request(cached_messages) self._update_usage_stats(response) + # Record success with circuit breaker + circuit_breaker.record_success() + # Emit LLM response event duration_ms = (time.time() - start_time) * 1000 if tracer and self.agent_id and request_id: @@ -361,10 +374,12 @@ async def generate( # noqa: PLR0912, PLR0915 cached_tokens=self._last_request_stats.cached_tokens, ) + # Extract content with validation content = "" - if ( - response.choices - and hasattr(response.choices[0], "message") + if not response.choices: + logger.warning("LLM returned empty choices, using empty content") + elif ( + hasattr(response.choices[0], "message") and response.choices[0].message ): content = getattr(response.choices[0].message, "content", "") or "" @@ -375,7 +390,13 @@ async def generate( # noqa: PLR0912, PLR0915 function_end_index = content.find("") + len("") content = content[:function_end_index] - tool_invocations = parse_tool_invocations(content) + # Parse tool invocations with graceful degradation + tool_invocations = None + try: + tool_invocations = parse_tool_invocations(content) + except Exception as e: # noqa: BLE001 + logger.warning(f"Failed to parse tool invocations: {e}") + # Continue with None tool_invocations - agent will handle text response return LLMResponse( scan_id=scan_id, @@ -389,6 +410,9 @@ async def generate( # noqa: PLR0912, PLR0915 last_error = e error_type = type(e).__name__ + # Record failure with circuit breaker (transient errors count) + circuit_breaker.record_failure(e) + if attempt < MAX_RETRIES: # Log retry attempt logger.warning( @@ -420,14 +444,15 @@ async def generate( # noqa: PLR0912, PLR0915 litellm.InvalidRequestError, litellm.BadRequestError, ) as e: - # Non-retryable errors - fail immediately + # Non-retryable errors - fail immediately (don't affect circuit breaker) error_type = type(e).__name__ error_msg = self._get_error_message(e) self._emit_llm_error(tracer, request_id, start_time, error_msg, error_type, str(e)) raise LLMRequestFailedError(f"LLM request failed: {error_msg}", str(e)) from e except Exception as e: - # Unknown errors - fail immediately + # Unknown errors - record with circuit breaker and fail + circuit_breaker.record_failure(e) error_type = type(e).__name__ self._emit_llm_error(tracer, request_id, start_time, str(e), error_type, str(e)) raise LLMRequestFailedError(f"LLM request failed: {error_type}", str(e)) from e @@ -508,6 +533,23 @@ def _should_include_reasoning_effort(self) -> bool: return model_matches(self.config.model_name, REASONING_EFFORT_PATTERNS) + def _should_use_streaming(self) -> bool: + """Check if streaming should be used for this request.""" + if not self.config.enable_streaming: + return False + + # Some models don't support streaming well + if not self.config.model_name: + return False + + # Disable streaming for reasoning models (they often don't support it well) + model_lower = self.config.model_name.lower() + no_stream_patterns = ["o1", "o3", "o4"] + if any(pat in model_lower for pat in no_stream_patterns): + return False + + return True + async def _make_request( self, messages: list[dict[str, Any]], @@ -524,14 +566,82 @@ async def _make_request( if self._should_include_reasoning_effort(): completion_args["reasoning_effort"] = "high" + # Check cache first + cache = get_global_cache() + cached_response = cache.get( + model=completion_args["model"], + messages=completion_args["messages"], + ) + if cached_response is not None: + logger.debug("Using cached LLM response") + self._total_stats.requests += 1 + self._last_request_stats = RequestStats(requests=1) + return cached_response + queue = get_global_queue() - response = await queue.make_request(completion_args) + + if self._should_use_streaming(): + response = await self._make_streaming_request(queue, completion_args) + else: + response = await queue.make_request(completion_args) + + # Cache the response + cache.put( + model=completion_args["model"], + messages=completion_args["messages"], + response=response, + ) self._total_stats.requests += 1 self._last_request_stats = RequestStats(requests=1) return response + async def _make_streaming_request( + self, + queue: Any, + completion_args: dict[str, Any], + ) -> ModelResponse: + """Make a streaming request with early termination on tag.""" + + def stop_on_function_end(content: str) -> bool: + return "" in content + + content, usage_chunk = await queue.make_streaming_request( + completion_args, + on_chunk=None, # Could add callback for real-time display + stop_condition=stop_on_function_end, + ) + + # Build a synthetic ModelResponse from streamed content + # This maintains compatibility with the rest of the code + from litellm import ModelResponse as LiteLLMResponse + from litellm.utils import Usage, Choices, Message + + message = Message(content=content, role="assistant") + choice = Choices(index=0, message=message, finish_reason="stop") + + # Use usage from final chunk if available, otherwise estimate + if usage_chunk and hasattr(usage_chunk, "usage") and usage_chunk.usage: + usage = usage_chunk.usage + else: + # Estimate tokens (rough approximation) + usage = Usage( + prompt_tokens=0, # Will be updated from actual response + completion_tokens=len(content) // 4, # ~4 chars per token + total_tokens=len(content) // 4, + ) + + response = LiteLLMResponse( + id="stream-" + str(time.time()), + choices=[choice], + created=int(time.time()), + model=self.config.model_name, + usage=usage, + ) + + return response + def _update_usage_stats(self, response: ModelResponse) -> None: try: if hasattr(response, "usage") and response.usage: diff --git a/strix/llm/request_queue.py b/strix/llm/request_queue.py index 770639a67..a86065a73 100644 --- a/strix/llm/request_queue.py +++ b/strix/llm/request_queue.py @@ -1,6 +1,9 @@ import asyncio import logging +import os +import threading import time +from collections.abc import AsyncIterator, Callable from typing import Any import litellm @@ -25,37 +28,64 @@ def should_retry_exception(exception: Exception) -> bool: 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: asyncio.Semaphore | None = None + def __init__(self, max_concurrent: int | None = None, delay_between_requests: float | None = None): + self.max_concurrent = max_concurrent or int(os.environ.get("LLM_RATE_LIMIT_CONCURRENT", "6")) + self.delay_between_requests = delay_between_requests or float( + os.environ.get("LLM_RATE_LIMIT_DELAY", "1.0") + ) + # Use threading primitives - they work across all event loops + self._semaphore = threading.BoundedSemaphore(self.max_concurrent) self._last_request_time = 0.0 - self._lock: asyncio.Lock | None = None - - def _get_semaphore(self) -> asyncio.Semaphore: - """Lazily create semaphore in the current event loop.""" - if self._semaphore is None: - self._semaphore = asyncio.Semaphore(self.max_concurrent) - return self._semaphore - - def _get_lock(self) -> asyncio.Lock: - """Lazily create lock in the current event loop.""" - if self._lock is None: - self._lock = asyncio.Lock() - return self._lock + self._lock = threading.Lock() + + async def _acquire_slot(self) -> float: + """Acquire a request slot and return the sleep time needed.""" + self._semaphore.acquire() + 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 + return sleep_needed + + def _release_slot(self) -> None: + """Release the request slot.""" + self._semaphore.release() async def make_request(self, completion_args: dict[str, Any]) -> ModelResponse: - async with self._get_semaphore(): - async with self._get_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 - + sleep_needed = await self._acquire_slot() + try: if sleep_needed > 0: await asyncio.sleep(sleep_needed) return await self._reliable_request(completion_args) + finally: + self._release_slot() + + async def make_streaming_request( + self, + completion_args: dict[str, Any], + on_chunk: Callable[[str], None] | None = None, + stop_condition: Callable[[str], bool] | None = None, + ) -> tuple[str, ModelResponse | None]: + """Make a streaming request, accumulating content and optionally stopping early. + + Args: + completion_args: Arguments to pass to litellm completion + on_chunk: Optional callback for each content chunk + stop_condition: Optional function that returns True to stop streaming early + + Returns: + Tuple of (accumulated_content, final_response_with_usage_or_None) + """ + sleep_needed = await self._acquire_slot() + try: + if sleep_needed > 0: + await asyncio.sleep(sleep_needed) + + return await self._streaming_request(completion_args, on_chunk, stop_condition) + finally: + self._release_slot() @retry( # type: ignore[misc] stop=stop_after_attempt(5), @@ -70,6 +100,58 @@ async def _reliable_request(self, completion_args: dict[str, Any]) -> ModelRespo self._raise_unexpected_response() raise RuntimeError("Unreachable code") + async def _streaming_request( + self, + completion_args: dict[str, Any], + on_chunk: Callable[[str], None] | None = None, + stop_condition: Callable[[str], bool] | None = None, + ) -> tuple[str, ModelResponse | None]: + """Execute streaming request with early termination support.""" + accumulated_content = "" + final_response: ModelResponse | None = None + + # Start streaming + stream = completion(**completion_args, stream=True) + + try: + async for chunk in self._iter_stream(stream): + # Extract content from chunk + if hasattr(chunk, "choices") and chunk.choices: + delta = getattr(chunk.choices[0], "delta", None) + if delta: + content = getattr(delta, "content", None) + if content: + accumulated_content += content + if on_chunk: + on_chunk(content) + + # Check for early termination + if stop_condition and stop_condition(accumulated_content): + logger.debug("Streaming stopped early due to stop condition") + break + + # Capture usage from final chunk if available + if hasattr(chunk, "usage") and chunk.usage: + final_response = chunk + + except GeneratorExit: + pass + + return accumulated_content, final_response + + async def _iter_stream(self, stream: Any) -> AsyncIterator[Any]: + """Iterate over stream, handling both sync and async generators.""" + if hasattr(stream, "__anext__"): + # Async generator + async for chunk in stream: + yield chunk + else: + # Sync generator - wrap in async + for chunk in stream: + yield chunk + # Yield control to event loop periodically + await asyncio.sleep(0) + def _raise_unexpected_response(self) -> None: raise RuntimeError("Unexpected response type") diff --git a/strix/llm/response_cache.py b/strix/llm/response_cache.py new file mode 100644 index 000000000..c73ba16e2 --- /dev/null +++ b/strix/llm/response_cache.py @@ -0,0 +1,160 @@ +"""LLM Response Cache for reducing redundant API calls. + +Provides an in-memory LRU cache with TTL expiration for LLM responses. +Useful for caching identical requests during parallel agent operations. +""" + +import hashlib +import json +import logging +import os +import threading +import time +from collections import OrderedDict +from dataclasses import dataclass +from typing import Any + +from litellm import ModelResponse + + +logger = logging.getLogger(__name__) + + +@dataclass +class CacheEntry: + """A cached LLM response with metadata.""" + + response: ModelResponse + created_at: float + hits: int = 0 + + +class ResponseCache: + """Thread-safe LRU cache for LLM responses with TTL expiration.""" + + def __init__( + self, + max_size: int | None = None, + ttl_seconds: float | None = None, + enabled: bool | None = None, + ): + self.max_size = max_size or int(os.environ.get("LLM_CACHE_MAX_SIZE", "100")) + self.ttl_seconds = ttl_seconds or float(os.environ.get("LLM_CACHE_TTL", "3600")) + + if enabled is not None: + self.enabled = enabled + else: + self.enabled = os.environ.get("LLM_CACHE_ENABLED", "true").lower() == "true" + + self._cache: OrderedDict[str, CacheEntry] = OrderedDict() + self._lock = threading.Lock() + self._stats = {"hits": 0, "misses": 0, "evictions": 0} + + def _generate_key(self, model: str, messages: list[dict[str, Any]]) -> str: + """Generate a deterministic cache key from request parameters.""" + # Create a hashable representation of the request + key_data = { + "model": model, + "messages": messages, + } + key_json = json.dumps(key_data, sort_keys=True, default=str) + return hashlib.sha256(key_json.encode()).hexdigest()[:32] + + def get(self, model: str, messages: list[dict[str, Any]]) -> ModelResponse | None: + """Get a cached response if available and not expired.""" + if not self.enabled: + return None + + key = self._generate_key(model, messages) + + with self._lock: + if key not in self._cache: + self._stats["misses"] += 1 + return None + + entry = self._cache[key] + + # Check TTL expiration + if time.time() - entry.created_at > self.ttl_seconds: + del self._cache[key] + self._stats["misses"] += 1 + self._stats["evictions"] += 1 + return None + + # Move to end (most recently used) + self._cache.move_to_end(key) + entry.hits += 1 + self._stats["hits"] += 1 + + logger.debug(f"Cache hit for key {key[:8]}... (hits: {entry.hits})") + return entry.response + + def put(self, model: str, messages: list[dict[str, Any]], response: ModelResponse) -> None: + """Cache a response.""" + if not self.enabled: + return + + key = self._generate_key(model, messages) + + with self._lock: + # Remove oldest entries if at capacity + while len(self._cache) >= self.max_size: + oldest_key = next(iter(self._cache)) + del self._cache[oldest_key] + self._stats["evictions"] += 1 + + self._cache[key] = CacheEntry( + response=response, + created_at=time.time(), + ) + + logger.debug(f"Cached response for key {key[:8]}... (cache size: {len(self._cache)})") + + def invalidate(self, model: str | None = None) -> int: + """Invalidate cache entries, optionally filtered by model. + + Returns the number of entries invalidated. + """ + with self._lock: + if model is None: + count = len(self._cache) + self._cache.clear() + return count + + # Would need to store model in entry to filter - for now just clear all + count = len(self._cache) + self._cache.clear() + return count + + def clear(self) -> None: + """Clear all cached entries.""" + with self._lock: + self._cache.clear() + self._stats = {"hits": 0, "misses": 0, "evictions": 0} + + @property + def stats(self) -> dict[str, int]: + """Get cache statistics.""" + with self._lock: + hit_rate = 0.0 + total = self._stats["hits"] + self._stats["misses"] + if total > 0: + hit_rate = self._stats["hits"] / total + + return { + **self._stats, + "size": len(self._cache), + "hit_rate": round(hit_rate, 3), + } + + +# Global cache instance +_global_cache: ResponseCache | None = None + + +def get_global_cache() -> ResponseCache: + """Get the global response cache instance.""" + global _global_cache # noqa: PLW0603 + if _global_cache is None: + _global_cache = ResponseCache() + return _global_cache diff --git a/strix/runtime/tool_server.py b/strix/runtime/tool_server.py index aeb61719c..258e42789 100644 --- a/strix/runtime/tool_server.py +++ b/strix/runtime/tool_server.py @@ -89,6 +89,9 @@ def agent_worker(_agent_id: str, request_queue: Queue[Any], response_queue: Queu from strix.tools.argument_parser import ArgumentConversionError, convert_arguments from strix.tools.registry import get_tool_by_name + consecutive_errors = 0 + max_consecutive_errors = 5 + while True: try: # Use timeout to prevent infinite blocking @@ -113,30 +116,79 @@ def agent_worker(_agent_id: str, request_queue: Queue[Any], response_queue: Queu result = tool_func(**converted_kwargs) response_queue.put({"result": result}) + consecutive_errors = 0 # Reset on success except (ArgumentConversionError, ValidationError) as e: response_queue.put({"error": f"Invalid arguments: {e}"}) except (RuntimeError, ValueError, ImportError) as e: response_queue.put({"error": f"Tool execution error: {e}"}) + except Exception as e: # noqa: BLE001 + # Catch-all for unexpected errors + consecutive_errors += 1 + response_queue.put({"error": f"Unexpected error ({type(e).__name__}): {e}"}) + if consecutive_errors >= max_consecutive_errors: + response_queue.put({"error": "Worker terminating due to repeated errors"}) + break + + except Exception as e: # noqa: BLE001 + # Critical error in worker loop itself + try: + response_queue.put({"error": f"Critical worker error: {e}"}) + except Exception: # noqa: BLE001 + pass # Queue might be broken + break # Exit worker on critical errors + + +def _create_agent_process(agent_id: str) -> tuple[Queue[Any], Queue[Any]]: + """Create a new worker process for an agent.""" + request_queue: Queue[Any] = Queue() + response_queue: Queue[Any] = Queue() + + process = Process( + target=agent_worker, args=(agent_id, request_queue, response_queue), daemon=True + ) + process.start() + + agent_processes[agent_id] = {"process": process, "pid": process.pid} + agent_queues[agent_id] = {"request": request_queue, "response": response_queue} - except (RuntimeError, ValueError, ImportError) as e: - response_queue.put({"error": f"Worker error: {e}"}) + return request_queue, response_queue + + +def _cleanup_dead_process(agent_id: str) -> None: + """Clean up resources for a dead worker process.""" + if agent_id in agent_processes: + try: + process = agent_processes[agent_id]["process"] + if process.is_alive(): + process.terminate() + process.join(timeout=1) + except Exception: # noqa: BLE001 + pass + del agent_processes[agent_id] + + if agent_id in agent_queues: + del agent_queues[agent_id] def ensure_agent_process(agent_id: str) -> tuple[Queue[Any], Queue[Any]]: - if agent_id not in agent_processes: - request_queue: Queue[Any] = Queue() - response_queue: Queue[Any] = Queue() + """Ensure a healthy worker process exists for the agent, restarting if needed.""" + if agent_id in agent_processes: + process = agent_processes[agent_id]["process"] + if not process.is_alive(): + # Process died, clean up and restart + logging.getLogger(__name__).warning( + f"Agent worker {agent_id} (pid {agent_processes[agent_id]['pid']}) died, restarting" + ) + _cleanup_dead_process(agent_id) + return _create_agent_process(agent_id) - process = Process( - target=agent_worker, args=(agent_id, request_queue, response_queue), daemon=True - ) - process.start() + return agent_queues[agent_id]["request"], agent_queues[agent_id]["response"] - agent_processes[agent_id] = {"process": process, "pid": process.pid} - agent_queues[agent_id] = {"request": request_queue, "response": response_queue} + return _create_agent_process(agent_id) - return agent_queues[agent_id]["request"], agent_queues[agent_id]["response"] + +RESPONSE_TIMEOUT = 180 # 3 minutes max wait for tool response @app.post("/execute", response_model=ToolExecutionResponse) @@ -151,7 +203,16 @@ async def execute_tool( try: loop = asyncio.get_event_loop() - response = await loop.run_in_executor(None, response_queue.get) + + # Use timeout to prevent indefinite blocking + def get_response_with_timeout() -> dict[str, Any]: + from queue import Empty + try: + return response_queue.get(timeout=RESPONSE_TIMEOUT) + except Empty: + return {"error": f"Tool execution timed out after {RESPONSE_TIMEOUT}s"} + + response = await loop.run_in_executor(None, get_response_with_timeout) if "error" in response: return ToolExecutionResponse(error=response["error"]) @@ -159,6 +220,8 @@ async def execute_tool( except (RuntimeError, ValueError, OSError) as e: return ToolExecutionResponse(error=f"Worker error: {e}") + except Exception as e: # noqa: BLE001 + return ToolExecutionResponse(error=f"Unexpected error: {type(e).__name__}: {e}") @app.post("/register_agent") diff --git a/strix/tools/executor.py b/strix/tools/executor.py index 15c656c32..c2f991309 100644 --- a/strix/tools/executor.py +++ b/strix/tools/executor.py @@ -1,4 +1,6 @@ +import asyncio import inspect +import logging import os import time import traceback @@ -359,6 +361,25 @@ def _get_tracer_and_agent_id(agent_state: Any | None) -> tuple[Any | None, str]: return tracer, agent_id +# Tools that must run sequentially (state-modifying or terminal) +SEQUENTIAL_TOOLS = frozenset({ + "finish_scan", + "agent_finish", + "create_agent", + "send_message_to_agent", + "create_vulnerability_report", + "save_progress", + "create_note", + "update_note", + "str_replace_editor", +}) + + +def _is_parallelizable(tool_name: str) -> bool: + """Check if a tool can be executed in parallel with others.""" + return tool_name not in SEQUENTIAL_TOOLS + + async def process_tool_invocations( tool_invocations: list[dict[str, Any]], conversation_history: list[dict[str, Any]], @@ -370,10 +391,64 @@ async def process_tool_invocations( tracer, agent_id = _get_tracer_and_agent_id(agent_state) - for tool_inv in tool_invocations: - observation_xml, images, tool_should_finish = await _execute_single_tool( - tool_inv, agent_state, tracer, agent_id + # Separate parallelizable and sequential tools while preserving order + parallel_batch: list[tuple[int, dict[str, Any]]] = [] + results: dict[int, tuple[str, list[dict[str, Any]], bool]] = {} + + async def execute_parallel_batch() -> None: + """Execute accumulated parallel tools concurrently.""" + nonlocal parallel_batch + if not parallel_batch: + return + + async def run_indexed(idx: int, inv: dict[str, Any]) -> tuple[int, tuple[str, list[dict[str, Any]], bool]]: + try: + result = await _execute_single_tool(inv, agent_state, tracer, agent_id) + return idx, result + except Exception as e: # noqa: BLE001 + # Return error result instead of raising + tool_name = inv.get("toolName", "unknown") + error_msg = f"Error executing {tool_name}: {type(e).__name__}: {e}" + error_xml = ( + f"\n{tool_name}\n" + f"{error_msg}\n" + ) + return idx, (error_xml, [], False) + + batch_results = await asyncio.gather( + *[run_indexed(idx, inv) for idx, inv in parallel_batch], + return_exceptions=True ) + + for item in batch_results: + if isinstance(item, Exception): + # Should not happen with our try/except above, but handle anyway + logging.getLogger(__name__).error(f"Unexpected error in parallel batch: {item}") + continue + idx, result = item + results[idx] = result + + parallel_batch = [] + + for idx, tool_inv in enumerate(tool_invocations): + tool_name = tool_inv.get("toolName", "unknown") + + if _is_parallelizable(tool_name): + # Accumulate parallelizable tools + parallel_batch.append((idx, tool_inv)) + else: + # Flush any pending parallel tools first + await execute_parallel_batch() + # Execute sequential tool immediately + result = await _execute_single_tool(tool_inv, agent_state, tracer, agent_id) + results[idx] = result + + # Flush remaining parallel tools + await execute_parallel_batch() + + # Collect results in original order + for idx in range(len(tool_invocations)): + observation_xml, images, tool_should_finish = results[idx] observation_parts.append(observation_xml) all_images.extend(images)