openhack 0.1.0__py3-none-any.whl

openhack/agents/session.py

@@ -0,0 +1,286 @@
+"""
+Session management for vulnerability scanning.
+"""
+
+import threading
+import time
+from typing import Any, Optional
+from uuid import uuid4
+from enum import Enum
+from dataclasses import dataclass, field
+
+
+class SessionStatus(str, Enum):
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    PAUSED = "paused"
+
+
+@dataclass
+class Finding:
+    """Represents a single vulnerability finding."""
+    category: str
+    severity: str
+    title: str
+    description: str
+    file_path: str
+    id: str = field(default_factory=lambda: str(uuid4()))
+    line_number: Optional[int] = None
+    code_snippet: Optional[str] = None
+    poc: Optional[str] = None
+    fix: Optional[str] = None
+    cvss_score: Optional[float] = None
+    confidence: str = "medium"
+    validated: bool = False
+    source: Optional[str] = None
+
+    def fingerprint(self) -> str:
+        import hashlib
+        file_norm = (self.file_path or "").strip().lower().split(":")[0]
+        cat_norm = self.category.strip().lower()
+        raw = f"{cat_norm}::{file_norm}"
+        return hashlib.sha256(raw.encode()).hexdigest()[:16]
+
+    def to_dict(self) -> dict:
+        d = {
+            "id": self.id,
+            "category": self.category,
+            "severity": self.severity,
+            "title": self.title,
+            "description": self.description,
+            "filePath": self.file_path,
+            "lineNumber": self.line_number,
+            "relevantCode": self.code_snippet,
+            "poc": self.poc,
+            "recommendation": self.fix,
+            "cvssScore": self.cvss_score,
+            "confidence": self.confidence,
+            "validated": self.validated,
+            "vulnerabilityType": self._generate_vulnerability_type(),
+            "fingerprint": self.fingerprint(),
+        }
+        if self.source:
+            d["verificationSource"] = self.source
+        return d
+
+    def _generate_vulnerability_type(self) -> str:
+        category_lower = self.category.lower().replace(" ", "_").replace("-", "_")
+        if "xss" in category_lower:
+            if "dangerouslysetinnerhtml" in self.description.lower():
+                return "xss_dangerously_set_html"
+            elif "innerhtml" in self.description.lower():
+                return "xss_innerhtml"
+            elif "document.write" in self.description.lower():
+                return "xss_document_write"
+            return f"xss_{category_lower}"
+        elif "sql" in category_lower or "injection" in category_lower:
+            if "raw" in self.description.lower():
+                return "sql_injection_raw_query"
+            return "sql_injection"
+        elif "idor" in category_lower:
+            return "idor_direct_object_reference"
+        elif "ssrf" in category_lower:
+            return "ssrf_server_side_request"
+        elif "csrf" in category_lower:
+            return "csrf_missing_token"
+        elif "auth" in category_lower:
+            return "auth_bypass"
+        return category_lower
+
+
+@dataclass
+class TraceEntry:
+    """A single trace entry for debugging/logging."""
+    timestamp: float
+    agent: str
+    event_type: str
+    content: Any
+    tool_name: Optional[str] = None
+    tool_input: Optional[dict] = None
+    tool_output: Optional[Any] = None
+
+
+class Session:
+    """Session to track scan state (in-memory)."""
+
+    def __init__(
+        self,
+        target_dir: str,
+        scan_id: Optional[str] = None,
+        project_context: Optional[dict] = None,
+        trace_id: Optional[str] = None,
+        on_trace: Optional[Any] = None,
+    ):
+        self.id = scan_id or str(uuid4())
+        self.trace_id = trace_id or (self.id[:8] if self.id else str(uuid4())[:8])
+        self.target_dir = target_dir
+        self.project_context = project_context
+        self.status = SessionStatus.RUNNING
+        self.created_at = time.time()
+        self.updated_at = time.time()
+        self.current_agent: Optional[str] = None
+        self.current_step: Optional[str] = None
+        self.findings: list[Finding] = []
+        self.trace: list[TraceEntry] = []
+        self.context: dict = {}
+        self.total_cost: float = 0.0
+        self.total_tokens: int = 0
+        self.total_input_tokens: int = 0
+        self.total_output_tokens: int = 0
+        self.step_costs: dict[str, float] = {}
+        self.step_tokens: dict[str, int] = {}
+        self.step_input_tokens: dict[str, int] = {}
+        self.step_output_tokens: dict[str, int] = {}
+        self._on_trace = on_trace
+        self._user_instructions: list[str] = []
+        self._instructions_lock = threading.Lock()
+        self._instructions_version: int = 0
+        self.cancelled: bool = False
+        # Pause control: an asyncio.Event that's *set* when running (default)
+        # and *cleared* when paused. Agents call `await wait_if_paused()`
+        # between iterations, which blocks while the event is cleared.
+        # Lazily created on first access because Session may be instantiated
+        # outside an event loop (e.g. in serialization tests).
+        self._pause_event: Optional[Any] = None
+
+    def _ensure_pause_event(self) -> Any:
+        if self._pause_event is None:
+            import asyncio
+            self._pause_event = asyncio.Event()
+            self._pause_event.set()  # default: not paused
+        return self._pause_event
+
+    @property
+    def paused(self) -> bool:
+        return self._pause_event is not None and not self._pause_event.is_set()
+
+    def pause(self) -> None:
+        """Block agent loops at their next safe checkpoint."""
+        self._ensure_pause_event().clear()
+
+    def resume(self) -> None:
+        """Unblock paused agent loops."""
+        self._ensure_pause_event().set()
+
+    async def wait_if_paused(self) -> None:
+        """If the session is paused, await until resumed. No-op when not paused."""
+        await self._ensure_pause_event().wait()
+
+    def cancel(self) -> None:
+        """Set the cancellation flag so all agents break out of their loops.
+
+        Also resumes the pause event so paused agents wake up and see the
+        cancellation flag instead of blocking forever.
+        """
+        self.cancelled = True
+        if self._pause_event is not None:
+            self._pause_event.set()
+
+    def add_user_instruction(self, text: str) -> None:
+        """Thread-safe: queue an instruction from the user during a running scan."""
+        with self._instructions_lock:
+            self._user_instructions.append(text)
+            self._instructions_version += 1
+        self.add_trace(agent="user", event_type="user_instruction", content=text)
+
+    def get_new_instructions(self, seen_version: int) -> tuple[list[str], int]:
+        """Thread-safe: return instructions added since *seen_version*.
+
+        Unlike drain, this does NOT clear instructions — every agent that calls
+        this will independently see every instruction added after its own
+        watermark, which is critical for the swarm pattern where many agents
+        run concurrently.
+
+        Returns (new_instructions, current_version).
+        """
+        with self._instructions_lock:
+            current = self._instructions_version
+            if seen_version >= current:
+                return [], current
+            new = self._user_instructions[seen_version:]
+            return list(new), current
+
+    def get_all_instructions(self) -> list[str]:
+        """Thread-safe: return a snapshot of all accumulated instructions."""
+        with self._instructions_lock:
+            return list(self._user_instructions)
+
+    def add_trace(
+        self,
+        agent: str,
+        event_type: str,
+        content: Any,
+        tool_name: Optional[str] = None,
+        tool_input: Optional[dict] = None,
+        tool_output: Optional[Any] = None,
+    ) -> TraceEntry:
+        entry = TraceEntry(
+            timestamp=time.time(),
+            agent=agent,
+            event_type=event_type,
+            content=content,
+            tool_name=tool_name,
+            tool_input=tool_input,
+            tool_output=tool_output,
+        )
+        self.trace.append(entry)
+        self.updated_at = time.time()
+        if self._on_trace:
+            self._on_trace(entry)
+        return entry
+
+    def add_finding(self, finding: Finding) -> None:
+        self.findings.append(finding)
+        self.updated_at = time.time()
+
+    def get_findings_dict(self) -> list[dict]:
+        return [f.to_dict() for f in self.findings]
+
+    def record_step_cost(
+        self,
+        step_name: str,
+        cost: float,
+        tokens: int,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+    ) -> None:
+        self.step_costs[step_name] = cost
+        self.step_tokens[step_name] = tokens
+        self.step_input_tokens[step_name] = input_tokens
+        self.step_output_tokens[step_name] = output_tokens
+        self.total_input_tokens += input_tokens
+        self.total_output_tokens += output_tokens
+        self.updated_at = time.time()
+
+    def restore_from_checkpoint(self, checkpoint_data: dict) -> None:
+        """Restore session state from checkpoint data."""
+        self.total_cost = checkpoint_data.get("total_cost", 0.0)
+        self.total_tokens = checkpoint_data.get("total_tokens", 0)
+        self.total_input_tokens = checkpoint_data.get("total_input_tokens", 0)
+        self.total_output_tokens = checkpoint_data.get("total_output_tokens", 0)
+        for step, cost in checkpoint_data.get("step_costs", {}).items():
+            self.step_costs[step] = cost
+        for step, tokens in checkpoint_data.get("step_tokens", {}).items():
+            self.step_tokens[step] = tokens
+        for step, tokens in checkpoint_data.get("step_input_tokens", {}).items():
+            self.step_input_tokens[step] = tokens
+        for step, tokens in checkpoint_data.get("step_output_tokens", {}).items():
+            self.step_output_tokens[step] = tokens
+
+    def get_cost_breakdown(self) -> dict:
+        return {
+            "total_cost": self.total_cost,
+            "total_tokens": self.total_tokens,
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "steps": {
+                step: {
+                    "cost": self.step_costs.get(step, 0),
+                    "tokens": self.step_tokens.get(step, 0),
+                    "input_tokens": self.step_input_tokens.get(step, 0),
+                    "output_tokens": self.step_output_tokens.get(step, 0),
+                }
+                for step in self.step_costs
+            }
+        }

