shadowcat 2.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (95) hide show
  1. agent/__init__.py +17 -0
  2. agent/benchmark/__init__.py +11 -0
  3. agent/benchmark/cli.py +179 -0
  4. agent/benchmark/config.py +15 -0
  5. agent/benchmark/docker.py +192 -0
  6. agent/benchmark/registry.py +99 -0
  7. agent/core/__init__.py +0 -0
  8. agent/core/agent.py +362 -0
  9. agent/core/backend.py +1667 -0
  10. agent/core/config.py +106 -0
  11. agent/core/controller.py +638 -0
  12. agent/core/events.py +177 -0
  13. agent/core/langfuse.py +320 -0
  14. agent/core/phantom.py +2327 -0
  15. agent/core/planner.py +493 -0
  16. agent/core/profiling.py +58 -0
  17. agent/core/sanitizer.py +104 -0
  18. agent/core/session.py +228 -0
  19. agent/core/tracer.py +137 -0
  20. agent/interface/__init__.py +0 -0
  21. agent/interface/components/__init__.py +0 -0
  22. agent/interface/components/activity_feed.py +202 -0
  23. agent/interface/components/renderers.py +149 -0
  24. agent/interface/components/splash.py +112 -0
  25. agent/interface/main.py +1126 -0
  26. agent/interface/styles.tcss +421 -0
  27. agent/interface/tui.py +508 -0
  28. agent/parsing/html_distiller.py +230 -0
  29. agent/parsing/tool_parser.py +115 -0
  30. agent/prompts/__init__.py +0 -0
  31. agent/prompts/pentesting.py +238 -0
  32. agent/rag_module/knowledge_base.py +116 -0
  33. agent/tests/benchmark_phantom.py +455 -0
  34. agent/tools/__init__.py +14 -0
  35. agent/tools/base.py +99 -0
  36. agent/tools/executor.py +345 -0
  37. agent/tools/registry.py +47 -0
  38. backend/README.md +244 -0
  39. backend/__init__.py +10 -0
  40. backend/api/__init__.py +1 -0
  41. backend/api/routes_scan.py +932 -0
  42. backend/authz.py +75 -0
  43. backend/cli.py +261 -0
  44. backend/compliance/__init__.py +1 -0
  45. backend/compliance/pdpa_mapping.py +16 -0
  46. backend/core/__init__.py +1 -0
  47. backend/core/coverage.py +93 -0
  48. backend/core/events.py +166 -0
  49. backend/core/llm_client.py +614 -0
  50. backend/core/orchestrator.py +402 -0
  51. backend/core/scan_memory.py +236 -0
  52. backend/crawler/__init__.py +1 -0
  53. backend/crawler/csrf.py +75 -0
  54. backend/crawler/headless.py +232 -0
  55. backend/crawler/js_analyzer.py +133 -0
  56. backend/crawler/spider.py +550 -0
  57. backend/crawler/subdomains.py +279 -0
  58. backend/daemon.py +252 -0
  59. backend/db/README.md +86 -0
  60. backend/db/__init__.py +17 -0
  61. backend/db/engine.py +87 -0
  62. backend/db/models.py +206 -0
  63. backend/db/repository.py +316 -0
  64. backend/db/schema.sql +247 -0
  65. backend/modes/__init__.py +5 -0
  66. backend/modes/base.py +178 -0
  67. backend/modes/ctf.py +39 -0
  68. backend/modes/enterprise.py +210 -0
  69. backend/modes/general.py +226 -0
  70. backend/modes/registry.py +34 -0
  71. backend/reporting/__init__.py +0 -0
  72. backend/reporting/generator.py +285 -0
  73. backend/schemas/__init__.py +1 -0
  74. backend/schemas/api.py +423 -0
  75. backend/tools/__init__.py +62 -0
  76. backend/tools/dirbrute_tool.py +304 -0
  77. backend/tools/general_report_tool.py +135 -0
  78. backend/tools/http_tool.py +351 -0
  79. backend/tools/nuclei_tool.py +20 -0
  80. backend/tools/report_tool.py +145 -0
  81. backend/tools/shell_tools.py +23 -0
  82. backend/verification/__init__.py +1 -0
  83. backend/verification/evidence_store.py +125 -0
  84. backend/verification/general_oracle.py +369 -0
  85. backend/verification/idor_oracle.py +131 -0
  86. backend/waf/__init__.py +0 -0
  87. backend/waf/detector.py +147 -0
  88. backend/waf/evasion.py +117 -0
  89. backend/webui/index.html +713 -0
  90. backend/workspace.py +182 -0
  91. shadowcat-2.0.0.dist-info/METADATA +360 -0
  92. shadowcat-2.0.0.dist-info/RECORD +95 -0
  93. shadowcat-2.0.0.dist-info/WHEEL +4 -0
  94. shadowcat-2.0.0.dist-info/entry_points.txt +4 -0
  95. shadowcat-2.0.0.dist-info/licenses/LICENSE.md +21 -0
