fix: redundant idor analysis

Author: fzn0x

watchtower/agents/logic_analysis.py

@@ -1,52 +1,128 @@
 import json
 from pydantic import BaseModel, Field
-from typing import List
+from typing import List, Optional
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.output_parsers import PydanticOutputParser
 from watchtower.core.state import AgentState
 from watchtower.agents.planner import get_llm
 
+
+# ── Schema ────────────────────────────────────────────────────────────────────
+
 class IDORReasoningOutput(BaseModel):
-    confidence_score: int = Field(description="Confidence score from 1-10 that IDOR vulnerabilities exist here.", ge=1, le=10)
-    potential_business_impact: str = Field(description="The potential business impact if these endpoints contain IDOR vulnerabilities.")
-    testing_plan: List[str] = Field(description="A comprehensive step-by-step token/UserID swapping testing plan.")
+    confidence_score: int = Field(
+        description="Confidence score from 1-10 that IDOR vulnerabilities exist here.",
+        ge=1,
+        le=10,
+    )
+    potential_business_impact: str = Field(
+        description="The potential business impact if these endpoints contain IDOR vulnerabilities."
+    )
+    testing_plan: List[str] = Field(
+        description="A comprehensive step-by-step token/UserID swapping testing plan."
+    )
+    affected_endpoints: Optional[List[str]] = Field(
+        default=None,
+        description="Specific endpoints or parameters flagged as high-risk for IDOR."
+    )
+
+
+# ── Constants ─────────────────────────────────────────────────────────────────
+
+RELEVANT_TOOLS = {"httpx", "kiterunner", "arjun", "gobuster", "ffuf"}
+
+SYSTEM_PROMPT = """\
+You are a Senior Security Researcher specialising in Insecure Direct Object \
+Reference (IDOR) and business logic flaws.
+
+Your responsibilities:
+- Analyse discovered endpoints and parameters for IDOR risk.
+- Assess likelihood of token/UserID swapping vulnerabilities.
+- Produce a concrete, step-by-step testing methodology.
+- Identify specific high-risk endpoints when visible in the data.
+
+Rules:
+- Base every conclusion strictly on the supplied recon data.
+- Never fabricate endpoints or parameters not present in the input.
+- Output ONLY valid JSON matching the requested schema — no prose, no markdown fences.\
+"""
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _build_recon_summary(observations: list) -> str:
+    """Concatenate output from relevant recon tools into a single block."""
+    lines = []
+    for obs in observations:
+        if obs.get("tool") in RELEVANT_TOOLS:
+            lines.append(
+                f"[Tool: {obs['tool']}]\n{obs.get('output', '(no output)')}"
+            )
+    return "\n\n".join(lines) if lines else ""
+
+
+def _build_finding(result: IDORReasoningOutput) -> dict:
+    plan_text = "\n".join(f"  {i+1}. {step}" for i, step in enumerate(result.testing_plan))
+    endpoints_text = (
+        "\n".join(f"  - {ep}" for ep in result.affected_endpoints)
+        if result.affected_endpoints
+        else "  None flagged explicitly."
+    )
+
+    return {
+        "title": "IDOR Logic Analysis Plan",
+        "severity": "Info",
+        "confidence": result.confidence_score,
+        "description": (
+            f"Confidence: {result.confidence_score}/10\n\n"
+            f"Business Impact:\n  {result.potential_business_impact}\n\n"
+            f"High-Risk Endpoints:\n{endpoints_text}\n\n"
+            f"Testing Plan:\n{plan_text}"
+        ),
+        "evidence": "Generated by Logic Analysis Agent",
+    }
+
+
+# ── Node ──────────────────────────────────────────────────────────────────────
 
 def logic_analysis_node(state: AgentState) -> dict:
-    observations = state.get("observations", [])
+    observations: list = state.get("observations", [])
+
     if not observations:
         return {"messages": ["No observations found for logic analysis."]}
-        
-    # Gather output from recent tools, especially web recon tools like httpx, kiterunner, arjun
-    recon_data = "\n".join(
-        f"Tool: {obs.get('tool')}\nOutput: {obs.get('output')}"
-        for obs in observations if obs.get("tool") in ["httpx", "kiterunner", "arjun", "gobuster", "ffuf"]
-    )
-    
+
+    latest_tool = observations[-1].get("tool")
+    if latest_tool not in RELEVANT_TOOLS:
+        return {
+            "messages": [f"Skipping logic analysis — last tool '{latest_tool}' is not a recon trigger."]
+        }
+
+    recon_summary = _build_recon_summary(observations)
+    if not recon_summary:
+        return {"messages": ["Relevant recon tools ran but produced no usable output."]}
+
     parser = PydanticOutputParser(pydantic_object=IDORReasoningOutput)
-    
-    prompt = f"""
-You are a Senior Security Researcher specializing in Insecure Direct Object Reference (IDOR) and business logic flaws.
-Review the following web recon endpoints and parameters discovered on the target.
 
-Web Recon Data:
-{recon_data}
+    user_prompt = (
+        "Analyse the following web recon data for IDOR vulnerabilities "
+        "and respond with a JSON object that matches the schema exactly.\n\n"
+        f"--- RECON DATA ---\n{recon_summary}\n--- END RECON DATA ---\n\n"
+        f"{parser.get_format_instructions()}"
+    )
+
+    messages = [
+        SystemMessage(content=SYSTEM_PROMPT),
+        HumanMessage(content=user_prompt),
+    ]
 
-Analyze the data for potential IDOR vulnerabilities. Generate a strict JSON output matching the expected schema.
-Assess the likelihood of IDOR, describe the potential business impact, and provide a detailed step-by-step methodology to test for token and UserID swapping.
-{parser.get_format_instructions()}
-"""
-    
     try:
         llm = get_llm()
         chain = llm | parser
-        result = chain.invoke([HumanMessage(content=prompt)])
-        
-        finding = {
-            "title": "IDOR Logic Analysis Plan",
-            "severity": "Info",
-            "description": f"Confidence: {result.confidence_score}/10\nImpact: {result.potential_business_impact}\nPlan: {', '.join(result.testing_plan)}",
-            "evidence": "Generated by Logic Analysis Agent"
-        }
-        return {"findings": [finding]}
+        result: IDORReasoningOutput = chain.invoke(messages)
+
+        finding = _build_finding(result)
+        existing_findings: list = state.get("findings", [])
+        return {"findings": existing_findings + [finding]}
+
     except Exception as e:
-        return {"messages": [f"Logic analysis error: {e}"]}
+        return {"messages": [f"Logic analysis error: {e}"]}