openhack/agents/validator.py

@@ -0,0 +1,217 @@
+"""
+Validator agent for confirming vulnerabilities.
+"""
+
+import json
+import logging
+from typing import Optional
+
+from .base import BaseAgent
+from .llm import Message, ToolResult
+from openhack.prompts import VALIDATOR_PROMPT, VALIDATOR_TOOL_INSTRUCTIONS, format_project_context
+
+logger = logging.getLogger(__name__)
+
+
+VALIDATE_FINDING_TOOL = {
+    "name": "validate_finding",
+    "description": "Report the validation result for the potential vulnerability.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "finding_index": {"type": "integer", "description": "Index (1-based) of the finding"},
+            "status": {"type": "string", "enum": ["confirmed", "false_positive", "needs_more_info"]},
+            "confidence": {"type": "string", "enum": ["high", "medium", "low"]},
+            "cvss_score": {"type": "number", "description": "CVSS 3.1 score (0.0 - 10.0)"},
+            "evidence": {"type": "string", "description": "Evidence supporting the validation"},
+            "poc": {"type": "string", "description": "Proof of concept"},
+            "fix": {"type": "string", "description": "Recommended fix"}
+        },
+        "required": ["finding_index", "status", "confidence"]
+    }
+}
+
+FINISH_VALIDATION_TOOL = {
+    "name": "finish_validation",
+    "description": "Call after validating all findings. Signals validation completion.",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "summary": {"type": "string"},
+            "total_confirmed": {"type": "integer"},
+            "total_false_positives": {"type": "integer"}
+        },
+        "required": ["summary", "total_confirmed", "total_false_positives"]
+    }
+}
+
+
+class ValidatorAgent(BaseAgent):
+    name = "validator"
+    description = "Validating and confirming vulnerabilities"
+
+    def __init__(self, *args, original_finding_index=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.validated_findings: list[dict] = []
+        self.false_positives: list[dict] = []
+        self.original_finding_index = original_finding_index
+
+        if original_finding_index is not None:
+            self.name = f"validator:finding_{original_finding_index}"
+            self.description = f"Validating finding {original_finding_index}"
+
+    def get_system_prompt(self, context: dict) -> str:
+        findings = context.get("hunter", {}).get("findings", [])
+        project_context = context.get("project_context", {})
+        project_context_str = format_project_context(project_context)
+
+        findings_text = ""
+        for i, f in enumerate(findings, 1):
+            findings_text += f"""
+### Finding {i}
+- **Category**: {f.get('category', 'Unknown')}
+- **Severity**: {f.get('severity', 'Unknown')}
+- **File**: {f.get('file_path', 'Unknown')}
+- **Line**: {f.get('line_number', 'Unknown')}
+- **Description**: {f.get('description', 'No description')}
+- **Code**:
+```
+{f.get('code_snippet', 'No code snippet')}
+```
+"""
+
+        base_prompt = VALIDATOR_PROMPT.format(
+            findings=findings_text or "No findings to validate",
+            project_context=project_context_str
+        )
+        base_prompt += VALIDATOR_TOOL_INSTRUCTIONS
+        return base_prompt
+
+    def get_tools(self) -> list[dict]:
+        return super().get_tools() + [VALIDATE_FINDING_TOOL, FINISH_VALIDATION_TOOL]
+
+    def _handle_validate_finding(self, args: dict) -> dict:
+        status = args.get("status", "").lower()
+        if self.original_finding_index is not None:
+            original_index = self.original_finding_index
+        else:
+            original_index = args.get("finding_index", 1) - 1
+
+        validation = {
+            "original_index": original_index,
+            "status": status,
+            "confidence": args.get("confidence", "medium").lower(),
+            "cvss_score": args.get("cvss_score"),
+            "evidence": args.get("evidence", ""),
+            "poc": args.get("poc"),
+            "fix": args.get("fix"),
+        }
+
+        if status == "confirmed":
+            self.validated_findings.append(validation)
+        else:
+            self.false_positives.append(validation)
+
+        return {"status": "recorded", "finding_index": args.get("finding_index")}
+
+    def _handle_finish_validation(self, args: dict) -> dict:
+        return {
+            "status": "validation_complete",
+            "confirmed": len(self.validated_findings),
+            "false_positives": len(self.false_positives),
+        }
+
+    async def run(self, task: str, context: Optional[dict] = None) -> dict:
+        context = context or {}
+        self.session.current_agent = self.name
+        self.validated_findings = []
+        self.false_positives = []
+
+        system_prompt = self.get_system_prompt(context)
+        self.messages = [Message(role="user", content=task)]
+        self._seed_existing_instructions()
+
+        max_iterations = 15 if self.original_finding_index is not None else 50
+        iteration = 0
+
+        while iteration < max_iterations:
+            if self.session.cancelled:
+                break
+            iteration += 1
+
+            self._inject_pending_instructions()
+
+            response = await self.llm.chat(
+                messages=self.messages, tools=self.get_tools(), system=system_prompt,
+            )
+
+            self.session.total_cost += response.cost
+            if response.usage:
+                self.session.total_tokens += response.usage.get("total_tokens", 0)
+                self.context_manager.update_usage(response.usage.get("input_tokens", 0))
+
+            if response.content:
+                self.session.add_trace(agent=self.name, event_type="thinking", content=response.content)
+
+            if not response.tool_calls:
+                return self._build_result(response.content or "")
+
+            assistant_msg = Message(
+                role="assistant", content=response.content,
+                tool_calls=[
+                    {"id": tc.id, "type": "function", "function": {"name": tc.name, "arguments": json.dumps(tc.arguments)}}
+                    for tc in response.tool_calls
+                ],
+                reasoning_content=getattr(response, 'reasoning_content', None),
+            )
+            self.messages.append(assistant_msg)
+
+            should_finish = False
+            for tool_call in response.tool_calls:
+                self.session.add_trace(
+                    agent=self.name, event_type="tool_call",
+                    content=f"Calling {tool_call.name}",
+                    tool_name=tool_call.name, tool_input=tool_call.arguments,
+                )
+
+                if tool_call.name == "validate_finding":
+                    result = self._handle_validate_finding(tool_call.arguments)
+                elif tool_call.name == "finish_validation":
+                    result = self._handle_finish_validation(tool_call.arguments)
+                    should_finish = True
+                else:
+                    result = self.tools.execute_tool(tool_call.name, tool_call.arguments)
+
+                self.session.add_trace(
+                    agent=self.name, event_type="tool_result",
+                    content=f"Result from {tool_call.name}",
+                    tool_name=tool_call.name, tool_output=result,
+                )
+
+                raw_content = json.dumps(result) if isinstance(result, dict) else str(result)
+                truncated_content = self.context_manager.truncate_tool_result(tool_call.name, raw_content)
+                tool_result = ToolResult(
+                    tool_call_id=tool_call.id,
+                    content=truncated_content,
+                )
+                self.messages.append(tool_result.to_message())
+
+            if should_finish:
+                return self._build_result(response.content or "")
+
+            if self.context_manager.needs_compaction():
+                self.messages = self.context_manager.compact_messages(self.messages)
+                logger.info(f"[{self.name}] Compacted message history ({self.context_manager.last_input_tokens} input tokens)")
+
+        return self._build_result("Max iterations reached")
+
+    def _build_result(self, summary: str) -> dict:
+        return {
+            "raw_output": summary,
+            "validated_findings": self.validated_findings,
+            "false_positives": self.false_positives,
+            "type": "validation_complete",
+        }
+
+    def _parse_final_response(self, content: str) -> dict:
+        return self._build_result(content)

openhack/agents/validator_swarm.py

@@ -0,0 +1,106 @@
+"""
+Validator swarm agent that spawns one sub-validator per finding.
+"""
+
+import asyncio
+import logging
+from typing import Optional
+
+from .validator import ValidatorAgent
+from .llm import LLMClient
+from .session import Session
+from openhack.tools.registry import ToolRegistry
+from openhack.config import settings
+
+logger = logging.getLogger(__name__)
+
+
+class ValidatorSwarmAgent:
+    name = "validator_swarm"
+    description = "Validator swarm coordinator"
+
+    def __init__(self, llm: LLMClient, tools: ToolRegistry, session: Session):
+        self.llm = llm
+        self.tools = tools
+        self.session = session
+        self.total_cost: float = 0.0
+        self.total_tokens: int = 0
+        self.total_input_tokens: int = 0
+        self.total_output_tokens: int = 0
+
+    def _create_llm_for_sub_validator(self) -> LLMClient:
+        model = settings.validator_model_id or self.llm.model
+        return LLMClient(model=model, temperature=0.0, max_tokens=8192, provider=self.llm.provider, prompt_cache_key=self.llm.prompt_cache_key)
+
+    def _build_sub_context(self, finding: dict, full_context: dict) -> dict:
+        return {
+            "hunter": {"findings": [finding]},
+            "recon": full_context.get("recon", {}),
+            "project_context": full_context.get("project_context", {}),
+        }
+
+    async def run(self, task: str, context: Optional[dict] = None) -> dict:
+        context = context or {}
+        findings = context.get("hunter", {}).get("findings", [])
+
+        if not findings:
+            return {"raw_output": "No findings to validate", "validated_findings": [], "false_positives": [], "type": "validation_complete"}
+
+        self.session.add_trace(agent=self.name, event_type="swarm_start", content={"findings_count": len(findings)})
+
+        sub_validators: list[tuple[int, ValidatorAgent, dict]] = []
+        for idx, finding in enumerate(findings):
+            llm = self._create_llm_for_sub_validator()
+            validator = ValidatorAgent(llm, self.tools, self.session, original_finding_index=idx)
+            sub_context = self._build_sub_context(finding, context)
+            sub_validators.append((idx, validator, sub_context))
+
+        semaphore = asyncio.Semaphore(settings.max_concurrent_validators)
+
+        async def run_sub_validator(finding_idx, validator, sub_context):
+            async with semaphore:
+                try:
+                    sub_task = "Validate this potential vulnerability. Confirm whether it is real, generate a PoC, and suggest a fix."
+                    result = await validator.run(sub_task, sub_context)
+                    return finding_idx, result
+                except Exception as e:
+                    logger.error(f"Sub-validator for finding {finding_idx} failed: {e}")
+                    return finding_idx, {"validated_findings": [], "false_positives": [], "type": "validation_failed"}
+
+        tasks = [
+            asyncio.create_task(run_sub_validator(idx, validator, sub_ctx))
+            for idx, validator, sub_ctx in sub_validators
+        ]
+        try:
+            results = await asyncio.gather(*tasks)
+        except asyncio.CancelledError:
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+
+        all_validated: list[dict] = []
+        all_false_positives: list[dict] = []
+
+        for finding_idx, result in results:
+            all_validated.extend(result.get("validated_findings", []))
+            all_false_positives.extend(result.get("false_positives", []))
+
+        for _, validator, _ in sub_validators:
+            self.total_cost += validator.llm.total_cost
+            self.total_tokens += validator.llm.total_tokens
+            self.total_input_tokens += validator.llm.total_input_tokens
+            self.total_output_tokens += validator.llm.total_output_tokens
+
+        self.session.add_trace(
+            agent=self.name, event_type="swarm_complete",
+            content={"total_confirmed": len(all_validated), "total_false_positives": len(all_false_positives),
+                     "total_cost": self.total_cost, "total_tokens": self.total_tokens},
+        )
+
+        return {
+            "raw_output": f"Validated {len(findings)} findings: {len(all_validated)} confirmed, {len(all_false_positives)} false positives",
+            "validated_findings": all_validated,
+            "false_positives": all_false_positives,
+            "type": "validation_complete",
+        }