Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
14 changes: 10 additions & 4 deletions src/exploit_iq_commons/data_models/checker_status.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,10 +192,16 @@ class VulnerabilityIntel(BaseModel):
description="Human-readable rationale behind identify_phase_fix_signal (e.g. NVRs compared)."
)

def format_for_prompt(self) -> str:
def format_for_prompt(self, *, include_pattern_lists: bool = True) -> str:
"""Format VulnerabilityIntel for injection into L1 agent runtime prompt.

Uses UPPERCASE labels so they can be referenced as anchors in thought prompts.

Args:
include_pattern_lists: When False, omit VULNERABLE_PATTERNS / FIX_PATTERNS
call-site lists (e.g. when FILE_CHANGES already carries per-file
removed/added patterns). Object fields are unchanged; functions and
keywords still emit. Default True preserves legacy full formatting.
"""
lines = []
if self.is_downstream_patch_available:
Expand All @@ -209,11 +215,11 @@ def format_for_prompt(self) -> str:
lines.append(f"VULNERABLE_FUNCTIONS: {', '.join(self.vulnerable_functions)}")
if self.vulnerable_variables:
lines.append(f"VULNERABLE_VARIABLES: {', '.join(self.vulnerable_variables)}")
if self.vulnerable_patterns:
if include_pattern_lists and self.vulnerable_patterns:
lines.append("VULNERABLE_PATTERNS:")
for p in self.vulnerable_patterns:
lines.append(f" - {p}")
if self.fix_patterns:
if include_pattern_lists and self.fix_patterns:
lines.append("FIX_PATTERNS:")
for p in self.fix_patterns:
lines.append(f" - {p}")
Expand Down
325 changes: 324 additions & 1 deletion src/vuln_analysis/functions/code_agent_graph_defs.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,8 @@ class CodeAgentState(MessagesState):
thought: NotRequired[CheckerThought | None]
observation: NotRequired[Observation | None]
vulnerability_intel: NotRequired["VulnerabilityIntel | None"]
# New intel flow result (Issue 08: per-file focused observation context)
intel_gathering_result: NotRequired["IntelGatheringResult | None"]
arch_mismatch_reason: NotRequired[str | None]
# Reference Mining state (populated when no patch found via existing flows)
reference_hints: NotRequired["ReferenceHints | None"]
Expand Down Expand Up @@ -140,6 +142,327 @@ class ParsedPatch(BaseModel):
files: list[PatchFile]


class FileChangeSummary(BaseModel):
"""Per-file summary of changes from a patch for tracing/debugging."""
file_path: str
is_new_file: bool = False
is_deleted_file: bool = False
removed_lines_count: int = 0
added_lines_count: int = 0
hunks_count: int = 0
functions_touched: list[str] = Field(default_factory=list)
key_removed_patterns: list[str] = Field(default_factory=list)
key_added_patterns: list[str] = Field(default_factory=list)
estimated_tokens: int = 0


class CVEUnderstanding(BaseModel):
"""CVE-level intelligence extracted without patch (Prompt 1 output)."""
vulnerability_type: str = Field(
default="",
description="Category: buffer_overflow, integer_overflow, use_after_free, etc."
)
affected_component: str = Field(
default="",
description="High-level component affected (e.g., 'memory allocator', 'XML parser')"
)
attack_vector: Literal["network", "local", "physical", "adjacent", "unknown"] = Field(
default="unknown",
description="How the vulnerability is triggered"
)
cve_keywords: list[str] = Field(
default_factory=list,
description="Search terms ordered by specificity (most specific first)"
)
root_cause: str = Field(
default="",
description="Technical explanation of why the code is vulnerable"
)
component_names: list[str] = Field(
default_factory=list,
description=(
"Module/component names mentioned verbatim in CVE/advisory "
"(e.g. mod_http2). Empty if none explicit."
),
)
affected_bitness: Literal["32-bit", "64-bit", "both"] = Field(
default="both",
description="Which bitness is affected: 32-bit only, 64-bit only, or both",
)
affected_architectures: list[str] | None = Field(
default=None,
description="CPU families affected (e.g. ['x86']). None means all architectures.",
)


class PerFileIntel(BaseModel):
"""Per-file intelligence from patch analysis (Prompt 2/3 output)."""
file_path: str
vulnerable_functions: list[str] = Field(default_factory=list)
vulnerable_patterns: list[str] = Field(default_factory=list)
fix_patterns: list[str] = Field(default_factory=list)
search_keywords: list[str] = Field(default_factory=list)


