Skip to content

Pikpikcu/airecon #118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pikpikcu merged 2 commits into from
Apr 23, 2026
Merged

Pikpikcu/airecon #118

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
50dde07
Repeat dataset_search hint every 4 iterations instead of once per phase
pikpikcu Apr 23, 2026
1c69d94
Include RECON phase in dataset_search hints
pikpikcu Apr 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 78 additions & 26 deletions airecon/proxy/agent/loop_exploration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -324,24 +324,41 @@ 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_key = phase.value

# Only inject once per phase
if getattr(self, "_dataset_hint_injected_phase", "") == phase_key:
return ""
- Phase is REPORT
- dataset_search was already called recently

# 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.
if phase not in (PipelinePhase.ANALYSIS, PipelinePhase.EXPLOIT):
Covered phases and their installed categories:
RECON → 'recon' (playbooks), 'network', 'api', 'web'
ANALYSIS → 'pentest', 'bug-bounty', 'web', 'api'
EXPLOIT → 'pentest', 'vulnerability', 'bug-bounty'
"""
# REPORT phase needs no knowledge-base lookup
if phase == PipelinePhase.REPORT:
return ""

# Check datasets actually installed
Expand All @@ -353,7 +370,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 ""

# ── Collect session context ───────────────────────────────────────
session_techs: list[str] = []
if self._session: # type: ignore[attr-defined]
techs = getattr(self._session, "technologies", {}) or {}
Expand All @@ -362,8 +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:
# Short FTS5-friendly queries — 2-3 tokens for best recall
# ── 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:
Expand All @@ -374,12 +426,11 @@ 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}}'
Expand All @@ -388,19 +439,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 recon playbooks, CVEs, exploit techniques, payloads, WAF bypasses, nuclei templates."
)

# ── Web search RECON phase hint ────────────────────────────────────
Expand Down
Loading