From 50dde071f8b4a1d0fde521bfd2fa526e81431003 Mon Sep 17 00:00:00 2001 From: pikpikcu <60111811+pikpikcu@users.noreply.github.com> Date: Wed, 22 Apr 2026 23:45:24 -0400 Subject: [PATCH 1/2] Repeat dataset_search hint every 4 iterations instead of once per phase The previous one-shot injection per phase was too easy for the LLM to miss. Add _dataset_search_called_recently() to check the last 10 tool results, and _dataset_last_hint_iter to track when the hint last fired. The hint now re-injects every _DATASET_HINT_REPEAT_EVERY=4 iterations during ANALYSIS and EXPLOIT as long as dataset_search has not been called recently. Guidance wording changed from advisory to REQUIRED to improve LLM compliance. --- airecon/proxy/agent/loop_exploration.py | 73 ++++++++++++++++--------- 1 file changed, 48 insertions(+), 25 deletions(-) diff --git a/airecon/proxy/agent/loop_exploration.py b/airecon/proxy/agent/loop_exploration.py index e426afd..0b6c9df 100644 --- a/airecon/proxy/agent/loop_exploration.py +++ b/airecon/proxy/agent/loop_exploration.py @@ -324,23 +324,35 @@ def _build_exploration_directive(self, phase: PipelinePhase) -> str: # ── Dataset knowledge base phase hint ────────────────────────────── + # Re-inject the dataset hint every N iterations if the LLM has not called + # dataset_search recently. A single injection per phase is not enough — + # small models routinely skip a hint that appears only once. + _DATASET_HINT_REPEAT_EVERY: int = 4 + + def _dataset_search_called_recently(self, window: int = 10) -> bool: + """Return True if dataset_search appears in the last `window` tool results.""" + seen = 0 + for msg in reversed(self.state.conversation): # type: ignore[attr-defined] + role = msg.get("role", "") + if role == "tool": + if msg.get("name") == "dataset_search": + return True + seen += 1 + if seen >= window: + break + return False + def _build_dataset_hint(self, phase: PipelinePhase) -> str: - """Return a one-time dataset_search suggestion when phase changes. + """Return a dataset_search reminder when the LLM has not used it recently. - Fires once per phase transition (and at iteration 1). Returns "" when: + Fires on every phase transition AND every _DATASET_HINT_REPEAT_EVERY + iterations while dataset_search has not been called in the last 10 tool + results. Silenced when: - No datasets installed at ~/.airecon/datasets/ - - Already injected for this phase - - Phase is REPORT (no lookup needed) + - Phase is RECON or REPORT + - dataset_search was already called recently """ - phase_key = phase.value - - # Only inject once per phase - if getattr(self, "_dataset_hint_injected_phase", "") == phase_key: - return "" - - # Only useful in ANALYSIS and EXPLOIT — RECON dataset results - # are too generic (LLM already knows tool usage) and FTS5 keyword - # matching returns poor results for broad recon methodology queries. + # Only useful in ANALYSIS and EXPLOIT if phase not in (PipelinePhase.ANALYSIS, PipelinePhase.EXPLOIT): return "" @@ -353,7 +365,20 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: if not has_db: return "" - # Build context-aware example queries from session state + phase_key = phase.value + current_iter: int = getattr(self.state, "iteration", 1) # type: ignore[attr-defined] + last_hint_iter: int = getattr(self, "_dataset_last_hint_iter", -self._DATASET_HINT_REPEAT_EVERY) + is_new_phase = getattr(self, "_dataset_hint_injected_phase", "") != phase_key + + # Suppress if: not a new phase AND not enough iterations since last hint + if not is_new_phase and (current_iter - last_hint_iter) < self._DATASET_HINT_REPEAT_EVERY: + return "" + + # Suppress if the LLM already called dataset_search recently + if self._dataset_search_called_recently(): + return "" + + # ── Build context-aware example queries ─────────────────────────── session_techs: list[str] = [] if self._session: # type: ignore[attr-defined] techs = getattr(self._session, "technologies", {}) or {} @@ -363,7 +388,6 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: tech_str = session_techs[0] if session_techs else "" if phase == PipelinePhase.ANALYSIS: - # Short FTS5-friendly queries — 2-3 tokens for best recall if tech_str and tested_vulns: example = f'{{"query": "{tech_str} {tested_vulns[0]} exploit", "limit": 3}}' elif tech_str: @@ -374,12 +398,10 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: else: example = '{"query": "SSRF bypass payload", "category": "bug-bounty", "limit": 3}' guidance = ( - "Query the knowledge base for specific vulnerability techniques, " - "bypass methods, and payloads for detected technologies and injection points." + "REQUIRED: call dataset_search NOW to retrieve vulnerability techniques, " + "bypass methods, and payloads for detected technologies before continuing analysis." ) - else: # EXPLOIT - # CVE queries and payload bypasses — highest value use cases if tested_vulns: vuln = tested_vulns[0].replace("_", " ") example = f'{{"query": "{vuln} exploit payload", "category": "pentest", "limit": 3}}' @@ -388,19 +410,20 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: else: example = '{"query": "WAF bypass exploit payload", "category": "pentest", "limit": 3}' guidance = ( - "Query the knowledge base for CVE exploitation details, PoC techniques, " - "WAF bypass payloads, and nuclei template structure before crafting exploits." + "REQUIRED: call dataset_search NOW for CVE exploitation details, PoC techniques, " + "and WAF bypass payloads before crafting exploits." ) - # Mark injected for this phase + # Update tracking state self._dataset_hint_injected_phase = phase_key # type: ignore[attr-defined] + self._dataset_last_hint_iter = current_iter # type: ignore[attr-defined] return ( - f"[SYSTEM: KNOWLEDGE BASE — Phase={phase_key}]\n" + f"[SYSTEM: KNOWLEDGE BASE — Phase={phase_key} iter={current_iter}]\n" f"{guidance}\n" - f"Example: `dataset_search: {example}`\n" + f"Example query: {example}\n" "Keep queries short (2-3 specific tokens) for best FTS5 recall. " - "Covers CVEs, exploit techniques, payloads, WAF bypasses, and nuclei templates." + "Covers CVEs, exploit techniques, payloads, WAF bypasses, nuclei templates." ) # ── Web search RECON phase hint ──────────────────────────────────── From 1c69d94bb3c41c02bc22b04d9ec6e044f9825600 Mon Sep 17 00:00:00 2001 From: pikpikcu <60111811+pikpikcu@users.noreply.github.com> Date: Wed, 22 Apr 2026 23:54:50 -0400 Subject: [PATCH 2/2] Include RECON phase in dataset_search hints RECON phase now triggers dataset_search suggestions for recon-specific knowledge bases (recon-playbook, network-recon, api-security, web-vuln-patterns). Previously only ANALYSIS and EXPLOIT phases received hints, leaving recon DBs unused during reconnaissance. Also suppresses hints when dataset_search was called recently (last 10 tool results) to avoid redundant prompting. --- airecon/proxy/agent/loop_exploration.py | 41 +++++++++++++++++++++---- 1 file changed, 35 insertions(+), 6 deletions(-) diff --git a/airecon/proxy/agent/loop_exploration.py b/airecon/proxy/agent/loop_exploration.py index 0b6c9df..cdb3d58 100644 --- a/airecon/proxy/agent/loop_exploration.py +++ b/airecon/proxy/agent/loop_exploration.py @@ -349,11 +349,16 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: iterations while dataset_search has not been called in the last 10 tool results. Silenced when: - No datasets installed at ~/.airecon/datasets/ - - Phase is RECON or REPORT + - Phase is REPORT - dataset_search was already called recently + + Covered phases and their installed categories: + RECON → 'recon' (playbooks), 'network', 'api', 'web' + ANALYSIS → 'pentest', 'bug-bounty', 'web', 'api' + EXPLOIT → 'pentest', 'vulnerability', 'bug-bounty' """ - # Only useful in ANALYSIS and EXPLOIT - if phase not in (PipelinePhase.ANALYSIS, PipelinePhase.EXPLOIT): + # REPORT phase needs no knowledge-base lookup + if phase == PipelinePhase.REPORT: return "" # Check datasets actually installed @@ -378,7 +383,7 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: if self._dataset_search_called_recently(): return "" - # ── Build context-aware example queries ─────────────────────────── + # ── Collect session context ─────────────────────────────────────── session_techs: list[str] = [] if self._session: # type: ignore[attr-defined] techs = getattr(self._session, "technologies", {}) or {} @@ -387,7 +392,30 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: tested_vulns = sorted(self._get_tested_vuln_classes())[:3] tech_str = session_techs[0] if session_techs else "" - if phase == PipelinePhase.ANALYSIS: + # ── Detect target type from session ────────────────────────────── + target: str = "" + if self._session: # type: ignore[attr-defined] + target = str(getattr(self._session, "target", "") or "") + is_api_target = any(x in target.lower() for x in ("/api/", "/rest/", "/graphql", "/v1/", "/v2/", "/v3/")) + + # ── Build phase-specific examples and guidance ──────────────────── + if phase == PipelinePhase.RECON: + # airecon-recon-playbook.db (cat: recon), airecon-network-recon.db (cat: network), + # airecon-api-security.db (cat: api), airecon-web-vuln-patterns.db (cat: web) + if tech_str: + example = f'{{"query": "{tech_str} enumeration attack surface", "category": "recon", "limit": 3}}' + elif is_api_target: + example = '{"query": "API endpoint discovery authentication bypass", "category": "api", "limit": 3}' + else: + example = '{"query": "subdomain enumeration port scan fingerprint", "category": "network", "limit": 3}' + guidance = ( + "REQUIRED: call dataset_search NOW to get recon methodology and tool sequences " + "from installed playbooks (categories: recon, network, api, web). " + "The knowledge base contains AIRecon-specific recon order, nmap sequences, " + "DNS enumeration, service fingerprinting, and API discovery playbooks." + ) + + elif phase == PipelinePhase.ANALYSIS: if tech_str and tested_vulns: example = f'{{"query": "{tech_str} {tested_vulns[0]} exploit", "limit": 3}}' elif tech_str: @@ -401,6 +429,7 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: "REQUIRED: call dataset_search NOW to retrieve vulnerability techniques, " "bypass methods, and payloads for detected technologies before continuing analysis." ) + else: # EXPLOIT if tested_vulns: vuln = tested_vulns[0].replace("_", " ") @@ -423,7 +452,7 @@ def _build_dataset_hint(self, phase: PipelinePhase) -> str: f"{guidance}\n" f"Example query: {example}\n" "Keep queries short (2-3 specific tokens) for best FTS5 recall. " - "Covers CVEs, exploit techniques, payloads, WAF bypasses, nuclei templates." + "Covers recon playbooks, CVEs, exploit techniques, payloads, WAF bypasses, nuclei templates." ) # ── Web search RECON phase hint ────────────────────────────────────