class IntelGatheringResult(BaseModel):
"""Internal object holding all gathered intelligence, queryable per-file.

This is the NEW internal representation. It can generate a VulnerabilityIntel
for backward compatibility with existing consumers.
"""

cve_understanding: CVEUnderstanding | None = None
per_file_intel: dict[str, PerFileIntel] = Field(default_factory=dict)
file_summaries: list[FileChangeSummary] = Field(default_factory=list)
patch_source: str = ""
total_files_in_patch: int = 0
files_analyzed_by_llm: int = 0
prompt_count: int = 0

def get_intel_for_file(self, file_path: str) -> PerFileIntel | None:
"""Query intel for a specific file (for focused observation context).

Matches by exact path or basename equality only. Does not use substring
containment (avoids 'rsync' matching grep target 'rsync.c').
"""
if not file_path:
return None

query = file_path.replace("\\", "/").lower().lstrip("./")
query_basename = Path(query).name

if file_path in self.per_file_intel:
return self.per_file_intel[file_path]

for key, intel in self.per_file_intel.items():
key_norm = key.replace("\\", "/").lower().lstrip("./")
if query == key_norm:
return intel
if query_basename == Path(key_norm).name:
return intel

return None

def get_functions_for_file(self, file_path: str) -> list[str]:
"""Get all function names relevant to a file (hunk headers + LLM)."""
functions: list[str] = []

query_basename = Path(file_path).name.lower()
for fs in self.file_summaries:
summary_basename = Path(fs.file_path).name.lower()
if query_basename == summary_basename:
functions.extend(fs.functions_touched)

file_intel = self.get_intel_for_file(file_path)
if file_intel:
functions.extend(file_intel.vulnerable_functions)

seen: set[str] = set()
return [f for f in functions if not (f in seen or seen.add(f))]

def to_vulnerability_intel(self) -> VulnerabilityIntel:
"""Convert to VulnerabilityIntel for backward compatibility.

Aggregates all per-file intel into the flat structure expected by
existing consumers.
"""
affected_files = list(self.per_file_intel.keys())

all_functions: list[str] = []
all_vulnerable_patterns: list[str] = []
all_fix_patterns: list[str] = []
all_keywords: list[str] = []

for intel in self.per_file_intel.values():
all_functions.extend(intel.vulnerable_functions)
all_vulnerable_patterns.extend(intel.vulnerable_patterns)
all_fix_patterns.extend(intel.fix_patterns)
all_keywords.extend(intel.search_keywords)

cve = self.cve_understanding
return VulnerabilityIntel(
affected_files=affected_files,
vulnerable_functions=list(dict.fromkeys(all_functions)),
vulnerable_patterns=all_vulnerable_patterns[:10],
fix_patterns=all_fix_patterns[:10],
search_keywords=list(dict.fromkeys(all_keywords))[:10],
Comment on lines +286 to +288

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@RedTanny Why only 10 items?? if there are more it will ignore them

vulnerability_type=cve.vulnerability_type if cve else "",
root_cause=cve.root_cause if cve else "",
component_names=list(cve.component_names) if cve else [],
affected_bitness=cve.affected_bitness if cve else "both",
affected_architectures=(
list(cve.affected_architectures)
if cve and cve.affected_architectures is not None
else None
),
)

def format_focused_context(self, grep_target: str) -> str:
"""Format intel for a specific file for inclusion in observation prompt."""
file_intel = self.get_intel_for_file(grep_target)
if not file_intel:
return ""

lines = [f"**Patch context for `{file_intel.file_path}`:**"]

functions = self.get_functions_for_file(grep_target)
if functions:
lines.append(f"- Functions modified: {', '.join(functions[:5])}")

if file_intel.vulnerable_patterns:
lines.append("- Vulnerable patterns (removed lines):")
for p in file_intel.vulnerable_patterns[:3]:
lines.append(f" - `{p[:80]}`")

if file_intel.fix_patterns:
lines.append("- Fix patterns (added lines):")
for p in file_intel.fix_patterns[:3]:
lines.append(f" - `{p[:80]}`")

return "\n".join(lines)


DEFAULT_MAX_PATTERNS = 3
DEFAULT_MIN_PATTERN_LENGTH = 10
FILE_CHANGE_PATTERN_MAX_CHARS = 100
FILE_CHANGES_THOUGHT_HEADER = (
"FILE_CHANGES (from fixed patch — orient UNPATCHED search):"
)
# Soft cap for thought prompt only; matches L1 max_iterations default.
# Full file_summaries stay in memory for focused observation context.
DEFAULT_FILE_CHANGES_THOUGHT_MAX_FILES = 10