agent/core/phantom.py ADDED
@@ -0,0 +1,2327 @@
1
+ """PHANTOM — Probabilistic Hallucination-Anchored Navigation.
2
+
3
+ Stage 1: Fingerprint Crystallization.
4
+
5
+ The goal of this stage is to collapse a noisy set of HTTP signals (headers,
6
+ cookies, body snippets from a handful of well-known paths) into a single
7
+ compact ``FingerprintVector`` that downstream PHANTOM stages can reason over
8
+ without re-issuing the underlying requests. One concurrent burst of recon +
9
+ one cheap LLM call replaces dozens of agent-driven probes.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import asyncio
15
+ import json
16
+ import re
17
+ from collections.abc import Awaitable, Callable
18
+ from typing import Any
19
+ from urllib.parse import urljoin, urlparse
20
+
21
+ import httpx
22
+ from pydantic import BaseModel, Field, ValidationError, field_validator
23
+
24
+ from agent.core.profiling import span
25
+
26
+ LLMCallable = Callable[[str], Awaitable[str]]
27
+
28
+ # Same async callable shape, but with an optional ``json_schema`` kwarg the
29
+ # integrator can plumb into the provider's strict structured-output mode
30
+ # (OpenAI ``response_format={"type":"json_schema",...}``, Anthropic tool-use
31
+ # forcing, vLLM/llama.cpp guided JSON). Used by :func:`generate_mam` to
32
+ # constrain MAM output against :data:`COMPACT_MAM_JSON_SCHEMA`. If the
33
+ # supplied callable doesn't accept the kwarg, :func:`generate_mam`
34
+ # transparently falls back to a plain ``(prompt,) -> str`` call — so older
35
+ # test stubs and integrators that haven't migrated still work.
36
+ StructuredLLMCallable = Callable[..., Awaitable[str]]
37
+
38
+ DEFAULT_RECON_PATHS: tuple[str, ...] = (
39
+ "/",
40
+ "/robots.txt",
41
+ "/favicon.ico",
42
+ "/.well-known/security.txt",
43
+ "/sitemap.xml",
44
+ )
45
+
46
+ _HEADERS_OF_INTEREST: tuple[str, ...] = (
47
+ "server",
48
+ "x-powered-by",
49
+ "x-aspnet-version",
50
+ "x-aspnetmvc-version",
51
+ "x-generator",
52
+ "x-drupal-cache",
53
+ "x-runtime",
54
+ "via",
55
+ "set-cookie",
56
+ "content-type",
57
+ "www-authenticate",
58
+ )
59
+
60
+ _BODY_SNIPPET_CHARS = 2048
61
+ _PER_REQUEST_TIMEOUT_SECONDS = 3.0
62
+ _FAVICON_HASH_BYTES = 4096
63
+
64
+
65
+ class FingerprintVector(BaseModel):
66
+ """Crystallized view of a target's tech stack.
67
+
68
+ Every field is a free-form short string (the LLM picks the vocabulary) so
69
+ that downstream stages can do cheap string matching without a fixed
70
+ taxonomy. ``overall_confidence`` is the model's self-reported credence in
71
+ the *combination* of fields, not an average — a single high-confidence
72
+ field paired with several unknowns should still pull this down.
73
+ """
74
+
75
+ framework: str = Field(
76
+ description="Server-side framework (e.g. 'Django', 'Express', 'ASP.NET', 'unknown')."
77
+ )
78
+ version_guess: str = Field(
79
+ description="Best-effort version string for the framework. 'unknown' if no signal."
80
+ )
81
+ frontend: str = Field(
82
+ description="Frontend stack (e.g. 'React', 'Vue', 'server-rendered', 'unknown')."
83
+ )
84
+ auth_pattern: str = Field(
85
+ description="Auth scheme hint (e.g. 'session-cookie', 'JWT', 'Basic', 'none-observed')."
86
+ )
87
+ likely_orm: str = Field(
88
+ description="ORM / persistence hint (e.g. 'Django ORM', 'ActiveRecord', 'Prisma', 'unknown')."
89
+ )
90
+ overall_confidence: float = Field(
91
+ ge=0.0,
92
+ le=1.0,
93
+ description="Self-reported confidence in the overall fingerprint, in [0.0, 1.0].",
94
+ )
95
+
96
+ @field_validator("framework", "version_guess", "frontend", "auth_pattern", "likely_orm")
97
+ @classmethod
98
+ def _non_empty(cls, v: str) -> str:
99
+ # Empty strings would silently degrade downstream string matching; force
100
+ # the LLM (and any test stub) to commit to at least 'unknown'.
101
+ if not v or not v.strip():
102
+ return "unknown"
103
+ return v.strip()
104
+
105
+
106
+ class PathSignal(BaseModel):
107
+ """Single-path result inside the raw signal dict.
108
+
109
+ Kept as a model rather than a free-form dict so the crystallization prompt
110
+ has a stable shape regardless of which paths timed out.
111
+ """
112
+
113
+ path: str
114
+ status: int | None = None
115
+ headers: dict[str, str] = Field(default_factory=dict)
116
+ set_cookie: list[str] = Field(default_factory=list)
117
+ body_snippet: str = ""
118
+ body_length: int = 0
119
+ error: str | None = None
120
+
121
+
122
+ async def gather_raw_signals(
123
+ target_url: str,
124
+ *,
125
+ paths: tuple[str, ...] = DEFAULT_RECON_PATHS,
126
+ timeout_seconds: float = _PER_REQUEST_TIMEOUT_SECONDS,
127
+ client: httpx.AsyncClient | None = None,
128
+ ) -> dict[str, Any]:
129
+ """Fetch ``paths`` concurrently from ``target_url`` and return raw signals.
130
+
131
+ Each path is fetched in parallel via :func:`asyncio.gather`. A failure on
132
+ one path (timeout, connection refused, non-2xx, etc.) never blocks the
133
+ others — the failed entry simply carries an ``error`` field.
134
+
135
+ The optional ``client`` parameter lets callers reuse a long-lived
136
+ :class:`httpx.AsyncClient` (and its connection pool). When ``None``, a
137
+ one-shot client is created and closed before return.
138
+ """
139
+ if not target_url:
140
+ raise ValueError("target_url must be a non-empty URL")
141
+
142
+ parsed = urlparse(target_url)
143
+ if not parsed.scheme or not parsed.netloc:
144
+ raise ValueError(f"target_url must include scheme and host (got {target_url!r})")
145
+ base = f"{parsed.scheme}://{parsed.netloc}"
146
+
147
+ owns_client = False
148
+ if client is not None:
149
+ http = client
150
+ else:
151
+ from agent.core.backend import get_shared_http_client
152
+
153
+ http = get_shared_http_client()
154
+
155
+ try:
156
+ with span(f"phantom.gather_raw_signals [{len(paths)} paths, base={base}]"):
157
+ coros = [_fetch_one(http, base, p, timeout_seconds) for p in paths]
158
+ results = await asyncio.gather(*coros, return_exceptions=False)
159
+ finally:
160
+ if owns_client:
161
+ await http.aclose()
162
+
163
+ return {
164
+ "target_url": target_url,
165
+ "base": base,
166
+ "paths": [r.model_dump() for r in results],
167
+ }
168
+
169
+
170
+ async def _fetch_one(
171
+ http: httpx.AsyncClient,
172
+ base: str,
173
+ path: str,
174
+ timeout_seconds: float,
175
+ ) -> PathSignal:
176
+ """Single-path fetch with hard timeout and graceful failure."""
177
+ url = urljoin(base + "/", path.lstrip("/"))
178
+ try:
179
+ resp = await asyncio.wait_for(
180
+ http.get(url),
181
+ timeout=timeout_seconds,
182
+ )
183
+ except TimeoutError:
184
+ return PathSignal(path=path, error="timeout")
185
+ except httpx.TimeoutException:
186
+ return PathSignal(path=path, error="timeout")
187
+ except httpx.RequestError as e:
188
+ return PathSignal(path=path, error=f"{e.__class__.__name__}: {e}")
189
+
190
+ headers = {k: v for k, v in resp.headers.items() if k.lower() in _HEADERS_OF_INTEREST}
191
+ set_cookie = (
192
+ list(resp.headers.get_list("set-cookie")) if hasattr(resp.headers, "get_list") else []
193
+ )
194
+
195
+ body_snippet = _safe_text_snippet(resp, path)
196
+
197
+ return PathSignal(
198
+ path=path,
199
+ status=resp.status_code,
200
+ headers=headers,
201
+ set_cookie=set_cookie,
202
+ body_snippet=body_snippet,
203
+ body_length=len(resp.content),
204
+ )
205
+
206
+
207
+ def _safe_text_snippet(resp: httpx.Response, path: str) -> str:
208
+ """Return a short textual snippet suited to the path.
209
+
210
+ Binary paths (favicon.ico) get a hex digest of the leading bytes — useful
211
+ for favicon-hash style fingerprinting later — rather than mojibake text.
212
+ """
213
+ ctype = (resp.headers.get("content-type") or "").lower()
214
+ if path.endswith(".ico") or ("image" in ctype) or ("octet-stream" in ctype):
215
+ head = resp.content[:_FAVICON_HASH_BYTES]
216
+ return f"<binary {len(resp.content)}B, leading-sha={_short_hex(head)}>"
217
+ text = resp.text
218
+ if len(text) > _BODY_SNIPPET_CHARS:
219
+ return text[:_BODY_SNIPPET_CHARS] + "…[truncated]"
220
+ return text
221
+
222
+
223
+ def _short_hex(data: bytes) -> str:
224
+ import hashlib
225
+
226
+ return hashlib.sha256(data).hexdigest()[:16]
227
+
228
+
229
+ # ---------------------------------------------------------------------------
230
+ # Crystallization: raw signals -> FingerprintVector via a fast/cheap LLM
231
+ # ---------------------------------------------------------------------------
232
+
233
+
234
+ _CRYSTALLIZATION_SYSTEM_RULES = """\
235
+ You are a web-stack fingerprinting assistant. You will receive raw HTTP recon
236
+ signals (headers, Set-Cookie values, and short body snippets) from a small
237
+ set of well-known paths on a single target.
238
+
239
+ Your only job is to output a STRICT JSON object — no prose, no markdown
240
+ fences, no commentary — that matches this schema exactly:
241
+
242
+ {
243
+ "framework": "<server-side framework or 'unknown'>",
244
+ "version_guess": "<version string or 'unknown'>",
245
+ "frontend": "<frontend stack or 'unknown'>",
246
+ "auth_pattern": "<auth hint or 'none-observed'>",
247
+ "likely_orm": "<ORM hint or 'unknown'>",
248
+ "overall_confidence": <float in [0.0, 1.0]>
249
+ }
250
+
251
+ Rules:
252
+ - Every string field MUST be present. Use 'unknown' when there is no signal.
253
+ - overall_confidence reflects credence in the WHOLE fingerprint, not any
254
+ single field. If most fields are 'unknown', confidence must be low.
255
+ - Do not invent versions. If no version string is visible, return 'unknown'.
256
+ - Do not output anything outside the JSON object.
257
+ """
258
+
259
+
260
+ def build_crystallization_prompt(raw_signals: dict[str, Any]) -> str:
261
+ """Construct the user-message prompt for the crystallization LLM call.
262
+
263
+ The prompt embeds ``raw_signals`` as compact JSON and re-states the
264
+ schema. Kept pure (no I/O) so it can be unit-tested deterministically and
265
+ snapshot-compared.
266
+ """
267
+ serialized = json.dumps(raw_signals, separators=(",", ":"), default=str)
268
+ return (
269
+ f"{_CRYSTALLIZATION_SYSTEM_RULES}\n"
270
+ f"--- BEGIN RAW SIGNALS (JSON) ---\n"
271
+ f"{serialized}\n"
272
+ f"--- END RAW SIGNALS ---\n"
273
+ f"Return the JSON object now."
274
+ )
275
+
276
+
277
+ _JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
278
+
279
+
280
+ def parse_fingerprint_response(response_text: str) -> FingerprintVector:
281
+ """Parse the LLM's JSON output into a validated ``FingerprintVector``.
282
+
283
+ Tolerant of common LLM output quirks (leading/trailing prose, ```json
284
+ fences) by extracting the first balanced ``{...}`` block before parsing.
285
+ Raises :class:`ValueError` with the original text included if parsing or
286
+ schema validation fails — the caller can decide whether to retry or fall
287
+ back to a default low-confidence vector.
288
+ """
289
+ if not response_text or not response_text.strip():
290
+ raise ValueError("Empty response from crystallization LLM")
291
+
292
+ candidate = response_text.strip()
293
+ if not candidate.startswith("{"):
294
+ match = _JSON_OBJECT_RE.search(candidate)
295
+ if match is None:
296
+ raise ValueError(f"No JSON object found in LLM response: {response_text!r}")
297
+ candidate = match.group(0)
298
+
299
+ try:
300
+ payload = json.loads(candidate)
301
+ except json.JSONDecodeError as e:
302
+ raise ValueError(
303
+ f"Crystallization LLM returned invalid JSON ({e}): {response_text!r}"
304
+ ) from e
305
+
306
+ try:
307
+ return FingerprintVector.model_validate(payload)
308
+ except ValidationError as e:
309
+ raise ValueError(
310
+ f"Crystallization LLM output failed schema validation: {e}\nRaw payload: {payload!r}"
311
+ ) from e
312
+
313
+
314
+ async def crystallize_fingerprint(
315
+ raw_signals: dict[str, Any],
316
+ llm_call: LLMCallable,
317
+ ) -> FingerprintVector:
318
+ """Turn raw signals into a ``FingerprintVector`` via one LLM call.
319
+
320
+ ``llm_call`` is an injected async callable ``(prompt: str) -> str`` so the
321
+ concrete backend (Haiku via Anthropic, an OpenRouter-routed cheap model,
322
+ or a unit-test stub) is selected at the integration site rather than
323
+ hardwired here. This mirrors the project's existing ``AgentBackend``
324
+ injection pattern.
325
+ """
326
+ prompt = build_crystallization_prompt(raw_signals)
327
+ response_text = await llm_call(prompt)
328
+ return parse_fingerprint_response(response_text)
329
+
330
+
331
+ # ---------------------------------------------------------------------------
332
+ # Stage 2: Code Mirroring — FingerprintVector -> MirroredApplicationManifest
333
+ #
334
+ # Given a crystallized fingerprint, a reasoning LLM is asked to hallucinate
335
+ # the most probable internal architecture: which routes likely exist, which
336
+ # default framework misconfigurations are likely exposed, and which
337
+ # vulnerability classes each predicted endpoint might harbor. The output is
338
+ # a budget the agent can spend on targeted probes, instead of brute-forcing
339
+ # wordlists.
340
+ # ---------------------------------------------------------------------------
341
+
342
+
343
+ class PredictedEndpoint(BaseModel):
344
+ """One hypothesized route on the target.
345
+
346
+ ``probability`` is the LLM's credence that this exact path+method exists
347
+ and returns ``expected_status``; downstream stages use it to rank probe
348
+ order. ``rationale`` must reference the fingerprint signal that justified
349
+ the guess so a human reviewer can sanity-check (or veto) the prediction.
350
+ """
351
+
352
+ path: str = Field(description="Absolute URL path beginning with '/' (e.g. '/api/v1/users').")
353
+ method: str = Field(
354
+ description="HTTP method in upper case (GET, POST, PUT, PATCH, DELETE, etc.)."
355
+ )
356
+ expected_status: int = Field(
357
+ ge=100,
358
+ le=599,
359
+ description="Most likely HTTP status returned by an unauthenticated probe.",
360
+ )
361
+ probability: float = Field(
362
+ ge=0.0,
363
+ le=1.0,
364
+ description="Credence that this endpoint exists as predicted, in [0.0, 1.0].",
365
+ )
366
+ rationale: str = Field(
367
+ description=(
368
+ "Short justification tied to a fingerprint signal or framework convention "
369
+ "(e.g. 'Laravel default debug route', 'Django admin convention')."
370
+ )
371
+ )
372
+ potential_vulnerability_class: str | None = Field(
373
+ default=None,
374
+ description=(
375
+ "Suspected vulnerability class if probed naively (e.g. 'IDOR', 'SSRF', "
376
+ "'info-disclosure', 'auth-bypass'). None if no specific class is implied."
377
+ ),
378
+ )
379
+
380
+ @field_validator("path")
381
+ @classmethod
382
+ def _leading_slash(cls, v: str) -> str:
383
+ v = v.strip()
384
+ if not v:
385
+ raise ValueError("path must be non-empty")
386
+ return v if v.startswith("/") else "/" + v
387
+
388
+ @field_validator("method")
389
+ @classmethod
390
+ def _upper_method(cls, v: str) -> str:
391
+ v = v.strip().upper()
392
+ if not v:
393
+ raise ValueError("method must be non-empty")
394
+ return v
395
+
396
+ @field_validator("rationale")
397
+ @classmethod
398
+ def _non_empty_rationale(cls, v: str) -> str:
399
+ if not v or not v.strip():
400
+ raise ValueError("rationale must be non-empty (every prediction needs justification)")
401
+ return v.strip()
402
+
403
+ @field_validator("potential_vulnerability_class")
404
+ @classmethod
405
+ def _normalize_vuln_class(cls, v: str | None) -> str | None:
406
+ # Models routinely emit "none"/"n/a"/"" instead of a true JSON null;
407
+ # collapse those to None so downstream filters work uniformly.
408
+ if v is None:
409
+ return None
410
+ stripped = v.strip()
411
+ if not stripped or stripped.lower() in {"none", "n/a", "na", "null", "unknown"}:
412
+ return None
413
+ return stripped
414
+
415
+
416
+ class MirroredApplicationManifest(BaseModel):
417
+ """Hallucinated mirror of the target application's internal surface."""
418
+
419
+ predicted_endpoints: list[PredictedEndpoint] = Field(
420
+ default_factory=list,
421
+ description="Endpoints the LLM expects to exist, ranked by `probability`.",
422
+ )
423
+ global_assumptions: list[str] = Field(
424
+ default_factory=list,
425
+ description=(
426
+ "Framework- or deployment-wide assumptions underlying the predictions "
427
+ "(e.g. 'Assumes default Laravel Sanctum configuration', "
428
+ "'Assumes APP_DEBUG=true based on stack trace')."
429
+ ),
430
+ )
431
+
432
+
433
+ # ---------------------------------------------------------------------------
434
+ # Stage 2 (Compact): same MAM contract, ~half the output tokens.
435
+ #
436
+ # The original schema forces a non-empty rationale per endpoint. On a complex
437
+ # target the model emits 60-120 endpoints; rationale + verbose field names
438
+ # cost ~2-3k OUTPUT tokens per MAM call, and output tokens dominate Stage-2
439
+ # latency. The compact wire format below shortens every field name and drops
440
+ # the mandatory rationale. Expansion back into the full PredictedEndpoint
441
+ # shape happens server-side in :func:`expand_compact`, so downstream stages
442
+ # (rank_probes, cheap_sweep, Stage 4) see the same MirroredApplicationManifest
443
+ # type they always did — only the LLM-facing wire format changes.
444
+ # ---------------------------------------------------------------------------
445
+
446
+
447
+ class PredictedEndpointCompact(BaseModel):
448
+ """Compact wire format the LLM emits in :func:`generate_mam`.
449
+
450
+ Field names are single letters because output tokens are the dominant
451
+ cost of Stage 2; ``"p"`` over ``"path"`` saves ~3 tokens per endpoint
452
+ times ~60 endpoints/target = ~180 tokens per call, every call.
453
+ """
454
+
455
+ p: str = Field(description="Absolute path starting with '/'")
456
+ m: str = Field(default="GET", description="HTTP method, upper case")
457
+ s: int = Field(default=200, ge=100, le=599, description="Expected HTTP status code")
458
+ c: float = Field(ge=0.0, le=1.0, description="Confidence/probability in [0.0, 1.0]")
459
+ v: str | None = Field(
460
+ default=None,
461
+ description="Vulnerability class label (kebab/snake) or null",
462
+ )
463
+
464
+ @field_validator("p")
465
+ @classmethod
466
+ def _leading_slash(cls, v: str) -> str:
467
+ v = v.strip()
468
+ if not v:
469
+ raise ValueError("p must be non-empty")
470
+ return v if v.startswith("/") else "/" + v
471
+
472
+ @field_validator("m")
473
+ @classmethod
474
+ def _upper_method(cls, v: str) -> str:
475
+ return (v or "").strip().upper() or "GET"
476
+
477
+ @field_validator("v")
478
+ @classmethod
479
+ def _normalize_vuln(cls, v: str | None) -> str | None:
480
+ # Models routinely emit "none"/"unknown"/"" instead of true JSON null
481
+ # even when the schema declares the field nullable; collapse those to
482
+ # None so downstream severity scoring works uniformly.
483
+ if v is None:
484
+ return None
485
+ stripped = v.strip()
486
+ if not stripped or stripped.lower() in {"none", "n/a", "na", "null", "unknown"}:
487
+ return None
488
+ return stripped
489
+
490
+
491
+ class CompactMirroredManifest(BaseModel):
492
+ """Compact analogue of :class:`MirroredApplicationManifest`."""
493
+
494
+ endpoints: list[PredictedEndpointCompact] = Field(default_factory=list)
495
+ assumptions: list[str] = Field(default_factory=list)
496
+
497
+
498
+ # JSON Schema passed to provider strict-output modes. Hand-written (not
499
+ # Pydantic-derived) so it satisfies OpenAI's strict-mode constraints exactly:
500
+ # every object has ``additionalProperties: false`` and lists every property
501
+ # in ``required``. Models can still emit JSON null for ``v`` thanks to the
502
+ # ``["string","null"]`` union — that's OpenAI's official nullable shape in
503
+ # strict mode.
504
+ COMPACT_MAM_JSON_SCHEMA: dict[str, Any] = {
505
+ "name": "MirroredApplicationManifest",
506
+ "strict": True,
507
+ "schema": {
508
+ "type": "object",
509
+ "additionalProperties": False,
510
+ "required": ["endpoints", "assumptions"],
511
+ "properties": {
512
+ "endpoints": {
513
+ "type": "array",
514
+ "items": {
515
+ "type": "object",
516
+ "additionalProperties": False,
517
+ "required": ["p", "m", "s", "c", "v"],
518
+ "properties": {
519
+ "p": {
520
+ "type": "string",
521
+ "description": "Absolute path starting with /",
522
+ },
523
+ "m": {
524
+ "type": "string",
525
+ "enum": [
526
+ "GET",
527
+ "POST",
528
+ "PUT",
529
+ "PATCH",
530
+ "DELETE",
531
+ "HEAD",
532
+ "OPTIONS",
533
+ ],
534
+ },
535
+ "s": {"type": "integer", "minimum": 100, "maximum": 599},
536
+ "c": {"type": "number", "minimum": 0.0, "maximum": 1.0},
537
+ "v": {
538
+ "type": ["string", "null"],
539
+ "description": (
540
+ "Vulnerability class label "
541
+ "(e.g. 'info-disclosure', 'IDOR', 'SSRF', "
542
+ "'auth-bypass', 'rce') or null."
543
+ ),
544
+ },
545
+ },
546
+ },
547
+ },
548
+ "assumptions": {
549
+ "type": "array",
550
+ "items": {"type": "string"},
551
+ },
552
+ },
553
+ },
554
+ }
555
+
556
+
557
+ def expand_compact(
558
+ compact: CompactMirroredManifest,
559
+ *,
560
+ source: str = "llm-compact",
561
+ ) -> MirroredApplicationManifest:
562
+ """Expand compact wire format into a full :class:`MirroredApplicationManifest`.
563
+
564
+ The original schema requires a non-empty ``rationale`` per endpoint — we
565
+ trade that field for tokens in compact mode, so synthesize a short
566
+ ``source`` placeholder ("llm-compact" / "deterministic") that downstream
567
+ debuggers can grep on without forcing the model to re-generate
568
+ justifications it didn't need to produce.
569
+ """
570
+ endpoints = [
571
+ PredictedEndpoint(
572
+ path=ep.p,
573
+ method=ep.m,
574
+ expected_status=ep.s,
575
+ probability=ep.c,
576
+ rationale=source,
577
+ potential_vulnerability_class=ep.v,
578
+ )
579
+ for ep in compact.endpoints
580
+ ]
581
+ return MirroredApplicationManifest(
582
+ predicted_endpoints=endpoints,
583
+ global_assumptions=compact.assumptions,
584
+ )
585
+
586
+
587
+ # ---------------------------------------------------------------------------
588
+ # Deterministic rulebook
589
+ #
590
+ # Framework conventions are constants — there's no value in burning LLM tokens
591
+ # to re-derive them every call, and a rulebook gives us EXHAUSTIVE coverage
592
+ # of the boring-but-high-value endpoints (the ones CTF designers love to use
593
+ # precisely because they look low-signal to a generation model).
594
+ #
595
+ # ALWAYS_PROBE is the framework-agnostic floor (dotfiles, common admin/API
596
+ # routes, backup leaks). FRAMEWORK_RULES is keyed by a case-insensitive
597
+ # substring of FingerprintVector.framework so a fingerprint of "Spring Boot 2.7"
598
+ # still matches the "spring" pack.
599
+ #
600
+ # Rationale strings start with "static:" so a debugger reading a merged MAM
601
+ # can immediately tell which endpoints came from this table vs the LLM.
602
+ # Probabilities reflect "how often is this *actually* reachable in the wild",
603
+ # NOT severity — severity weighting happens in rank_probes via
604
+ # potential_vulnerability_class.
605
+ # ---------------------------------------------------------------------------
606
+
607
+
608
+ ALWAYS_PROBE: tuple[PredictedEndpoint, ...] = (
609
+ PredictedEndpoint(
610
+ path="/.git/config",
611
+ method="GET",
612
+ expected_status=200,
613
+ probability=0.15,
614
+ rationale="static: VCS leak",
615
+ potential_vulnerability_class="info-disclosure",
616
+ ),
617
+ PredictedEndpoint(
618
+ path="/.git/HEAD",
619
+ method="GET",
620
+ expected_status=200,
621
+ probability=0.15,
622
+ rationale="static: VCS leak",
623
+ potential_vulnerability_class="info-disclosure",
624
+ ),
625
+ PredictedEndpoint(
626
+ path="/.env",
627
+ method="GET",
628
+ expected_status=200,
629
+ probability=0.10,
630
+ rationale="static: env leak",
631
+ potential_vulnerability_class="info-disclosure",
632
+ ),
633
+ PredictedEndpoint(
634
+ path="/.DS_Store",
635
+ method="GET",
636
+ expected_status=200,
637
+ probability=0.05,
638
+ rationale="static: macOS metadata",
639
+ potential_vulnerability_class="info-disclosure",
640
+ ),
641
+ PredictedEndpoint(
642
+ path="/robots.txt",
643
+ method="GET",
644
+ expected_status=200,
645
+ probability=0.70,
646
+ rationale="static: convention",
647
+ potential_vulnerability_class=None,
648
+ ),
649
+ PredictedEndpoint(
650
+ path="/sitemap.xml",
651
+ method="GET",
652
+ expected_status=200,
653
+ probability=0.40,
654
+ rationale="static: convention",
655
+ potential_vulnerability_class=None,
656
+ ),
657
+ PredictedEndpoint(
658
+ path="/.well-known/security.txt",
659
+ method="GET",
660
+ expected_status=200,
661
+ probability=0.10,
662
+ rationale="static: convention",
663
+ potential_vulnerability_class=None,
664
+ ),
665
+ PredictedEndpoint(
666
+ path="/server-status",
667
+ method="GET",
668
+ expected_status=200,
669
+ probability=0.05,
670
+ rationale="static: Apache mod_status",
671
+ potential_vulnerability_class="info-disclosure",
672
+ ),
673
+ PredictedEndpoint(
674
+ path="/server-info",
675
+ method="GET",
676
+ expected_status=200,
677
+ probability=0.05,
678
+ rationale="static: Apache mod_info",
679
+ potential_vulnerability_class="info-disclosure",
680
+ ),
681
+ PredictedEndpoint(
682
+ path="/phpinfo.php",
683
+ method="GET",
684
+ expected_status=200,
685
+ probability=0.05,
686
+ rationale="static: PHP info leak",
687
+ potential_vulnerability_class="info-disclosure",
688
+ ),
689
+ PredictedEndpoint(
690
+ path="/swagger.json",
691
+ method="GET",
692
+ expected_status=200,
693
+ probability=0.15,
694
+ rationale="static: OpenAPI doc",
695
+ potential_vulnerability_class="info-disclosure",
696
+ ),
697
+ PredictedEndpoint(
698
+ path="/swagger-ui.html",
699
+ method="GET",
700
+ expected_status=200,
701
+ probability=0.10,
702
+ rationale="static: Swagger UI",
703
+ potential_vulnerability_class="info-disclosure",
704
+ ),
705
+ PredictedEndpoint(
706
+ path="/openapi.json",
707
+ method="GET",
708
+ expected_status=200,
709
+ probability=0.10,
710
+ rationale="static: OpenAPI doc",
711
+ potential_vulnerability_class="info-disclosure",
712
+ ),
713
+ PredictedEndpoint(
714
+ path="/graphql",
715
+ method="GET",
716
+ expected_status=200,
717
+ probability=0.10,
718
+ rationale="static: GraphQL endpoint",
719
+ potential_vulnerability_class="info-disclosure",
720
+ ),
721
+ PredictedEndpoint(
722
+ path="/admin",
723
+ method="GET",
724
+ expected_status=200,
725
+ probability=0.30,
726
+ rationale="static: admin convention",
727
+ potential_vulnerability_class="auth-bypass",
728
+ ),
729
+ PredictedEndpoint(
730
+ path="/login",
731
+ method="GET",
732
+ expected_status=200,
733
+ probability=0.40,
734
+ rationale="static: auth convention",
735
+ potential_vulnerability_class=None,
736
+ ),
737
+ PredictedEndpoint(
738
+ path="/api",
739
+ method="GET",
740
+ expected_status=200,
741
+ probability=0.40,
742
+ rationale="static: API convention",
743
+ potential_vulnerability_class=None,
744
+ ),
745
+ PredictedEndpoint(
746
+ path="/api/v1",
747
+ method="GET",
748
+ expected_status=200,
749
+ probability=0.30,
750
+ rationale="static: API convention",
751
+ potential_vulnerability_class=None,
752
+ ),
753
+ PredictedEndpoint(
754
+ path="/backup",
755
+ method="GET",
756
+ expected_status=200,
757
+ probability=0.05,
758
+ rationale="static: backup leak",
759
+ potential_vulnerability_class="info-disclosure",
760
+ ),
761
+ PredictedEndpoint(
762
+ path="/backup.zip",
763
+ method="GET",
764
+ expected_status=200,
765
+ probability=0.05,
766
+ rationale="static: backup leak",
767
+ potential_vulnerability_class="info-disclosure",
768
+ ),
769
+ PredictedEndpoint(
770
+ path="/backup.tar.gz",
771
+ method="GET",
772
+ expected_status=200,
773
+ probability=0.03,
774
+ rationale="static: backup leak",
775
+ potential_vulnerability_class="info-disclosure",
776
+ ),
777
+ )
778
+
779
+
780
+ FRAMEWORK_RULES: dict[str, tuple[PredictedEndpoint, ...]] = {
781
+ "django": (
782
+ PredictedEndpoint(
783
+ path="/admin/",
784
+ method="GET",
785
+ expected_status=200,
786
+ probability=0.60,
787
+ rationale="static: Django admin",
788
+ potential_vulnerability_class="auth-bypass",
789
+ ),
790
+ PredictedEndpoint(
791
+ path="/accounts/login/",
792
+ method="GET",
793
+ expected_status=200,
794
+ probability=0.40,
795
+ rationale="static: Django auth",
796
+ potential_vulnerability_class=None,
797
+ ),
798
+ PredictedEndpoint(
799
+ path="/__debug__/",
800
+ method="GET",
801
+ expected_status=200,
802
+ probability=0.10,
803
+ rationale="static: django-debug-toolbar leak",
804
+ potential_vulnerability_class="info-disclosure",
805
+ ),
806
+ PredictedEndpoint(
807
+ path="/static/admin/css/base.css",
808
+ method="GET",
809
+ expected_status=200,
810
+ probability=0.50,
811
+ rationale="static: Django admin CSS",
812
+ potential_vulnerability_class=None,
813
+ ),
814
+ ),
815
+ "rails": (
816
+ PredictedEndpoint(
817
+ path="/rails/info/routes",
818
+ method="GET",
819
+ expected_status=200,
820
+ probability=0.10,
821
+ rationale="static: Rails dev routes",
822
+ potential_vulnerability_class="info-disclosure",
823
+ ),
824
+ PredictedEndpoint(
825
+ path="/rails/info/properties",
826
+ method="GET",
827
+ expected_status=200,
828
+ probability=0.10,
829
+ rationale="static: Rails dev props",
830
+ potential_vulnerability_class="info-disclosure",
831
+ ),
832
+ PredictedEndpoint(
833
+ path="/rails/conductor",
834
+ method="GET",
835
+ expected_status=200,
836
+ probability=0.05,
837
+ rationale="static: ActionMailbox conductor",
838
+ potential_vulnerability_class="info-disclosure",
839
+ ),
840
+ ),
841
+ "laravel": (
842
+ PredictedEndpoint(
843
+ path="/telescope",
844
+ method="GET",
845
+ expected_status=200,
846
+ probability=0.10,
847
+ rationale="static: Laravel Telescope",
848
+ potential_vulnerability_class="info-disclosure",
849
+ ),
850
+ PredictedEndpoint(
851
+ path="/horizon",
852
+ method="GET",
853
+ expected_status=200,
854
+ probability=0.10,
855
+ rationale="static: Laravel Horizon",
856
+ potential_vulnerability_class="info-disclosure",
857
+ ),
858
+ PredictedEndpoint(
859
+ path="/_ignition/execute-solution",
860
+ method="POST",
861
+ expected_status=200,
862
+ probability=0.10,
863
+ rationale="static: Ignition (CVE-2021-3129)",
864
+ potential_vulnerability_class="rce",
865
+ ),
866
+ PredictedEndpoint(
867
+ path="/storage/logs/laravel.log",
868
+ method="GET",
869
+ expected_status=200,
870
+ probability=0.05,
871
+ rationale="static: Laravel log leak",
872
+ potential_vulnerability_class="info-disclosure",
873
+ ),
874
+ ),
875
+ "spring": (
876
+ PredictedEndpoint(
877
+ path="/actuator",
878
+ method="GET",
879
+ expected_status=200,
880
+ probability=0.40,
881
+ rationale="static: Spring Boot Actuator",
882
+ potential_vulnerability_class="info-disclosure",
883
+ ),
884
+ PredictedEndpoint(
885
+ path="/actuator/env",
886
+ method="GET",
887
+ expected_status=200,
888
+ probability=0.15,
889
+ rationale="static: Actuator env",
890
+ potential_vulnerability_class="info-disclosure",
891
+ ),
892
+ PredictedEndpoint(
893
+ path="/actuator/heapdump",
894
+ method="GET",
895
+ expected_status=200,
896
+ probability=0.10,
897
+ rationale="static: Actuator heapdump",
898
+ potential_vulnerability_class="info-disclosure",
899
+ ),
900
+ PredictedEndpoint(
901
+ path="/actuator/health",
902
+ method="GET",
903
+ expected_status=200,
904
+ probability=0.50,
905
+ rationale="static: Actuator health",
906
+ potential_vulnerability_class=None,
907
+ ),
908
+ PredictedEndpoint(
909
+ path="/actuator/info",
910
+ method="GET",
911
+ expected_status=200,
912
+ probability=0.40,
913
+ rationale="static: Actuator info",
914
+ potential_vulnerability_class=None,
915
+ ),
916
+ PredictedEndpoint(
917
+ path="/actuator/mappings",
918
+ method="GET",
919
+ expected_status=200,
920
+ probability=0.10,
921
+ rationale="static: Actuator route map",
922
+ potential_vulnerability_class="info-disclosure",
923
+ ),
924
+ PredictedEndpoint(
925
+ path="/actuator/beans",
926
+ method="GET",
927
+ expected_status=200,
928
+ probability=0.10,
929
+ rationale="static: Actuator beans",
930
+ potential_vulnerability_class="info-disclosure",
931
+ ),
932
+ PredictedEndpoint(
933
+ path="/actuator/configprops",
934
+ method="GET",
935
+ expected_status=200,
936
+ probability=0.10,
937
+ rationale="static: Actuator config",
938
+ potential_vulnerability_class="info-disclosure",
939
+ ),
940
+ PredictedEndpoint(
941
+ path="/actuator/threaddump",
942
+ method="GET",
943
+ expected_status=200,
944
+ probability=0.10,
945
+ rationale="static: Actuator threaddump",
946
+ potential_vulnerability_class="info-disclosure",
947
+ ),
948
+ PredictedEndpoint(
949
+ path="/h2-console",
950
+ method="GET",
951
+ expected_status=200,
952
+ probability=0.05,
953
+ rationale="static: H2 console (RCE risk)",
954
+ potential_vulnerability_class="rce",
955
+ ),
956
+ PredictedEndpoint(
957
+ path="/jolokia",
958
+ method="GET",
959
+ expected_status=200,
960
+ probability=0.05,
961
+ rationale="static: Jolokia JMX-over-HTTP",
962
+ potential_vulnerability_class="info-disclosure",
963
+ ),
964
+ ),
965
+ "express": (
966
+ PredictedEndpoint(
967
+ path="/static/",
968
+ method="GET",
969
+ expected_status=200,
970
+ probability=0.20,
971
+ rationale="static: Express static",
972
+ potential_vulnerability_class=None,
973
+ ),
974
+ ),
975
+ "flask": (
976
+ PredictedEndpoint(
977
+ path="/static/",
978
+ method="GET",
979
+ expected_status=200,
980
+ probability=0.40,
981
+ rationale="static: Flask static",
982
+ potential_vulnerability_class=None,
983
+ ),
984
+ PredictedEndpoint(
985
+ path="/console",
986
+ method="GET",
987
+ expected_status=200,
988
+ probability=0.05,
989
+ rationale="static: Werkzeug debug console (RCE)",
990
+ potential_vulnerability_class="rce",
991
+ ),
992
+ ),
993
+ "asp.net": (
994
+ PredictedEndpoint(
995
+ path="/web.config",
996
+ method="GET",
997
+ expected_status=200,
998
+ probability=0.05,
999
+ rationale="static: web.config leak",
1000
+ potential_vulnerability_class="info-disclosure",
1001
+ ),
1002
+ PredictedEndpoint(
1003
+ path="/trace.axd",
1004
+ method="GET",
1005
+ expected_status=200,
1006
+ probability=0.05,
1007
+ rationale="static: ASP.NET trace",
1008
+ potential_vulnerability_class="info-disclosure",
1009
+ ),
1010
+ PredictedEndpoint(
1011
+ path="/elmah.axd",
1012
+ method="GET",
1013
+ expected_status=200,
1014
+ probability=0.05,
1015
+ rationale="static: ELMAH error log",
1016
+ potential_vulnerability_class="info-disclosure",
1017
+ ),
1018
+ PredictedEndpoint(
1019
+ path="/default.aspx",
1020
+ method="GET",
1021
+ expected_status=200,
1022
+ probability=0.20,
1023
+ rationale="static: ASP.NET default page",
1024
+ potential_vulnerability_class=None,
1025
+ ),
1026
+ ),
1027
+ "wordpress": (
1028
+ PredictedEndpoint(
1029
+ path="/wp-login.php",
1030
+ method="GET",
1031
+ expected_status=200,
1032
+ probability=0.80,
1033
+ rationale="static: WP login",
1034
+ potential_vulnerability_class="auth-bypass",
1035
+ ),
1036
+ PredictedEndpoint(
1037
+ path="/wp-admin/",
1038
+ method="GET",
1039
+ expected_status=200,
1040
+ probability=0.70,
1041
+ rationale="static: WP admin",
1042
+ potential_vulnerability_class="auth-bypass",
1043
+ ),
1044
+ PredictedEndpoint(
1045
+ path="/wp-content/",
1046
+ method="GET",
1047
+ expected_status=200,
1048
+ probability=0.50,
1049
+ rationale="static: WP content",
1050
+ potential_vulnerability_class=None,
1051
+ ),
1052
+ PredictedEndpoint(
1053
+ path="/wp-json/wp/v2/users",
1054
+ method="GET",
1055
+ expected_status=200,
1056
+ probability=0.30,
1057
+ rationale="static: WP REST user enum",
1058
+ potential_vulnerability_class="info-disclosure",
1059
+ ),
1060
+ PredictedEndpoint(
1061
+ path="/xmlrpc.php",
1062
+ method="GET",
1063
+ expected_status=200,
1064
+ probability=0.40,
1065
+ rationale="static: WP XMLRPC",
1066
+ potential_vulnerability_class="auth-bypass",
1067
+ ),
1068
+ PredictedEndpoint(
1069
+ path="/wp-config.php.bak",
1070
+ method="GET",
1071
+ expected_status=200,
1072
+ probability=0.05,
1073
+ rationale="static: WP config backup",
1074
+ potential_vulnerability_class="info-disclosure",
1075
+ ),
1076
+ PredictedEndpoint(
1077
+ path="/readme.html",
1078
+ method="GET",
1079
+ expected_status=200,
1080
+ probability=0.30,
1081
+ rationale="static: WP version disclosure",
1082
+ potential_vulnerability_class="info-disclosure",
1083
+ ),
1084
+ ),
1085
+ "drupal": (
1086
+ PredictedEndpoint(
1087
+ path="/user/login",
1088
+ method="GET",
1089
+ expected_status=200,
1090
+ probability=0.50,
1091
+ rationale="static: Drupal login",
1092
+ potential_vulnerability_class=None,
1093
+ ),
1094
+ PredictedEndpoint(
1095
+ path="/CHANGELOG.txt",
1096
+ method="GET",
1097
+ expected_status=200,
1098
+ probability=0.15,
1099
+ rationale="static: Drupal changelog (version leak)",
1100
+ potential_vulnerability_class="info-disclosure",
1101
+ ),
1102
+ PredictedEndpoint(
1103
+ path="/core/CHANGELOG.txt",
1104
+ method="GET",
1105
+ expected_status=200,
1106
+ probability=0.15,
1107
+ rationale="static: Drupal 8+ changelog",
1108
+ potential_vulnerability_class="info-disclosure",
1109
+ ),
1110
+ PredictedEndpoint(
1111
+ path="/update.php",
1112
+ method="GET",
1113
+ expected_status=200,
1114
+ probability=0.10,
1115
+ rationale="static: Drupal update",
1116
+ potential_vulnerability_class=None,
1117
+ ),
1118
+ PredictedEndpoint(
1119
+ path="/sites/default/files/",
1120
+ method="GET",
1121
+ expected_status=200,
1122
+ probability=0.20,
1123
+ rationale="static: Drupal files dir",
1124
+ potential_vulnerability_class=None,
1125
+ ),
1126
+ ),
1127
+ "joomla": (
1128
+ PredictedEndpoint(
1129
+ path="/administrator/",
1130
+ method="GET",
1131
+ expected_status=200,
1132
+ probability=0.50,
1133
+ rationale="static: Joomla admin",
1134
+ potential_vulnerability_class="auth-bypass",
1135
+ ),
1136
+ PredictedEndpoint(
1137
+ path="/administrator/index.php",
1138
+ method="GET",
1139
+ expected_status=200,
1140
+ probability=0.40,
1141
+ rationale="static: Joomla admin login",
1142
+ potential_vulnerability_class="auth-bypass",
1143
+ ),
1144
+ PredictedEndpoint(
1145
+ path="/configuration.php-dist",
1146
+ method="GET",
1147
+ expected_status=200,
1148
+ probability=0.05,
1149
+ rationale="static: Joomla config template",
1150
+ potential_vulnerability_class="info-disclosure",
1151
+ ),
1152
+ ),
1153
+ }
1154
+
1155
+
1156
+ def deterministic_endpoints(fingerprint: FingerprintVector) -> list[PredictedEndpoint]:
1157
+ """Compute the rulebook-only endpoint list for a given fingerprint.
1158
+
1159
+ Returns ``ALWAYS_PROBE`` first, then every ``FRAMEWORK_RULES`` entry whose
1160
+ key appears as a case-insensitive substring of ``fingerprint.framework``.
1161
+ A single fingerprint can match multiple rule packs in principle, though
1162
+ the keys are designed not to overlap by accident.
1163
+
1164
+ Pure: no I/O, no side effects.
1165
+ """
1166
+ out: list[PredictedEndpoint] = list(ALWAYS_PROBE)
1167
+ fw_norm = (fingerprint.framework or "").lower().strip()
1168
+ if fw_norm and fw_norm != "unknown":
1169
+ for key, rules in FRAMEWORK_RULES.items():
1170
+ if key in fw_norm:
1171
+ out.extend(rules)
1172
+ return out
1173
+
1174
+
1175
+ def merge_endpoints(
1176
+ deterministic: list[PredictedEndpoint],
1177
+ llm_manifest: MirroredApplicationManifest,
1178
+ ) -> MirroredApplicationManifest:
1179
+ """Merge deterministic + LLM predictions, deduped by (method, path).
1180
+
1181
+ The LLM's prediction wins on conflict. Reasoning: the compact prompt
1182
+ explicitly tells the LLM NOT to re-predict deterministic paths, so if it
1183
+ does anyway, that's *target-specific evidence* (e.g. it spotted
1184
+ ``/.git/HEAD`` linked from robots.txt and wants to assign higher
1185
+ probability than our 0.15 default). Throwing that signal away would
1186
+ defeat the point of running the LLM at all on those paths.
1187
+
1188
+ Output is sorted by probability descending for human debug-readability;
1189
+ :func:`rank_probes` re-sorts by score downstream, so ordering here is
1190
+ cosmetic.
1191
+ """
1192
+ by_key: dict[tuple[str, str], PredictedEndpoint] = {
1193
+ (ep.method, ep.path): ep for ep in deterministic
1194
+ }
1195
+ for ep in llm_manifest.predicted_endpoints:
1196
+ by_key[(ep.method, ep.path)] = ep # LLM wins on conflict
1197
+ merged = sorted(by_key.values(), key=lambda e: e.probability, reverse=True)
1198
+ return MirroredApplicationManifest(
1199
+ predicted_endpoints=merged,
1200
+ global_assumptions=llm_manifest.global_assumptions,
1201
+ )
1202
+
1203
+
1204
+ _MIRRORING_SYSTEM_RULES_COMPACT = """\
1205
+ You are a web-stack code-mirroring assistant. You receive a crystallized
1206
+ fingerprint of one target web application.
1207
+
1208
+ A DETERMINISTIC RULESET runs ALONGSIDE you and ALREADY probes the following
1209
+ classes of paths. Do NOT re-predict them — they will be deduped out, wasting
1210
+ your output budget:
1211
+
1212
+ - VCS/secret leaks: /.git/*, /.env, /.DS_Store
1213
+ - Conventional docs: /robots.txt, /sitemap.xml, /.well-known/security.txt
1214
+ - Server defaults: /server-status, /server-info, /phpinfo.php
1215
+ - API docs: /swagger.json, /swagger-ui.html, /openapi.json, /graphql
1216
+ - Common admin/auth: /admin, /login, /api, /api/v1
1217
+ - Backup leaks: /backup, /backup.zip, /backup.tar.gz
1218
+ - The framework-specific default-routes pack for the fingerprinted stack
1219
+ (Django admin, Rails dev routes, Laravel telescope/horizon/ignition,
1220
+ Spring Boot actuator family, ASP.NET trace/elmah, WordPress/Drupal/Joomla
1221
+ cores, Flask/Werkzeug debug, etc.)
1222
+
1223
+ YOUR job: predict NOVEL, target-specific endpoints the deterministic ruleset
1224
+ would miss — application-specific REST resources implied by the fingerprint,
1225
+ custom admin paths, debug routes inferred from version-specific defaults that
1226
+ aren't already in the rulebook, etc.
1227
+
1228
+ Output ONLY a STRICT JSON object matching this exact schema:
1229
+
1230
+ {
1231
+ "endpoints": [
1232
+ {"p": "/path", "m": "GET", "s": 200, "c": 0.0, "v": null}
1233
+ ],
1234
+ "assumptions": ["..."]
1235
+ }
1236
+
1237
+ Field meanings (single-letter keys are intentional — output tokens are the
1238
+ bottleneck):
1239
+ p = absolute path starting with '/'
1240
+ m = HTTP method, upper case (GET, POST, PUT, PATCH, DELETE, HEAD, OPTIONS)
1241
+ s = expected HTTP status, integer 100..599
1242
+ c = your confidence/probability the path exists, float in [0.0, 1.0]
1243
+ Calibrate strictly: >=0.8 = near-certain, <=0.3 = speculative
1244
+ v = vulnerability class label (short kebab/snake, e.g. 'info-disclosure',
1245
+ 'IDOR', 'SSRF', 'auth-bypass', 'rce', 'ssti') or null if none implied
1246
+
1247
+ Rules:
1248
+ - Do NOT repeat any path covered by the deterministic ruleset above.
1249
+ - Rank `endpoints` by `c`, descending.
1250
+ - `assumptions` is for framework- or deployment-wide hypotheses you relied
1251
+ on (e.g. 'Assumes APP_DEBUG=true based on stack trace').
1252
+ - Output ONLY the JSON object — no prose, no markdown fences, no preamble.
1253
+ """
1254
+
1255
+
1256
+ def build_mirroring_prompt_compact(
1257
+ fingerprint: FingerprintVector,
1258
+ *,
1259
+ correction_context: str | None = None,
1260
+ ) -> str:
1261
+ """Compact-format counterpart to :func:`build_mirroring_prompt`.
1262
+
1263
+ Pure (no I/O) and snapshot-testable. Same correction-context contract as
1264
+ the legacy builder — when set, the block is appended verbatim so the LLM
1265
+ sees which previous predictions were disproven by real probes.
1266
+ """
1267
+ serialized = json.dumps(fingerprint.model_dump(), separators=(",", ":"), default=str)
1268
+ correction_block = ""
1269
+ if correction_context and correction_context.strip():
1270
+ correction_block = (
1271
+ f"--- BEGIN CORRECTION CONTEXT ---\n"
1272
+ f"{correction_context.strip()}\n"
1273
+ f"--- END CORRECTION CONTEXT ---\n"
1274
+ )
1275
+ return (
1276
+ f"{_MIRRORING_SYSTEM_RULES_COMPACT}\n"
1277
+ f"--- BEGIN FINGERPRINT (JSON) ---\n"
1278
+ f"{serialized}\n"
1279
+ f"--- END FINGERPRINT ---\n"
1280
+ f"{correction_block}"
1281
+ f"Return the JSON object now."
1282
+ )
1283
+
1284
+
1285
+ def parse_compact_mirroring_response(response_text: str) -> CompactMirroredManifest:
1286
+ """Parse the LLM's compact-schema JSON output.
1287
+
1288
+ Tolerant of markdown fences and leading/trailing prose so the function
1289
+ works even when ``json_schema`` strict mode is unavailable (e.g. a local
1290
+ vLLM build without guided JSON, or a provider that returns the schema
1291
+ correctly but still wraps it in a ```json fence). Raises
1292
+ :class:`ValueError` with the original text on parse/validation failure so
1293
+ the caller can retry or fall back to an empty manifest.
1294
+ """
1295
+ if not response_text or not response_text.strip():
1296
+ raise ValueError("Empty response from compact-mirroring LLM")
1297
+
1298
+ candidate = _strip_markdown_fence(response_text)
1299
+ if not candidate.startswith("{"):
1300
+ match = _JSON_OBJECT_RE.search(candidate)
1301
+ if match is None:
1302
+ raise ValueError(
1303
+ f"No JSON object found in compact-mirroring LLM response: {response_text!r}"
1304
+ )
1305
+ candidate = match.group(0)
1306
+
1307
+ try:
1308
+ payload = json.loads(candidate)
1309
+ except json.JSONDecodeError as e:
1310
+ raise ValueError(
1311
+ f"Compact-mirroring LLM returned invalid JSON ({e}): {response_text!r}"
1312
+ ) from e
1313
+
1314
+ try:
1315
+ return CompactMirroredManifest.model_validate(payload)
1316
+ except ValidationError as e:
1317
+ raise ValueError(
1318
+ f"Compact-mirroring LLM output failed schema validation: {e}\nRaw payload: {payload!r}"
1319
+ ) from e
1320
+
1321
+
1322
+ _MIRRORING_SYSTEM_RULES = """\
1323
+ You are an elite security architect performing a code-mirroring exercise. You
1324
+ are given a crystallized fingerprint of a single target web application
1325
+ (framework, version guess, frontend, auth pattern, ORM, overall confidence).
1326
+
1327
+ Your job is to hallucinate the application's most probable internal surface:
1328
+ the routes that are likely to exist, the default misconfigurations the stack
1329
+ is known to expose, and the vulnerability classes each predicted endpoint
1330
+ might harbor under default settings. Treat this as a hypothesis-generation
1331
+ step, not a scan — your output drives WHERE the agent probes next.
1332
+
1333
+ Concretely, draw on:
1334
+ - Framework conventions and scaffolded routes
1335
+ (e.g. Django: /admin/, /accounts/login/, /__debug__/;
1336
+ Rails: /rails/info/routes, /rails/conductor;
1337
+ Laravel: /telescope, /horizon, /_ignition/execute-solution;
1338
+ Spring Boot: /actuator, /actuator/env, /actuator/heapdump;
1339
+ Express: /api, default error pages exposing stack traces).
1340
+ - Default exposed files and dotfiles
1341
+ (e.g. /.git/config, /.env, /.DS_Store, /server-status, /server-info,
1342
+ /storage/logs/laravel.log, /web.config, /WEB-INF/web.xml).
1343
+ - Common API structures inferred from the auth pattern
1344
+ (e.g. JWT often implies /api/v1/auth/login, /api/v1/auth/refresh;
1345
+ session-cookie often implies /login, /logout, /csrf, /me).
1346
+ - ORM-implied resource routes
1347
+ (e.g. ActiveRecord -> RESTful /users, /posts; Prisma -> often /api/<model>).
1348
+
1349
+ Output a STRICT JSON object — no prose, no markdown fences, no preamble, no
1350
+ trailing commentary — that matches this schema exactly:
1351
+
1352
+ {
1353
+ "predicted_endpoints": [
1354
+ {
1355
+ "path": "<absolute path starting with '/'>",
1356
+ "method": "<HTTP method in upper case>",
1357
+ "expected_status": <int 100-599>,
1358
+ "probability": <float in [0.0, 1.0]>,
1359
+ "rationale": "<one-line justification tied to a fingerprint signal or framework convention>",
1360
+ "potential_vulnerability_class": "<short class name or null>"
1361
+ }
1362
+ ],
1363
+ "global_assumptions": [
1364
+ "<assumption 1>",
1365
+ "<assumption 2>"
1366
+ ]
1367
+ }
1368
+
1369
+ Rules:
1370
+ - Every predicted_endpoints entry MUST include every field above.
1371
+ - potential_vulnerability_class is null when no specific class is implied;
1372
+ otherwise a short kebab/snake label (e.g. 'info-disclosure', 'IDOR',
1373
+ 'SSRF', 'auth-bypass', 'rce').
1374
+ - probability must be calibrated: reserve >=0.8 for routes that are
1375
+ near-guaranteed under the fingerprinted stack's defaults; use <=0.3 for
1376
+ speculative guesses.
1377
+ - Rank predicted_endpoints by probability, descending.
1378
+ - Do not output anything outside the JSON object.
1379
+ """
1380
+
1381
+
1382
+ def build_mirroring_prompt(
1383
+ fingerprint: FingerprintVector,
1384
+ *,
1385
+ correction_context: str | None = None,
1386
+ ) -> str:
1387
+ """Construct the user-message prompt for the code-mirroring LLM call.
1388
+
1389
+ Pure (no I/O). The fingerprint is serialized as compact JSON so the
1390
+ prompt is deterministic and snapshot-testable. When ``correction_context``
1391
+ is supplied (Stage 4 belief update), it is appended as its own block so
1392
+ the LLM can see which previous predictions were disproven by real probes.
1393
+ """
1394
+ serialized = json.dumps(fingerprint.model_dump(), separators=(",", ":"), default=str)
1395
+ correction_block = ""
1396
+ if correction_context and correction_context.strip():
1397
+ correction_block = (
1398
+ f"--- BEGIN CORRECTION CONTEXT ---\n"
1399
+ f"{correction_context.strip()}\n"
1400
+ f"--- END CORRECTION CONTEXT ---\n"
1401
+ )
1402
+ return (
1403
+ f"{_MIRRORING_SYSTEM_RULES}\n"
1404
+ f"--- BEGIN FINGERPRINT (JSON) ---\n"
1405
+ f"{serialized}\n"
1406
+ f"--- END FINGERPRINT ---\n"
1407
+ f"{correction_block}"
1408
+ f"Return the JSON object now."
1409
+ )
1410
+
1411
+
1412
+ def _strip_markdown_fence(text: str) -> str:
1413
+ """Remove a single surrounding ```json ... ``` (or ``` ... ```) fence.
1414
+
1415
+ Reasoning models drop into fenced output more often than fingerprinting
1416
+ models, so Stage 2 strips fences explicitly before falling back to the
1417
+ Stage 1 regex extraction.
1418
+ """
1419
+ s = text.strip()
1420
+ if not s.startswith("```"):
1421
+ return s
1422
+ # Drop the opening fence line ("```" or "```json", etc.)
1423
+ first_nl = s.find("\n")
1424
+ if first_nl == -1:
1425
+ return s
1426
+ s = s[first_nl + 1 :]
1427
+ if s.rstrip().endswith("```"):
1428
+ s = s.rstrip()[: -len("```")]
1429
+ return s.strip()
1430
+
1431
+
1432
+ def parse_mirroring_response(response_text: str) -> MirroredApplicationManifest:
1433
+ """Parse the LLM's JSON output into a validated manifest.
1434
+
1435
+ Tolerant of markdown fences and leading/trailing prose. Raises
1436
+ :class:`ValueError` with the original text included if parsing or
1437
+ schema validation fails so the caller can retry or fall back to an
1438
+ empty manifest.
1439
+ """
1440
+ if not response_text or not response_text.strip():
1441
+ raise ValueError("Empty response from code-mirroring LLM")
1442
+
1443
+ candidate = _strip_markdown_fence(response_text)
1444
+ if not candidate.startswith("{"):
1445
+ match = _JSON_OBJECT_RE.search(candidate)
1446
+ if match is None:
1447
+ raise ValueError(f"No JSON object found in mirroring LLM response: {response_text!r}")
1448
+ candidate = match.group(0)
1449
+
1450
+ try:
1451
+ payload = json.loads(candidate)
1452
+ except json.JSONDecodeError as e:
1453
+ raise ValueError(f"Mirroring LLM returned invalid JSON ({e}): {response_text!r}") from e
1454
+
1455
+ try:
1456
+ return MirroredApplicationManifest.model_validate(payload)
1457
+ except ValidationError as e:
1458
+ raise ValueError(
1459
+ f"Mirroring LLM output failed schema validation: {e}\nRaw payload: {payload!r}"
1460
+ ) from e
1461
+
1462
+
1463
+ async def _call_llm_with_optional_schema(
1464
+ llm_call: StructuredLLMCallable,
1465
+ prompt: str,
1466
+ *,
1467
+ json_schema: dict[str, Any] | None,
1468
+ ) -> str:
1469
+ """Invoke ``llm_call`` with ``json_schema`` when supported, else without.
1470
+
1471
+ Integrators that wire structured-output mode (OpenAI/Anthropic/vLLM)
1472
+ define ``llm_call`` to accept the ``json_schema`` kwarg and pass it
1473
+ through to ``response_format``. Older test stubs and the basic
1474
+ :data:`LLMCallable` shape don't accept the kwarg; for those we catch the
1475
+ ``TypeError`` and re-call without ``json_schema``. The behavioural
1476
+ fallback means callers don't have to special-case "old SDK" vs "new SDK"
1477
+ — they just get whatever level of strictness their stub supports.
1478
+ """
1479
+ if json_schema is None:
1480
+ return await llm_call(prompt)
1481
+ try:
1482
+ return await llm_call(prompt, json_schema=json_schema)
1483
+ except TypeError as exc:
1484
+ # Distinguish "wrong kwarg" from "wrong type altogether" — the latter
1485
+ # is a real bug we don't want to swallow.
1486
+ msg = str(exc)
1487
+ if "json_schema" in msg or "unexpected keyword" in msg:
1488
+ return await llm_call(prompt)
1489
+ raise
1490
+
1491
+
1492
+ async def generate_mam(
1493
+ fingerprint: FingerprintVector,
1494
+ llm_call: StructuredLLMCallable,
1495
+ *,
1496
+ correction_context: str | None = None,
1497
+ use_compact_schema: bool = True,
1498
+ include_deterministic: bool = True,
1499
+ ) -> MirroredApplicationManifest:
1500
+ """Turn a fingerprint into a ``MirroredApplicationManifest``.
1501
+
1502
+ Default fast path (``use_compact_schema=True``, ``include_deterministic=True``):
1503
+ 1. Build the compact prompt — the deterministic rulebook is announced
1504
+ as already-handled so the LLM doesn't burn output tokens
1505
+ re-predicting it.
1506
+ 2. Call ``llm_call`` with :data:`COMPACT_MAM_JSON_SCHEMA`; the
1507
+ integrator is responsible for plumbing this into the provider's
1508
+ strict-output mode (OpenAI ``response_format``, Anthropic tool-use
1509
+ forcing, vLLM guided JSON). When ``llm_call`` doesn't accept the
1510
+ ``json_schema`` kwarg, we transparently fall back to a plain
1511
+ ``(prompt,) -> str`` call and the response is parsed via the
1512
+ regex-fence-and-extract path in
1513
+ :func:`parse_compact_mirroring_response`.
1514
+ 3. Expand compact -> full :class:`MirroredApplicationManifest`.
1515
+ 4. Merge with :func:`deterministic_endpoints` for the fingerprint;
1516
+ the LLM wins on (method, path) conflicts. See
1517
+ :func:`merge_endpoints` for the rationale.
1518
+
1519
+ Legacy path (``use_compact_schema=False``): pre-refactor behaviour —
1520
+ calls :func:`build_mirroring_prompt` / :func:`parse_mirroring_response`
1521
+ and (by default) still merges in the deterministic ruleset. Kept so
1522
+ existing tests and integrators that haven't migrated still work.
1523
+
1524
+ Set ``include_deterministic=False`` to skip the rulebook entirely —
1525
+ useful for snapshot tests that want to assert only on LLM output.
1526
+ """
1527
+ if use_compact_schema:
1528
+ prompt = build_mirroring_prompt_compact(fingerprint, correction_context=correction_context)
1529
+ response_text = await _call_llm_with_optional_schema(
1530
+ llm_call, prompt, json_schema=COMPACT_MAM_JSON_SCHEMA
1531
+ )
1532
+ compact = parse_compact_mirroring_response(response_text)
1533
+ llm_manifest = expand_compact(compact, source="llm-compact")
1534
+ else:
1535
+ prompt = build_mirroring_prompt(fingerprint, correction_context=correction_context)
1536
+ response_text = await _call_llm_with_optional_schema(llm_call, prompt, json_schema=None)
1537
+ llm_manifest = parse_mirroring_response(response_text)
1538
+
1539
+ if include_deterministic:
1540
+ return merge_endpoints(deterministic_endpoints(fingerprint), llm_manifest)
1541
+ return llm_manifest
1542
+
1543
+
1544
+ # ---------------------------------------------------------------------------
1545
+ # Stage 3: Posterior Construction & Probe Ranking
1546
+ #
1547
+ # Deterministic, LLM-free. Takes a MirroredApplicationManifest and ranks its
1548
+ # predicted endpoints by expected value:
1549
+ #
1550
+ # score = (probability * severity_weight) / cost_weight
1551
+ #
1552
+ # This is the budgeting layer — the agent walks the ranked list top-down and
1553
+ # stops when it runs out of probe budget, instead of issuing requests in MAM
1554
+ # insertion order (which reflects LLM whim, not pentest priority).
1555
+ # ---------------------------------------------------------------------------
1556
+
1557
+
1558
+ # Severity buckets. Tuples of (substring-to-match-case-insensitively, weight).
1559
+ # Order matters: the FIRST matching bucket wins, so the highest-severity
1560
+ # patterns are checked first and a vuln_class like "sql-injection-to-rce"
1561
+ # correctly collapses to the RCE-tier weight.
1562
+ _SEVERITY_BUCKETS: tuple[tuple[tuple[str, ...], float], ...] = (
1563
+ (("rce", "sqli", "command injection"), 10.0),
1564
+ (("xss", "csrf", "traversal"), 7.0),
1565
+ (("disclosure", "leak", "log"), 5.0),
1566
+ )
1567
+ _DEFAULT_SEVERITY_WEIGHT = 1.0
1568
+
1569
+ # Method cost buckets. GET/HEAD are bare reads; the mutating verbs typically
1570
+ # need a body and may trigger side effects, so they cost more to probe.
1571
+ _CHEAP_METHODS: frozenset[str] = frozenset({"GET", "HEAD"})
1572
+ _MUTATING_METHODS: frozenset[str] = frozenset({"POST", "PUT", "PATCH", "DELETE"})
1573
+ _CHEAP_METHOD_COST = 1.0
1574
+ _MUTATING_METHOD_COST = 1.5
1575
+
1576
+
1577
+ def estimate_severity(vuln_class: str | None) -> float:
1578
+ """Map a predicted vulnerability class to a severity weight.
1579
+
1580
+ Case-insensitive substring match against ordered buckets — see
1581
+ ``_SEVERITY_BUCKETS``. Returns the default weight when ``vuln_class`` is
1582
+ ``None`` or matches no bucket; this matches the Stage 2 contract where
1583
+ "no specific class is implied" is encoded as ``None``.
1584
+ """
1585
+ if vuln_class is None:
1586
+ return _DEFAULT_SEVERITY_WEIGHT
1587
+ needle = vuln_class.lower()
1588
+ for patterns, weight in _SEVERITY_BUCKETS:
1589
+ if any(p in needle for p in patterns):
1590
+ return weight
1591
+ return _DEFAULT_SEVERITY_WEIGHT
1592
+
1593
+
1594
+ def estimate_cost(method: str) -> float:
1595
+ """Map an HTTP method to a probe cost weight.
1596
+
1597
+ Unknown methods fall back to the cheap-read cost: most exotic verbs
1598
+ (OPTIONS, TRACE, PROPFIND) are still bare requests, and we'd rather not
1599
+ silently demote a high-severity prediction just because the LLM picked
1600
+ an uncommon verb.
1601
+ """
1602
+ m = method.strip().upper()
1603
+ if m in _MUTATING_METHODS:
1604
+ return _MUTATING_METHOD_COST
1605
+ if m in _CHEAP_METHODS:
1606
+ return _CHEAP_METHOD_COST
1607
+ return _CHEAP_METHOD_COST
1608
+
1609
+
1610
+ class RankedProbe(PredictedEndpoint):
1611
+ """A ``PredictedEndpoint`` enriched with ranking weights and a final score.
1612
+
1613
+ Inherits every field (and validator) from :class:`PredictedEndpoint` so
1614
+ the ranked list can be consumed anywhere a predicted endpoint is
1615
+ expected, while also carrying the heuristic inputs that produced the
1616
+ score — useful for debugging unexpected orderings.
1617
+ """
1618
+
1619
+ severity_weight: float = Field(
1620
+ ge=0.0,
1621
+ description="Severity weight derived from `potential_vulnerability_class`.",
1622
+ )
1623
+ cost_weight: float = Field(
1624
+ gt=0.0,
1625
+ description="Probe cost weight derived from the HTTP `method`. Strictly > 0.",
1626
+ )
1627
+ score: float = Field(
1628
+ ge=0.0,
1629
+ description="Expected value: (probability * severity_weight) / cost_weight.",
1630
+ )
1631
+
1632
+
1633
+ def rank_probes(mam: MirroredApplicationManifest) -> list[RankedProbe]:
1634
+ """Score and rank every predicted endpoint in ``mam`` by expected value.
1635
+
1636
+ Pure: same input always yields the same output, no I/O, no LLM calls.
1637
+ The returned list is sorted by ``score`` descending; ties preserve the
1638
+ MAM's original ordering (Python's sort is stable), which means the LLM's
1639
+ own confidence ranking acts as the tiebreaker.
1640
+ """
1641
+ ranked: list[RankedProbe] = []
1642
+ for ep in mam.predicted_endpoints:
1643
+ severity = estimate_severity(ep.potential_vulnerability_class)
1644
+ cost = estimate_cost(ep.method)
1645
+ score = (ep.probability * severity) / cost
1646
+ ranked.append(
1647
+ RankedProbe(
1648
+ **ep.model_dump(),
1649
+ severity_weight=severity,
1650
+ cost_weight=cost,
1651
+ score=score,
1652
+ )
1653
+ )
1654
+ ranked.sort(key=lambda r: r.score, reverse=True)
1655
+ return ranked
1656
+
1657
+
1658
+ # ---------------------------------------------------------------------------
1659
+ # Stage 4a: Two-tier probing — cheap exhaustive sweep, then targeted deep probe
1660
+ #
1661
+ # A single "top-N by score" pass is fast but lossy: it filters using the LLM's
1662
+ # self-reported confidence before any real evidence exists. CTF designers
1663
+ # routinely hide flags behind endpoints that look low-signal to a generation
1664
+ # model (a stray /robots.txt comment, a boring /backup directory). The fix is
1665
+ # to flip the order: probe everything cheaply first, then let the *cheap
1666
+ # evidence* — not just the LLM score — decide who gets a full deep probe.
1667
+ #
1668
+ # Tier 1 (cheap sweep):
1669
+ # Concurrent GET against EVERY ranked probe, semaphore-capped to avoid
1670
+ # hammering small CTF VMs, body streamed and truncated to ~4KB. Bounded by a
1671
+ # total wall-clock budget; partial results survive a timeout instead of all
1672
+ # being cancelled by `asyncio.wait_for`. Method is always GET regardless of
1673
+ # probe.method — HEAD has too much per-server variance, and POST/PUT/DELETE
1674
+ # could mutate state on a real target. A 405 from a GET on a POST endpoint
1675
+ # is itself useful signal that the route exists.
1676
+ #
1677
+ # Selection (pure):
1678
+ # A probe graduates to Tier 2 if ANY of:
1679
+ # - Tier-1 status is in INTERESTING_STATUS_CODES (200/301/302/401/403/405/500/...).
1680
+ # - Tier-1 body matched the flag-shaped regex — direct flag detection
1681
+ # cuts straight to the answer.
1682
+ # - Path matched a keyword pattern (admin/.git/.env/swagger/etc.). This
1683
+ # is a BOOST, not an override — top-N by score already wins below.
1684
+ # Plus the top-N by Stage-3 score regardless (the LLM's bet always gets a
1685
+ # deep probe even if Tier 1 looked uninteresting). Cap at ``max_deep_probes``.
1686
+ #
1687
+ # Tier 2 (deep probe):
1688
+ # Existing `_execute_probe` using probe.method, same semaphore. Results feed
1689
+ # the unchanged surprise heuristic.
1690
+ #
1691
+ # Note: this optimizes the *HTTP* leg of PHANTOM. If profiling shows the
1692
+ # bottleneck is LLM generation (Stage 1 crystallization or Stage 2 mirroring),
1693
+ # the win is elsewhere — measure before tuning these constants.
1694
+ # ---------------------------------------------------------------------------
1695
+
1696
+
1697
+ _DEFAULT_TOP_N = 10
1698
+ _DEFAULT_SURPRISE_TOP_K = 3
1699
+ _PROBE_BODY_SNIPPET_CHARS = 200
1700
+
1701
+ _DEFAULT_CHEAP_TIMEOUT_SECONDS = 4.0
1702
+ _DEFAULT_CHEAP_BUDGET_SECONDS = 15.0
1703
+ _DEFAULT_CONCURRENCY = 8
1704
+ _DEFAULT_MAX_DEEP_PROBES = 20
1705
+ _CHEAP_BODY_BYTES = 4096
1706
+ _CHEAP_BODY_SNIPPET_CHARS = 512
1707
+
1708
+ # Status codes that historically carry signal on CTF / lab targets. 405 is in
1709
+ # here on purpose: a GET against a POST-only endpoint that returns 405 still
1710
+ # proves the route exists, which is exactly what we want Tier 2 to verify.
1711
+ INTERESTING_STATUS_CODES: frozenset[int] = frozenset(
1712
+ {200, 201, 202, 204, 301, 302, 307, 308, 401, 403, 405, 500, 501, 503}
1713
+ )
1714
+
1715
+ # High-value path substrings used as a Tier-2 promotion BOOST (not an override
1716
+ # of the LLM score). Lower-case, substring match. The list intentionally
1717
+ # leans heavy on CTF/HTB defaults — dotfiles, framework debug routes, common
1718
+ # leaks. ``flag`` is included for completeness but in practice the cheap-sweep
1719
+ # flag-regex check catches the real wins.
1720
+ DEFAULT_KEYWORD_PATTERNS: tuple[str, ...] = (
1721
+ "admin",
1722
+ "secret",
1723
+ "api",
1724
+ "login",
1725
+ "flag",
1726
+ "debug",
1727
+ "dev",
1728
+ "test",
1729
+ "backup",
1730
+ "config",
1731
+ ".git",
1732
+ ".env",
1733
+ ".bak",
1734
+ ".old",
1735
+ ".zip",
1736
+ ".sql",
1737
+ ".tar",
1738
+ ".ds_store",
1739
+ "phpinfo",
1740
+ "console",
1741
+ "swagger",
1742
+ "graphql",
1743
+ "actuator",
1744
+ "metrics",
1745
+ "health",
1746
+ "wp-",
1747
+ "phpmyadmin",
1748
+ "web.config",
1749
+ "id_rsa",
1750
+ "private",
1751
+ "internal",
1752
+ "hidden",
1753
+ )
1754
+
1755
+ # Mirrors ``AgentController.FLAG_PATTERNS`` so a flag visible in a Tier-1
1756
+ # snippet immediately promotes that endpoint to Tier 2. Kept local (not
1757
+ # imported from controller) to avoid a phantom -> controller dependency.
1758
+ _FLAG_SHAPED_RE = re.compile(
1759
+ r"(?:flag|FLAG|HTB|THM|CTF|picoCTF|HackTheBox)\{[^}\s]+\}|\b[a-f0-9]{32}\b",
1760
+ re.IGNORECASE,
1761
+ )
1762
+
1763
+
1764
+ class CheapProbeResult(BaseModel):
1765
+ """Tier-1 sweep result: a streamed-and-truncated GET against one endpoint.
1766
+
1767
+ Treated as *evidence* feeding the selection step, not a verdict. A status
1768
+ of ``0`` means the request failed before producing a response (timeout,
1769
+ connection refused, TLS error). ``contains_flag_shape`` and
1770
+ ``keyword_hits`` are precomputed here so the pure selector doesn't have
1771
+ to re-scan paths and bodies.
1772
+ """
1773
+
1774
+ probe: RankedProbe = Field(description="The ranked probe that was swept.")
1775
+ actual_status: int = Field(description="HTTP status returned, or 0 on transport error.")
1776
+ content_type: str = Field(default="", description="Response Content-Type header.")
1777
+ body_length: int = Field(default=0, description="Bytes downloaded (capped by the sweep).")
1778
+ body_snippet: str = Field(
1779
+ default="",
1780
+ description=f"First ~{_CHEAP_BODY_SNIPPET_CHARS} chars of the body.",
1781
+ )
1782
+ contains_flag_shape: bool = Field(
1783
+ default=False,
1784
+ description="True iff body_snippet matched the flag-shaped regex.",
1785
+ )
1786
+ keyword_hits: tuple[str, ...] = Field(
1787
+ default_factory=tuple,
1788
+ description="Keyword substrings matched against the (lower-cased) path.",
1789
+ )
1790
+ error: str | None = Field(
1791
+ default=None,
1792
+ description="Bracketed error string when the request failed pre-response.",
1793
+ )
1794
+
1795
+
1796
+ async def _cheap_probe_one(
1797
+ http: httpx.AsyncClient,
1798
+ base: str,
1799
+ probe: RankedProbe,
1800
+ *,
1801
+ timeout_seconds: float,
1802
+ body_cap_bytes: int,
1803
+ keyword_patterns: tuple[str, ...],
1804
+ ) -> CheapProbeResult:
1805
+ """Stream one GET, truncate the body, and classify the result.
1806
+
1807
+ Always uses GET regardless of ``probe.method``: in the cheap sweep we
1808
+ care about "does this route exist and what does it return?", not about
1809
+ semantically mutating the target. A 405 here is itself signal.
1810
+ """
1811
+ url = urljoin(base + "/", probe.path.lstrip("/"))
1812
+ path_lower = probe.path.lower()
1813
+ keyword_hits = tuple(k for k in keyword_patterns if k in path_lower)
1814
+
1815
+ try:
1816
+ async with http.stream("GET", url, timeout=timeout_seconds) as resp:
1817
+ body = bytearray()
1818
+ async for chunk in resp.aiter_bytes():
1819
+ body.extend(chunk)
1820
+ if len(body) >= body_cap_bytes:
1821
+ break
1822
+ snippet = bytes(body).decode("utf-8", errors="replace")[:_CHEAP_BODY_SNIPPET_CHARS]
1823
+ return CheapProbeResult(
1824
+ probe=probe,
1825
+ actual_status=resp.status_code,
1826
+ content_type=resp.headers.get("content-type", ""),
1827
+ body_length=len(body),
1828
+ body_snippet=snippet,
1829
+ contains_flag_shape=bool(_FLAG_SHAPED_RE.search(snippet)),
1830
+ keyword_hits=keyword_hits,
1831
+ )
1832
+ except (TimeoutError, httpx.TimeoutException):
1833
+ return CheapProbeResult(
1834
+ probe=probe,
1835
+ actual_status=0,
1836
+ keyword_hits=keyword_hits,
1837
+ error="timeout",
1838
+ )
1839
+ except httpx.RequestError as e:
1840
+ return CheapProbeResult(
1841
+ probe=probe,
1842
+ actual_status=0,
1843
+ keyword_hits=keyword_hits,
1844
+ error=f"{e.__class__.__name__}: {e}",
1845
+ )
1846
+
1847
+
1848
+ async def cheap_sweep(
1849
+ base: str,
1850
+ probes: list[RankedProbe],
1851
+ http: httpx.AsyncClient,
1852
+ *,
1853
+ timeout_seconds: float = _DEFAULT_CHEAP_TIMEOUT_SECONDS,
1854
+ total_budget_seconds: float = _DEFAULT_CHEAP_BUDGET_SECONDS,
1855
+ concurrency: int = _DEFAULT_CONCURRENCY,
1856
+ body_cap_bytes: int = _CHEAP_BODY_BYTES,
1857
+ keyword_patterns: tuple[str, ...] = DEFAULT_KEYWORD_PATTERNS,
1858
+ ) -> list[CheapProbeResult]:
1859
+ """Tier 1: cheap, exhaustive GET sweep with semaphore + total wall budget.
1860
+
1861
+ Probes are fanned out under an :class:`asyncio.Semaphore` so that no more
1862
+ than ``concurrency`` requests fly in parallel (CTF VMs and shared HTB
1863
+ boxes routinely fall over under uncapped ``asyncio.gather``). The whole
1864
+ sweep is bounded by ``total_budget_seconds``; tasks still pending at the
1865
+ deadline are cancelled, and the *partial* result list (completed tasks
1866
+ only, in submission order) is returned. Callers can detect a truncated
1867
+ sweep by comparing ``len(result) < len(probes)``.
1868
+
1869
+ Returns ``[]`` when ``probes`` is empty (no I/O performed).
1870
+ """
1871
+ if not probes:
1872
+ return []
1873
+
1874
+ semaphore = asyncio.Semaphore(max(1, concurrency))
1875
+
1876
+ async def _bounded(p: RankedProbe) -> CheapProbeResult:
1877
+ async with semaphore:
1878
+ return await _cheap_probe_one(
1879
+ http,
1880
+ base,
1881
+ p,
1882
+ timeout_seconds=timeout_seconds,
1883
+ body_cap_bytes=body_cap_bytes,
1884
+ keyword_patterns=keyword_patterns,
1885
+ )
1886
+
1887
+ tasks: list[asyncio.Task[CheapProbeResult]] = [asyncio.create_task(_bounded(p)) for p in probes]
1888
+
1889
+ try:
1890
+ # return_exceptions=True so a single rogue task can't blow up the
1891
+ # gather before the budget timer trips.
1892
+ await asyncio.wait_for(
1893
+ asyncio.gather(*tasks, return_exceptions=True),
1894
+ timeout=total_budget_seconds,
1895
+ )
1896
+ except TimeoutError:
1897
+ # wait_for cancels the inner gather, which cascade-cancels every
1898
+ # pending task. Drain them so we don't leave dangling coroutines.
1899
+ for t in tasks:
1900
+ if not t.done():
1901
+ t.cancel()
1902
+ await asyncio.gather(*tasks, return_exceptions=True)
1903
+
1904
+ results: list[CheapProbeResult] = []
1905
+ for t in tasks:
1906
+ if t.done() and not t.cancelled():
1907
+ try:
1908
+ results.append(t.result())
1909
+ except Exception:
1910
+ # _cheap_probe_one is supposed to catch every transport error
1911
+ # and return a CheapProbeResult; any leak is treated as a
1912
+ # non-result rather than killing the whole sweep.
1913
+ continue
1914
+ return results
1915
+
1916
+
1917
+ def select_for_deep_probe(
1918
+ cheap_results: list[CheapProbeResult],
1919
+ ranked: list[RankedProbe],
1920
+ *,
1921
+ interesting_status: frozenset[int] = INTERESTING_STATUS_CODES,
1922
+ top_n_forced: int = _DEFAULT_TOP_N,
1923
+ max_deep_probes: int = _DEFAULT_MAX_DEEP_PROBES,
1924
+ ) -> list[RankedProbe]:
1925
+ """Pick which probes deserve a full Tier-2 probe.
1926
+
1927
+ Pure: no I/O. Inputs are the Tier-1 evidence (``cheap_results``) and the
1928
+ full Stage-3 ranking (``ranked``). A probe is selected if it is in the
1929
+ top ``top_n_forced`` by score OR its Tier-1 result triggered ANY of:
1930
+ - status in ``interesting_status``,
1931
+ - body matched the flag-shaped regex (``contains_flag_shape``),
1932
+ - path matched a keyword pattern (``keyword_hits`` non-empty).
1933
+
1934
+ Returned order preserves the Stage-3 ranking (highest score first), is
1935
+ de-duplicated by (method, path), and is hard-capped at ``max_deep_probes``
1936
+ to bound Tier-2 cost.
1937
+ """
1938
+ by_key: dict[tuple[str, str], CheapProbeResult] = {
1939
+ (r.probe.method, r.probe.path): r for r in cheap_results
1940
+ }
1941
+
1942
+ selected: list[RankedProbe] = []
1943
+ seen: set[tuple[str, str]] = set()
1944
+
1945
+ # Pass 1 — forced top-N. The LLM's bet always gets a deep probe even if
1946
+ # Tier 1 looked uninteresting (handles the case where the response was
1947
+ # interesting in body but not in status: a 200 ratio of nothing useful).
1948
+ for probe in ranked[:top_n_forced]:
1949
+ key = (probe.method, probe.path)
1950
+ if key not in seen:
1951
+ selected.append(probe)
1952
+ seen.add(key)
1953
+ if len(selected) >= max_deep_probes:
1954
+ return selected
1955
+
1956
+ # Pass 2 — evidence-driven promotions in score order.
1957
+ for probe in ranked:
1958
+ if len(selected) >= max_deep_probes:
1959
+ break
1960
+ key = (probe.method, probe.path)
1961
+ if key in seen:
1962
+ continue
1963
+ cheap = by_key.get(key)
1964
+ if cheap is None:
1965
+ # No Tier-1 evidence (e.g. cheap sweep timed out before reaching
1966
+ # this probe) — defer to the LLM's score, which already had its
1967
+ # chance in Pass 1.
1968
+ continue
1969
+ promote = (
1970
+ cheap.actual_status in interesting_status
1971
+ or cheap.contains_flag_shape
1972
+ or bool(cheap.keyword_hits)
1973
+ )
1974
+ if promote:
1975
+ selected.append(probe)
1976
+ seen.add(key)
1977
+
1978
+ return selected
1979
+
1980
+
1981
+ # ---------------------------------------------------------------------------
1982
+ # Stage 4: Adaptive Verification with Belief Update
1983
+ #
1984
+ # The full PHANTOM loop. Stage 1-3 produce a ranked probe list deterministically;
1985
+ # Stage 4 fires probes against the target, evaluates the results, and — if
1986
+ # the surprise heuristic trips — re-prompts the mirroring LLM with a
1987
+ # correction context that names the disproven predictions. The new MAM is
1988
+ # re-ranked and the next batch fires. Iteration is capped by ``max_iterations``.
1989
+ #
1990
+ # When ``enable_two_tier=True`` (the default) the per-iteration probe step is
1991
+ # the Stage 4a two-tier flow: cheap sweep over the full ranked list, then
1992
+ # selection, then Tier-2 deep probes on the selected subset. When False,
1993
+ # the legacy single-tier flow runs: top-N by score, fired directly with
1994
+ # ``_execute_probe``.
1995
+ #
1996
+ # Concretely: surprise = "the K highest-probability deep-probe results in
1997
+ # this batch all returned a status other than the one predicted". That
1998
+ # signal is strong enough to suggest the fingerprint itself misled the
1999
+ # mirror, so we hand the evidence back to the LLM rather than just trying
2000
+ # more rows from the same (already-disproven) ranking.
2001
+ # ---------------------------------------------------------------------------
2002
+
2003
+
2004
+ class ProbeResult(BaseModel):
2005
+ """Outcome of firing a single ``RankedProbe`` at the target."""
2006
+
2007
+ probe: RankedProbe = Field(description="The probe that was executed.")
2008
+ actual_status: int = Field(
2009
+ description=(
2010
+ "HTTP status returned by the target, or 0 if the request itself "
2011
+ "failed (timeout, connection refused, TLS error, etc.)."
2012
+ ),
2013
+ )
2014
+ confirmed: bool = Field(
2015
+ description="True iff actual_status == probe.expected_status.",
2016
+ )
2017
+ response_snippet: str = Field(
2018
+ description=(
2019
+ f"First {_PROBE_BODY_SNIPPET_CHARS} characters of the response body, "
2020
+ "or a bracketed error string when the request failed before a response."
2021
+ ),
2022
+ )
2023
+
2024
+
2025
+ async def _execute_probe(
2026
+ http: httpx.AsyncClient,
2027
+ base: str,
2028
+ probe: RankedProbe,
2029
+ timeout_seconds: float,
2030
+ ) -> ProbeResult:
2031
+ """Fire one ranked probe and convert the response into a ``ProbeResult``.
2032
+
2033
+ Network/HTTP failures collapse to ``actual_status=0`` with a bracketed
2034
+ error string in ``response_snippet`` — that way the surprise heuristic
2035
+ uniformly treats "failed to reach" and "reached but wrong status" as
2036
+ 'did not match expected_status', without special-casing.
2037
+ """
2038
+ url = urljoin(base + "/", probe.path.lstrip("/"))
2039
+ try:
2040
+ resp = await asyncio.wait_for(
2041
+ http.request(probe.method, url),
2042
+ timeout=timeout_seconds,
2043
+ )
2044
+ except (TimeoutError, httpx.TimeoutException):
2045
+ return ProbeResult(
2046
+ probe=probe,
2047
+ actual_status=0,
2048
+ confirmed=False,
2049
+ response_snippet="<error: timeout>",
2050
+ )
2051
+ except httpx.RequestError as e:
2052
+ return ProbeResult(
2053
+ probe=probe,
2054
+ actual_status=0,
2055
+ confirmed=False,
2056
+ response_snippet=f"<error: {e.__class__.__name__}: {e}>",
2057
+ )
2058
+
2059
+ snippet = (resp.text or "")[:_PROBE_BODY_SNIPPET_CHARS]
2060
+ return ProbeResult(
2061
+ probe=probe,
2062
+ actual_status=resp.status_code,
2063
+ confirmed=(resp.status_code == probe.expected_status),
2064
+ response_snippet=snippet,
2065
+ )
2066
+
2067
+
2068
+ def _build_correction_context(failed: list[ProbeResult]) -> str:
2069
+ """Format previously-failed probes for the belief-update re-prompt.
2070
+
2071
+ Lists each failed probe with method, path, expected vs actual status,
2072
+ its prior probability, and the LLM's original rationale — that last
2073
+ field is the lever the LLM needs to notice 'oh, my framework assumption
2074
+ was wrong' rather than just shuffling paths.
2075
+ """
2076
+ lines = [
2077
+ (
2078
+ f"- {r.probe.method} {r.probe.path} -> expected {r.probe.expected_status}, "
2079
+ f"got {r.actual_status} (prior probability {r.probe.probability:.2f}; "
2080
+ f"original rationale: {r.probe.rationale})"
2081
+ )
2082
+ for r in failed
2083
+ ]
2084
+ return (
2085
+ "PREVIOUS ATTEMPT FAILED. The following high-probability predictions did "
2086
+ "NOT return the expected status when probed. Treat this as strong evidence "
2087
+ "that your prior framework or deployment assumptions were wrong — consider "
2088
+ "a different version, a non-default routing prefix, a different framework "
2089
+ "entirely, or that the target is behind a reverse proxy that rewrites paths. "
2090
+ "Do NOT re-predict any endpoint that has already been disproven below.\n\n"
2091
+ "Disproven predictions:\n" + "\n".join(lines)
2092
+ )
2093
+
2094
+
2095
+ async def execute_phantom_loop(
2096
+ target_url: str,
2097
+ llm_call: LLMCallable,
2098
+ *,
2099
+ max_iterations: int = 2,
2100
+ top_n: int = _DEFAULT_TOP_N,
2101
+ surprise_top_k: int = _DEFAULT_SURPRISE_TOP_K,
2102
+ timeout_seconds: float = _PER_REQUEST_TIMEOUT_SECONDS,
2103
+ client: httpx.AsyncClient | None = None,
2104
+ enable_two_tier: bool = True,
2105
+ cheap_timeout_seconds: float = _DEFAULT_CHEAP_TIMEOUT_SECONDS,
2106
+ cheap_budget_seconds: float = _DEFAULT_CHEAP_BUDGET_SECONDS,
2107
+ concurrency: int = _DEFAULT_CONCURRENCY,
2108
+ max_deep_probes: int = _DEFAULT_MAX_DEEP_PROBES,
2109
+ keyword_patterns: tuple[str, ...] = DEFAULT_KEYWORD_PATTERNS,
2110
+ include_deterministic: bool = True,
2111
+ use_compact_schema: bool = True,
2112
+ ) -> list[ProbeResult]:
2113
+ """Run the full PHANTOM pipeline against ``target_url``.
2114
+
2115
+ Pipeline per iteration:
2116
+ 1. (iteration 0 only) gather raw signals + crystallize fingerprint.
2117
+ 2. ``generate_mam`` — fresh on iteration 0, carrying a correction
2118
+ context on subsequent iterations.
2119
+ 3. ``rank_probes`` over the MAM.
2120
+ 4. Two-tier probe step (``enable_two_tier=True``):
2121
+ 4a. ``cheap_sweep`` GETs the full ranked list (minus already-tried),
2122
+ under a semaphore and total wall budget.
2123
+ 4b. ``select_for_deep_probe`` picks the Tier-2 subset (top-N by
2124
+ score plus evidence-driven promotions).
2125
+ 4c. ``_execute_probe`` fires the selected probes concurrently
2126
+ (same semaphore) using their predicted method.
2127
+ Legacy single-tier path (``enable_two_tier=False``): take the top-N
2128
+ ranked by score and fire them directly with ``_execute_probe``.
2129
+ 5. Surprise check: if the top-K deep-probe results (by probability)
2130
+ from this batch were all unconfirmed, build a correction context
2131
+ and loop. Otherwise return.
2132
+
2133
+ Returns every Tier-2 ``ProbeResult`` collected across all iterations, in
2134
+ the order they were executed. Tier-1 ``CheapProbeResult`` evidence is
2135
+ intentionally NOT in the return — callers that want it should invoke
2136
+ :func:`cheap_sweep` directly.
2137
+
2138
+ ``client`` may be supplied so callers (and tests) can share a single
2139
+ long-lived :class:`httpx.AsyncClient`; when ``None``, a one-shot client
2140
+ is created and reused across Stage 1 recon and Stage 4 probes so that
2141
+ the connection pool warms once.
2142
+
2143
+ ``include_deterministic`` and ``use_compact_schema`` are forwarded to
2144
+ :func:`generate_mam`; leave them ``True`` in production (merge the always-on
2145
+ rulebook; use the compact prompt/schema). Tests that feed a full
2146
+ ``MirroredApplicationManifest`` dump and want to assert only on the LLM's
2147
+ predicted endpoints set both to ``False``.
2148
+ """
2149
+ if max_iterations < 1:
2150
+ raise ValueError("max_iterations must be >= 1")
2151
+ if top_n < 1:
2152
+ raise ValueError("top_n must be >= 1")
2153
+ if surprise_top_k < 1:
2154
+ raise ValueError("surprise_top_k must be >= 1")
2155
+ if concurrency < 1:
2156
+ raise ValueError("concurrency must be >= 1")
2157
+ if max_deep_probes < 1:
2158
+ raise ValueError("max_deep_probes must be >= 1")
2159
+
2160
+ owns_client = False
2161
+ if client is not None:
2162
+ http = client
2163
+ else:
2164
+ from agent.core.backend import get_shared_http_client
2165
+
2166
+ http = get_shared_http_client()
2167
+
2168
+ all_results: list[ProbeResult] = []
2169
+ # Dedupe across iterations: even with the correction context telling the
2170
+ # LLM not to re-predict disproven endpoints, models drift — this is the
2171
+ # belt-and-braces guarantee that we never burn budget probing the same
2172
+ # (method, path) twice.
2173
+ tried: set[tuple[str, str]] = set()
2174
+ correction: str | None = None
2175
+ # One semaphore shared across Tier 1 and Tier 2 within a single iteration
2176
+ # so we never exceed ``concurrency`` in-flight requests against the target.
2177
+ semaphore = asyncio.Semaphore(concurrency)
2178
+
2179
+ try:
2180
+ raw_signals = await gather_raw_signals(
2181
+ target_url, timeout_seconds=timeout_seconds, client=http
2182
+ )
2183
+ base = raw_signals["base"]
2184
+ fingerprint = await crystallize_fingerprint(raw_signals, llm_call)
2185
+
2186
+ # Optional RAG hook for framework CVEs
2187
+ try:
2188
+ if fingerprint.framework and fingerprint.framework.lower() != "unknown":
2189
+ from agent.rag_module.knowledge_base import get_shared_kb
2190
+
2191
+ cves = await get_shared_kb().query_framework_cves(
2192
+ fingerprint.framework,
2193
+ fingerprint.version_guess,
2194
+ top_k=3,
2195
+ )
2196
+ if cves:
2197
+ rag_context = "\n".join(cves)
2198
+ if correction:
2199
+ correction += f"\n\n[RAG Framework Info]:\n{rag_context}"
2200
+ else:
2201
+ correction = f"[RAG Framework Info]:\n{rag_context}"
2202
+ except Exception:
2203
+ pass
2204
+
2205
+ for _iteration in range(max_iterations):
2206
+ mam = await generate_mam(
2207
+ fingerprint,
2208
+ llm_call,
2209
+ correction_context=correction,
2210
+ include_deterministic=include_deterministic,
2211
+ use_compact_schema=use_compact_schema,
2212
+ )
2213
+ ranked = rank_probes(mam)
2214
+ ranked = [r for r in ranked if (r.method, r.path) not in tried]
2215
+ if not ranked:
2216
+ # MAM produced nothing new; further iteration cannot help.
2217
+ break
2218
+
2219
+ if enable_two_tier:
2220
+ # Tier 1: cheap exhaustive GET sweep over the whole ranked list.
2221
+ cheap_results = await cheap_sweep(
2222
+ base,
2223
+ ranked,
2224
+ http,
2225
+ timeout_seconds=cheap_timeout_seconds,
2226
+ total_budget_seconds=cheap_budget_seconds,
2227
+ concurrency=concurrency,
2228
+ keyword_patterns=keyword_patterns,
2229
+ )
2230
+ # Tier 2 selection — top-N by score (the LLM's bet) plus any
2231
+ # cheap-sweep evidence that justifies a deeper look.
2232
+ batch = select_for_deep_probe(
2233
+ cheap_results,
2234
+ ranked,
2235
+ top_n_forced=top_n,
2236
+ max_deep_probes=max_deep_probes,
2237
+ )
2238
+ else:
2239
+ # Legacy single-tier behaviour, preserved for callers that
2240
+ # explicitly opt out (and to keep older tests honest).
2241
+ batch = ranked[:top_n]
2242
+
2243
+ if not batch:
2244
+ break
2245
+ for r in batch:
2246
+ tried.add((r.method, r.path))
2247
+
2248
+ results = await _execute_batch(http, base, batch, timeout_seconds, semaphore)
2249
+ all_results.extend(results)
2250
+
2251
+ # Surprise heuristic: rank THIS batch by probability (not by score)
2252
+ # and check whether the top-K all failed. Probability is the right
2253
+ # axis here because severity weighting can promote a low-probability
2254
+ # high-impact guess into the top-N — we don't want one such miss
2255
+ # to dominate the belief-update signal.
2256
+ by_prob = sorted(results, key=lambda r: r.probe.probability, reverse=True)
2257
+ top_k = by_prob[:surprise_top_k]
2258
+ if len(top_k) < surprise_top_k or any(r.confirmed for r in top_k):
2259
+ break
2260
+
2261
+ correction = _build_correction_context(top_k)
2262
+ finally:
2263
+ if owns_client:
2264
+ await http.aclose()
2265
+
2266
+ return all_results
2267
+
2268
+
2269
+ async def _execute_batch(
2270
+ http: httpx.AsyncClient,
2271
+ base: str,
2272
+ batch: list[RankedProbe],
2273
+ timeout_seconds: float,
2274
+ semaphore: asyncio.Semaphore,
2275
+ ) -> list[ProbeResult]:
2276
+ """Run Tier-2 probes concurrently under ``semaphore``.
2277
+
2278
+ Extracted so the two-tier and legacy paths share identical concurrency
2279
+ semantics — the only thing that varies between them is how ``batch`` was
2280
+ chosen.
2281
+ """
2282
+
2283
+ async def _bounded(p: RankedProbe) -> ProbeResult:
2284
+ async with semaphore:
2285
+ return await _execute_probe(http, base, p, timeout_seconds)
2286
+
2287
+ coros = [_bounded(p) for p in batch]
2288
+ return await asyncio.gather(*coros, return_exceptions=False)
2289
+
2290
+
2291
+ __all__ = [
2292
+ "ALWAYS_PROBE",
2293
+ "COMPACT_MAM_JSON_SCHEMA",
2294
+ "DEFAULT_KEYWORD_PATTERNS",
2295
+ "DEFAULT_RECON_PATHS",
2296
+ "FRAMEWORK_RULES",
2297
+ "INTERESTING_STATUS_CODES",
2298
+ "CheapProbeResult",
2299
+ "CompactMirroredManifest",
2300
+ "FingerprintVector",
2301
+ "LLMCallable",
2302
+ "MirroredApplicationManifest",
2303
+ "PathSignal",
2304
+ "PredictedEndpoint",
2305
+ "PredictedEndpointCompact",
2306
+ "ProbeResult",
2307
+ "RankedProbe",
2308
+ "StructuredLLMCallable",
2309
+ "build_crystallization_prompt",
2310
+ "build_mirroring_prompt",
2311
+ "build_mirroring_prompt_compact",
2312
+ "cheap_sweep",
2313
+ "crystallize_fingerprint",
2314
+ "deterministic_endpoints",
2315
+ "estimate_cost",
2316
+ "estimate_severity",
2317
+ "execute_phantom_loop",
2318
+ "expand_compact",
2319
+ "gather_raw_signals",
2320
+ "generate_mam",
2321
+ "merge_endpoints",
2322
+ "parse_compact_mirroring_response",
2323
+ "parse_fingerprint_response",
2324
+ "parse_mirroring_response",
2325
+ "rank_probes",
2326
+ "select_for_deep_probe",
2327
+ ]