def _file_change_sort_key(summary: FileChangeSummary) -> tuple[int, int]:
"""Sort key: largest total churn first, then highest removals (pre-move homes)."""
total = summary.removed_lines_count + summary.added_lines_count
return (total, summary.removed_lines_count)


def format_file_changes_for_thought(
summaries: list[FileChangeSummary] | None,
max_files: int | None = DEFAULT_FILE_CHANGES_THOUGHT_MAX_FILES,
) -> str:
"""Format per-file patch summaries for the L1 thought intel block.

Preserves delete/rename locality that flat VULNERABLE_PATTERNS / FIX_PATTERNS
lose (e.g. pure deletion in disk-io.c vs rewrite in extent_io.c).

Soft-caps to ``max_files`` (typically ``config.max_iterations``) after sorting by
change size so the agent sees the most actionable files first. Pass ``None`` for
no cap. Does not mutate the input list.
"""
if not summaries:
return ""

ordered = sorted(summaries, key=_file_change_sort_key, reverse=True)
omitted = 0
if max_files is not None and max_files >= 0 and len(ordered) > max_files:
omitted = len(ordered) - max_files
ordered = ordered[:max_files]

lines = [FILE_CHANGES_THOUGHT_HEADER]
for summary in ordered:
lines.append(
f" {summary.file_path}: "
f"-{summary.removed_lines_count} +{summary.added_lines_count}"
)
removed = summary.key_removed_patterns[:DEFAULT_MAX_PATTERNS]
added = summary.key_added_patterns[:DEFAULT_MAX_PATTERNS]
removed_text = (
"; ".join(p[:FILE_CHANGE_PATTERN_MAX_CHARS] for p in removed)
if removed
else "(none)"
)
added_text = (
"; ".join(p[:FILE_CHANGE_PATTERN_MAX_CHARS] for p in added)
if added
else "(none)"
)
lines.append(f" removed: {removed_text}")
lines.append(f" added: {added_text}")

if omitted:
lines.append(
f" ... +{omitted} more FILE_CHANGES omitted (cap={max_files})"
)

return "\n".join(lines)


def extract_file_change_summaries(
parsed_patch: ParsedPatch | None,
max_patterns: int = DEFAULT_MAX_PATTERNS,
min_pattern_length: int = DEFAULT_MIN_PATTERN_LENGTH,
) -> list[FileChangeSummary]:
"""Extract per-file change summaries from a parsed patch.

Args:
parsed_patch: Parsed patch structure
max_patterns: Max sample patterns per file
min_pattern_length: Min chars for a pattern to be "non-trivial"

Returns:
List of FileChangeSummary, one per file in patch
"""
if not parsed_patch:
return []

summaries = []
for pf in parsed_patch.files:
removed_lines: list[str] = []
added_lines: list[str] = []
for hunk in pf.hunks:
removed_lines.extend(hunk.removed_lines or [])
added_lines.extend(hunk.added_lines or [])

functions = [
hunk.section_header.strip()
for hunk in pf.hunks
if hunk.section_header and hunk.section_header.strip()
]

key_removed = [
line.strip()
for line in removed_lines
if len(line.strip()) >= min_pattern_length
][:max_patterns]

key_added = [
line.strip()
for line in added_lines
if len(line.strip()) >= min_pattern_length
][:max_patterns]

file_content = "\n".join(removed_lines + added_lines)
estimated_tokens = count_tokens(file_content)

summaries.append(
FileChangeSummary(
file_path=pf.clean_target_path,
is_new_file=pf.is_new_file,
is_deleted_file=pf.is_deleted_file,
removed_lines_count=len(removed_lines),
added_lines_count=len(added_lines),
hunks_count=len(pf.hunks),
functions_touched=functions,
key_removed_patterns=key_removed,
key_added_patterns=key_added,
estimated_tokens=estimated_tokens,
)
)

return summaries


def estimate_patch_tokens(parsed_patch: ParsedPatch | None) -> int:
"""Estimate total tokens in a parsed patch."""
if not parsed_patch:
return 0
summaries = extract_file_change_summaries(parsed_patch)
return sum(s.estimated_tokens for s in summaries)


class OSVPatchResult(BaseModel):
"""Result of fetching a patch from OSV/GitHub or intel references."""
cve_id: str
Expand Down Expand Up @@ -348,7 +671,7 @@ class AdvisoryContent(BaseModel):
CONFIDENCE_REVISION_FALLBACK = 0.85
CONFIDENCE_FUNCTION_MESSAGE = 0.75
CONFIDENCE_FUNCTION_PICKAXE = 0.70
CONFIDENCE_THRESHOLD_MIN = 0.50



class CommitSearchResult(BaseModel):
Expand Down
Loading