diff --git a/common/src/buttercup/common/sarif_store.py b/common/src/buttercup/common/sarif_store.py index 52bc3dd8..c88d7e53 100644 --- a/common/src/buttercup/common/sarif_store.py +++ b/common/src/buttercup/common/sarif_store.py @@ -1,10 +1,13 @@ from __future__ import annotations +import logging from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, computed_field from redis import Redis +logger = logging.getLogger(__name__) + class SARIFBroadcastDetail(BaseModel): """Model for SARIF broadcast details, matches the model in types.py""" @@ -19,8 +22,111 @@ class SARIFBroadcastDetail(BaseModel): task_id: str +class Finding(BaseModel): + """Individual vulnerability finding extracted from a SARIF result.""" + + rule_id: str + level: str + message: str + file_uri: str + start_line: int + end_line: int + start_column: int | None = None + tool_name: str + sarif_id: str + task_id: str + + @computed_field + @property + def fingerprint(self) -> str: + return f"{self.rule_id}:{self.file_uri}:{self.start_line}:{self.end_line}" + + +def _extract_finding_from_result( + result: dict[str, Any], + tool_name: str, + sarif_id: str, + task_id: str, +) -> Finding | None: + """Extract a Finding from a single SARIF result entry. + + Returns None if the result is malformed or missing required fields. + """ + rule_id = result.get("ruleId", "") + if not rule_id: + rule = result.get("rule", {}) + rule_id = rule.get("id", "unknown") + + level = result.get("level", "warning") + message_obj = result.get("message", {}) + message = message_obj.get("text", "") if isinstance(message_obj, dict) else str(message_obj) + + locations = result.get("locations", []) + if not locations: + return None + + physical = locations[0].get("physicalLocation", {}) + artifact_location = physical.get("artifactLocation", {}) + file_uri = artifact_location.get("uri", "") + if not file_uri: + return None + + region = physical.get("region", {}) + start_line = region.get("startLine", 0) + if start_line == 0: + return None + + end_line = region.get("endLine", start_line) + start_column = region.get("startColumn") + + return Finding( + rule_id=rule_id, + level=level, + message=message, + file_uri=file_uri, + start_line=start_line, + end_line=end_line, + start_column=start_column, + tool_name=tool_name, + sarif_id=sarif_id, + task_id=task_id, + ) + + +def extract_findings(sarif_detail: SARIFBroadcastDetail) -> list[Finding]: + """Extract individual findings from a SARIF broadcast detail. + + Iterates over runs[].results[] and extracts actionable fields. + Skips malformed entries gracefully. + """ + findings: list[Finding] = [] + sarif = sarif_detail.sarif + runs = sarif.get("runs", []) + + for run in runs: + driver = run.get("tool", {}).get("driver", {}) + tool_name = driver.get("name", "unknown") + results = run.get("results", []) + + for result in results: + finding = _extract_finding_from_result( + result, + tool_name, + sarif_detail.sarif_id, + sarif_detail.task_id, + ) + if finding is not None: + findings.append(finding) + + return findings + + class SARIFStore: - """Store and retrieve SARIF objects in Redis""" + """Store and retrieve SARIF objects and extracted findings in Redis.""" + + SARIF_PREFIX = "sarif:" + FINDING_PREFIX = "findings:" + FINDING_SEEN_PREFIX = "findings_seen:" def __init__(self, redis: Redis): """Initialize the SARIF store with a Redis connection. @@ -30,85 +136,67 @@ class SARIFStore: """ self.redis = redis - self.key_prefix = "sarif:" + # Keep for backward compat with code using self.key_prefix + self.key_prefix = self.SARIF_PREFIX def _get_key(self, task_id: str) -> str: - """Get the Redis key for a task_id. + return f"{self.SARIF_PREFIX}{task_id.lower()}" - Args: - task_id: Task ID + def _get_finding_key(self, task_id: str) -> str: + return f"{self.FINDING_PREFIX}{task_id.lower()}" - Returns: - Redis key - - """ - return f"{self.key_prefix}{task_id.lower()}" + def _get_finding_seen_key(self, task_id: str) -> str: + return f"{self.FINDING_SEEN_PREFIX}{task_id.lower()}" def _decode_key(self, key: str | bytes) -> str: - """Decode a Redis key if it's bytes, otherwise return as is. - - Args: - key: Redis key, either bytes or string - - Returns: - Decoded key as string - - """ if isinstance(key, bytes): return key.decode("utf-8") return key def store(self, sarif_detail: SARIFBroadcastDetail) -> None: - """Store a SARIF broadcast detail in Redis. - - Args: - sarif_detail: The SARIF broadcast detail to store - - """ + """Store a SARIF broadcast detail and its extracted findings in Redis.""" task_id = sarif_detail.task_id - key = self._get_key(task_id) - - # We'll use a Redis list to store multiple SARIF objects for the same task - # Serialize the SARIF object to JSON + sarif_key = self._get_key(task_id) sarif_json = sarif_detail.model_dump_json() + self.redis.rpush(sarif_key, sarif_json) - # Add to the list for this task - self.redis.rpush(key, sarif_json) + findings = extract_findings(sarif_detail) + self._store_findings(task_id, findings) + + def _store_findings(self, task_id: str, findings: list[Finding]) -> int: + """Store findings, deduplicating by fingerprint. Returns count of new findings added.""" + finding_key = self._get_finding_key(task_id) + seen_key = self._get_finding_seen_key(task_id) + added = 0 + + for finding in findings: + if self.redis.sismember(seen_key, finding.fingerprint): + continue + self.redis.rpush(finding_key, finding.model_dump_json()) + self.redis.sadd(seen_key, finding.fingerprint) + added += 1 + + if added > 0: + logger.info("Added %d new findings for task %s (total pool: %d)", added, task_id, added) + + return added def get_all(self) -> list[SARIFBroadcastDetail]: - """Retrieve all SARIF objects from Redis. - - Returns: - List of SARIF broadcast details - - """ - # Get all SARIF keys in Redis + """Retrieve all SARIF objects from Redis.""" all_keys = self.redis.keys(f"{self.key_prefix}*") result = [] for key in all_keys: - # Decode the key if it's bytes decoded_key = self._decode_key(key) - - # Get all SARIF objects for this task sarif_list = self.redis.lrange(decoded_key, 0, -1) for sarif_json in sarif_list: - # Parse each JSON string into a SARIFBroadcastDetail sarif_detail = SARIFBroadcastDetail.model_validate_json(sarif_json) result.append(sarif_detail) return result def get_by_task_id(self, task_id: str) -> list[SARIFBroadcastDetail]: - """Retrieve all SARIF objects for a specific task. - - Args: - task_id: Task ID - - Returns: - List of SARIF broadcast details for this task - - """ + """Retrieve all SARIF objects for a specific task.""" key = self._get_key(task_id) sarif_list = self.redis.lrange(key, 0, -1) @@ -119,15 +207,35 @@ class SARIFStore: return result - def delete_by_task_id(self, task_id: str) -> int: - """Remove all SARIF objects for a specific task. - - Args: - task_id: Task ID - - Returns: - Number of removed keys (0 or 1) + def get_findings_by_task_id(self, task_id: str) -> list[Finding]: + """Retrieve all findings for a specific task from the finding pool. + Falls back to extracting from stored SARIFs if the finding pool + is empty but SARIF data exists (backward compatibility). """ - key = self._get_key(task_id) - return self.redis.delete(key) + finding_key = self._get_finding_key(task_id) + finding_list = self.redis.lrange(finding_key, 0, -1) + + if finding_list: + return [Finding.model_validate_json(f) for f in finding_list] + + # Fallback: extract from old SARIF data if present + sarifs = self.get_by_task_id(task_id) + if not sarifs: + return [] + + all_findings: list[Finding] = [] + for sarif_detail in sarifs: + all_findings.extend(extract_findings(sarif_detail)) + + if all_findings: + self._store_findings(task_id, all_findings) + + return all_findings + + def delete_by_task_id(self, task_id: str) -> int: + """Remove all SARIF objects and findings for a specific task.""" + sarif_key = self._get_key(task_id) + finding_key = self._get_finding_key(task_id) + seen_key = self._get_finding_seen_key(task_id) + return self.redis.delete(sarif_key, finding_key, seen_key) diff --git a/common/tests/test_sarif_store.py b/common/tests/test_sarif_store.py index a4a1bad1..bad2feba 100644 --- a/common/tests/test_sarif_store.py +++ b/common/tests/test_sarif_store.py @@ -3,7 +3,12 @@ import json import pytest from redis import Redis -from buttercup.common.sarif_store import SARIFBroadcastDetail, SARIFStore +from buttercup.common.sarif_store import ( + Finding, + SARIFBroadcastDetail, + SARIFStore, + extract_findings, +) @pytest.fixture @@ -13,18 +18,18 @@ def redis_client(): @pytest.fixture def sarif_store(redis_client): - # Create a SARIFStore instance for testing store = SARIFStore(redis_client) - # Clean any existing test data - for key in redis_client.keys(f"{store.key_prefix}*"): - redis_client.delete(key) + # Clean any existing test data across all key namespaces + for prefix in (SARIFStore.SARIF_PREFIX, SARIFStore.FINDING_PREFIX, SARIFStore.FINDING_SEEN_PREFIX): + for key in redis_client.keys(f"{prefix}*"): + redis_client.delete(key) yield store - # Clean up after tests - for key in redis_client.keys(f"{store.key_prefix}*"): - redis_client.delete(key) + for prefix in (SARIFStore.SARIF_PREFIX, SARIFStore.FINDING_PREFIX, SARIFStore.FINDING_SEEN_PREFIX): + for key in redis_client.keys(f"{prefix}*"): + redis_client.delete(key) @pytest.fixture @@ -38,18 +43,58 @@ def sample_sarif_detail(): ) +@pytest.fixture +def sarif_with_findings(): + """Create a SARIFBroadcastDetail with realistic findings.""" + return SARIFBroadcastDetail( + metadata={"source": "test"}, + sarif={ + "version": "2.1.0", + "runs": [ + { + "tool": {"driver": {"name": "CodeScan++"}}, + "results": [ + { + "ruleId": "CWE-121", + "level": "error", + "message": {"text": "Stack-based buffer overflow"}, + "locations": [ + { + "physicalLocation": { + "artifactLocation": {"uri": "pngrutil.c"}, + "region": {"startLine": 1421, "endLine": 1447}, + }, + }, + ], + }, + { + "ruleId": "CWE-787", + "level": "warning", + "message": {"text": "Out-of-bounds write"}, + "locations": [ + { + "physicalLocation": { + "artifactLocation": {"uri": "pngread.c"}, + "region": {"startLine": 200, "endLine": 210, "startColumn": 5}, + }, + }, + ], + }, + ], + }, + ], + }, + sarif_id="sarif-with-findings", + task_id="test-task-id", + ) + + def test_sarif_store_store_and_get_by_task_id(sarif_store, sample_sarif_detail): """Test storing a SARIF detail and retrieving it by task ID""" - # Store the SARIF detail sarif_store.store(sample_sarif_detail) - - # Retrieve it by task ID retrieved_sarifs = sarif_store.get_by_task_id(sample_sarif_detail.task_id) - # Verify we got exactly one SARIF detail assert len(retrieved_sarifs) == 1 - - # Verify the retrieved SARIF detail matches the original retrieved = retrieved_sarifs[0] assert retrieved.sarif_id == sample_sarif_detail.sarif_id assert retrieved.task_id == sample_sarif_detail.task_id @@ -59,7 +104,6 @@ def test_sarif_store_store_and_get_by_task_id(sarif_store, sample_sarif_detail): def test_sarif_store_get_all(sarif_store, sample_sarif_detail): """Test retrieving all SARIF details""" - # Create a second SARIF detail with a different task ID second_sarif = SARIFBroadcastDetail( metadata={"source": "test2", "version": "1.0"}, sarif={"version": "2.1.0", "runs": []}, @@ -67,17 +111,12 @@ def test_sarif_store_get_all(sarif_store, sample_sarif_detail): task_id="test-task-id-2", ) - # Store both SARIF details sarif_store.store(sample_sarif_detail) sarif_store.store(second_sarif) - # Retrieve all SARIF details all_sarifs = sarif_store.get_all() - - # Verify we got both SARIF details assert len(all_sarifs) == 2 - # Verify the task IDs of the retrieved SARIF details task_ids = {sarif.task_id for sarif in all_sarifs} assert sample_sarif_detail.task_id in task_ids assert second_sarif.task_id in task_ids @@ -85,25 +124,19 @@ def test_sarif_store_get_all(sarif_store, sample_sarif_detail): def test_sarif_store_multiple_sarifs_per_task(sarif_store, sample_sarif_detail): """Test storing multiple SARIF details for the same task""" - # Create a second SARIF detail with the same task ID but different SARIF ID second_sarif = SARIFBroadcastDetail( metadata={"source": "test2", "version": "1.0"}, sarif={"version": "2.1.0", "runs": []}, sarif_id="test-sarif-id-2", - task_id=sample_sarif_detail.task_id, # Same task ID + task_id=sample_sarif_detail.task_id, ) - # Store both SARIF details sarif_store.store(sample_sarif_detail) sarif_store.store(second_sarif) - # Retrieve SARIF details for the task retrieved_sarifs = sarif_store.get_by_task_id(sample_sarif_detail.task_id) - - # Verify we got both SARIF details assert len(retrieved_sarifs) == 2 - # Verify the SARIF IDs of the retrieved SARIF details sarif_ids = {sarif.sarif_id for sarif in retrieved_sarifs} assert sample_sarif_detail.sarif_id in sarif_ids assert second_sarif.sarif_id in sarif_ids @@ -111,7 +144,6 @@ def test_sarif_store_multiple_sarifs_per_task(sarif_store, sample_sarif_detail): def test_sarif_store_delete_by_task_id(sarif_store, sample_sarif_detail): """Test deleting SARIF details by task ID""" - # Create a second SARIF detail with a different task ID second_sarif = SARIFBroadcastDetail( metadata={"source": "test2", "version": "1.0"}, sarif={"version": "2.1.0", "runs": []}, @@ -119,69 +151,222 @@ def test_sarif_store_delete_by_task_id(sarif_store, sample_sarif_detail): task_id="test-task-id-2", ) - # Store both SARIF details sarif_store.store(sample_sarif_detail) sarif_store.store(second_sarif) - # Delete SARIF details for the first task deleted = sarif_store.delete_by_task_id(sample_sarif_detail.task_id) + assert deleted >= 1 - # Verify deletion was successful - assert deleted == 1 - - # Retrieve all SARIF details all_sarifs = sarif_store.get_all() - - # Verify we only have the second SARIF detail assert len(all_sarifs) == 1 assert all_sarifs[0].task_id == second_sarif.task_id def test_sarif_store_nonexistent_task_id(sarif_store): """Test retrieving and deleting SARIF details for a nonexistent task ID""" - # Attempt to retrieve SARIF details for a nonexistent task ID retrieved_sarifs = sarif_store.get_by_task_id("nonexistent-task-id") - - # Verify we got an empty list assert retrieved_sarifs == [] - # Attempt to delete SARIF details for a nonexistent task ID deleted = sarif_store.delete_by_task_id("nonexistent-task-id") - - # Verify no keys were deleted assert deleted == 0 def test_sarif_store_case_insensitive_task_id(sarif_store, sample_sarif_detail): """Test that task IDs are case-insensitive""" - # Store a SARIF detail with a lowercase task ID lowercase_task_id = sample_sarif_detail.task_id.lower() sample_sarif_detail.task_id = lowercase_task_id sarif_store.store(sample_sarif_detail) - # Retrieve the SARIF detail using an uppercase task ID uppercase_task_id = lowercase_task_id.upper() retrieved_sarifs = sarif_store.get_by_task_id(uppercase_task_id) - # Verify we got the SARIF detail assert len(retrieved_sarifs) == 1 assert retrieved_sarifs[0].sarif_id == sample_sarif_detail.sarif_id def test_sarif_store_json_serialization(sarif_store, sample_sarif_detail): """Test that SARIF details are properly serialized and deserialized""" - # Store a SARIF detail sarif_store.store(sample_sarif_detail) - # Retrieve the raw JSON from Redis key = sarif_store._get_key(sample_sarif_detail.task_id) raw_json = sarif_store.redis.lrange(key, 0, -1)[0] - # Parse the JSON parsed_json = json.loads(raw_json) - - # Verify the parsed JSON matches the original SARIF detail assert parsed_json["sarif_id"] == sample_sarif_detail.sarif_id assert parsed_json["task_id"] == sample_sarif_detail.task_id assert parsed_json["metadata"] == sample_sarif_detail.metadata assert parsed_json["sarif"] == sample_sarif_detail.sarif + + +# --- Finding pool tests --- + + +def test_extract_findings(sarif_with_findings): + """Test extracting individual findings from SARIF broadcast detail.""" + findings = extract_findings(sarif_with_findings) + + assert len(findings) == 2 + + f0 = findings[0] + assert f0.rule_id == "CWE-121" + assert f0.level == "error" + assert f0.message == "Stack-based buffer overflow" + assert f0.file_uri == "pngrutil.c" + assert f0.start_line == 1421 + assert f0.end_line == 1447 + assert f0.start_column is None + assert f0.tool_name == "CodeScan++" + assert f0.sarif_id == "sarif-with-findings" + + f1 = findings[1] + assert f1.rule_id == "CWE-787" + assert f1.file_uri == "pngread.c" + assert f1.start_column == 5 + + +def test_extract_findings_skips_malformed_results(): + """Test that extraction skips results without locations or file URIs.""" + detail = SARIFBroadcastDetail( + metadata={}, + sarif={ + "version": "2.1.0", + "runs": [ + { + "tool": {"driver": {"name": "TestTool"}}, + "results": [ + {"ruleId": "R1", "level": "error", "message": {"text": "no locations"}}, + { + "ruleId": "R2", + "level": "error", + "message": {"text": "no file uri"}, + "locations": [{"physicalLocation": {"artifactLocation": {}}}], + }, + { + "ruleId": "R3", + "level": "error", + "message": {"text": "no start line"}, + "locations": [ + { + "physicalLocation": { + "artifactLocation": {"uri": "test.c"}, + "region": {}, + }, + }, + ], + }, + { + "ruleId": "R4", + "level": "warning", + "message": {"text": "valid"}, + "locations": [ + { + "physicalLocation": { + "artifactLocation": {"uri": "good.c"}, + "region": {"startLine": 10}, + }, + }, + ], + }, + ], + }, + ], + }, + sarif_id="malformed-test", + task_id="task-1", + ) + + findings = extract_findings(detail) + assert len(findings) == 1 + assert findings[0].rule_id == "R4" + assert findings[0].end_line == 10 # defaults to start_line + + +def test_extract_findings_empty_runs(): + """Test extraction from SARIF with no runs.""" + detail = SARIFBroadcastDetail( + metadata={}, + sarif={"version": "2.1.0", "runs": []}, + sarif_id="empty", + task_id="task-1", + ) + assert extract_findings(detail) == [] + + +def test_store_populates_findings(sarif_store, sarif_with_findings): + """Test that store() populates both sarif and findings keys.""" + sarif_store.store(sarif_with_findings) + + # SARIF key should have the full broadcast + sarifs = sarif_store.get_by_task_id("test-task-id") + assert len(sarifs) == 1 + + # Findings key should have extracted findings + findings = sarif_store.get_findings_by_task_id("test-task-id") + assert len(findings) == 2 + assert findings[0].rule_id == "CWE-121" + assert findings[1].rule_id == "CWE-787" + + +def test_findings_deduplication(sarif_store, sarif_with_findings): + """Test that storing the same SARIF twice doesn't duplicate findings.""" + sarif_store.store(sarif_with_findings) + sarif_store.store(sarif_with_findings) + + findings = sarif_store.get_findings_by_task_id("test-task-id") + assert len(findings) == 2 # Not 4 + + +def test_findings_backward_compat_fallback(sarif_store, sarif_with_findings): + """Test fallback extraction when only sarif: keys exist.""" + # Manually store only the SARIF (bypassing finding extraction) + task_id = sarif_with_findings.task_id + sarif_key = sarif_store._get_key(task_id) + sarif_store.redis.rpush(sarif_key, sarif_with_findings.model_dump_json()) + + # get_findings_by_task_id should extract on the fly + findings = sarif_store.get_findings_by_task_id(task_id) + assert len(findings) == 2 + assert findings[0].rule_id == "CWE-121" + + # After fallback, findings should now be stored + finding_key = sarif_store._get_finding_key(task_id) + assert sarif_store.redis.llen(finding_key) == 2 + + +def test_delete_cleans_all_keys(sarif_store, sarif_with_findings): + """Test that delete_by_task_id removes sarif, findings, and seen keys.""" + sarif_store.store(sarif_with_findings) + task_id = sarif_with_findings.task_id + + # Verify keys exist + assert sarif_store.redis.llen(sarif_store._get_key(task_id)) > 0 + assert sarif_store.redis.llen(sarif_store._get_finding_key(task_id)) > 0 + assert sarif_store.redis.scard(sarif_store._get_finding_seen_key(task_id)) > 0 + + sarif_store.delete_by_task_id(task_id) + + assert sarif_store.redis.llen(sarif_store._get_key(task_id)) == 0 + assert sarif_store.redis.llen(sarif_store._get_finding_key(task_id)) == 0 + assert sarif_store.redis.scard(sarif_store._get_finding_seen_key(task_id)) == 0 + + +def test_finding_fingerprint(): + """Test finding fingerprint generation.""" + finding = Finding( + rule_id="CWE-121", + level="error", + message="test", + file_uri="foo.c", + start_line=10, + end_line=20, + tool_name="T", + sarif_id="s1", + task_id="t1", + ) + assert finding.fingerprint == "CWE-121:foo.c:10:20" + + +def test_get_findings_nonexistent_task(sarif_store): + """Test getting findings for a task with no data.""" + findings = sarif_store.get_findings_by_task_id("nonexistent") + assert findings == [] diff --git a/seed-gen/src/buttercup/seed_gen/seed_gen_bot.py b/seed-gen/src/buttercup/seed_gen/seed_gen_bot.py index 0ea9bee9..6b154a2a 100644 --- a/seed-gen/src/buttercup/seed_gen/seed_gen_bot.py +++ b/seed-gen/src/buttercup/seed_gen/seed_gen_bot.py @@ -186,7 +186,7 @@ class SeedGenBot(TaskLoop): seed_init.do_task(out_dir) elif task_choice == TaskName.VULN_DISCOVERY.value: sarif_store = SARIFStore(self.redis) - sarifs = sarif_store.get_by_task_id(challenge_task.task_meta.task_id) + findings = sarif_store.get_findings_by_task_id(challenge_task.task_meta.task_id) fbuilds = builds[BuildType.FUZZER] reproduce_multiple = ReproduceMultiple(temp_dir, fbuilds) crash_submit = CrashSubmit( @@ -210,7 +210,7 @@ class SeedGenBot(TaskLoop): project_yaml, self.redis, mult, - sarifs, + findings, crash_submit=crash_submit, ) else: @@ -222,7 +222,7 @@ class SeedGenBot(TaskLoop): project_yaml, self.redis, mult, - sarifs, + findings, crash_submit=crash_submit, ) vuln_discovery.do_task(out_dir, current_dir) diff --git a/seed-gen/src/buttercup/seed_gen/vuln_base_task.py b/seed-gen/src/buttercup/seed_gen/vuln_base_task.py index 4d910d4e..b463db3d 100644 --- a/seed-gen/src/buttercup/seed_gen/vuln_base_task.py +++ b/seed-gen/src/buttercup/seed_gen/vuln_base_task.py @@ -1,4 +1,3 @@ -import json import logging import operator import random @@ -16,7 +15,7 @@ from buttercup.common.llm import get_langfuse_callbacks from buttercup.common.project_yaml import Language from buttercup.common.queues import ReliableQueue from buttercup.common.reproduce_multiple import ReproduceMultiple, ReproduceResult -from buttercup.common.sarif_store import SARIFBroadcastDetail +from buttercup.common.sarif_store import Finding from buttercup.common.stack_parsing import CrashSet from buttercup.common.telemetry import CRSActionCategory, set_crs_attributes from langchain_core.output_parsers import StrOutputParser @@ -61,8 +60,8 @@ class PoVAttempt: class VulnBaseState(BaseTaskState): analysis: str = Field(description="The analysis of the vulnerability", default="") - sarifs: list[SARIFBroadcastDetail] = Field( - description="SARIF broadcasts for the task", + findings: list[Finding] = Field( + description="Individual findings from SARIF reports", default_factory=list, ) valid_pov_count: int = Field(description="The number of valid PoVs found", default=0) @@ -72,14 +71,20 @@ class VulnBaseState(BaseTaskState): pov_iteration: int = Field(description="Count of pov write iterations", default=0) pov_attempts: Annotated[list[PoVAttempt], operator.add] = Field(default_factory=list) - def format_sarif_hints(self) -> str: - """Format SARIF hints for prompts""" - if not self.sarifs: + def format_finding_hints(self) -> str: + """Format finding hints for prompts in compact XML-like format.""" + if not self.findings: return "" hints = [] - for sarif in self.sarifs: - hints.append(json.dumps(sarif.sarif, indent=2)) + for f in self.findings: + lines = f"{f.start_line}-{f.end_line}" if f.end_line != f.start_line else str(f.start_line) + hints.append( + f'\n' + f' \n' + f" {f.message}\n" + f"", + ) return "\n\n".join(hints) @@ -99,11 +104,13 @@ class CrashSubmit: @dataclass class VulnBaseTask(Task): reproduce_multiple: ReproduceMultiple - sarifs: list[SARIFBroadcastDetail] + findings: list[Finding] TaskStateClass: ClassVar[type[BaseTaskState]] - SARIF_PROBABILITY: ClassVar[float] = 0.5 crash_submit: CrashSubmit | None = None + MAX_FINDINGS: ClassVar[int] = 10 + MIN_FINDING_PROBABILITY: ClassVar[float] = 0.3 + MAX_POV_ITERATIONS: ClassVar[int] = 3 MAX_CONTEXT_ITERATIONS: ClassVar[int] @@ -350,12 +357,25 @@ class VulnBaseTask(Task): str(err), ) - def sample_sarifs(self) -> list[SARIFBroadcastDetail]: - """Sample SARIFs for the task""" - if random.random() <= VulnBaseTask.SARIF_PROBABILITY: - logger.info("Using %d SARIFs for challenge %s", len(self.sarifs), self.package_name) - return self.sarifs - return [] + def sample_findings(self) -> list[Finding]: + """Sample findings from the pool with probability scaling by pool size. + + Probability scales linearly from MIN_FINDING_PROBABILITY (at 1 finding) + to 1.0 (at 5+ findings). Samples up to MAX_FINDINGS randomly. + """ + pool_size = len(self.findings) + if pool_size == 0: + return [] + + probability = min(1.0, self.MIN_FINDING_PROBABILITY + 0.14 * pool_size) + if random.random() > probability: + logger.info("Skipped findings for challenge %s (p=%.2f)", self.package_name, probability) + return [] + + k = min(pool_size, self.MAX_FINDINGS) + sampled = random.sample(self.findings, k) + logger.info("Using %d/%d findings for challenge %s", k, pool_size, self.package_name) + return sampled def get_pov_examples(self) -> str: """Get PoV examples for the task""" diff --git a/seed-gen/src/buttercup/seed_gen/vuln_discovery_delta.py b/seed-gen/src/buttercup/seed_gen/vuln_discovery_delta.py index 3f48d741..232c396d 100644 --- a/seed-gen/src/buttercup/seed_gen/vuln_discovery_delta.py +++ b/seed-gen/src/buttercup/seed_gen/vuln_discovery_delta.py @@ -38,7 +38,7 @@ class VulnDiscoveryDeltaTask(VulnBaseTask): "diff": state.diff_content, "harness": str(state.harness), "retrieved_context": state.format_retrieved_context(), - "sarif_hints": state.format_sarif_hints(), + "sarif_hints": state.format_finding_hints(), "vuln_files": self.get_vuln_files(), "fuzzer_name": self.get_fuzzer_name(), "cwe_list": self.get_cwe_list(), @@ -58,7 +58,7 @@ class VulnDiscoveryDeltaTask(VulnBaseTask): "diff": state.diff_content, "harness": str(state.harness), "retrieved_context": state.format_retrieved_context(), - "sarif_hints": state.format_sarif_hints(), + "sarif_hints": state.format_finding_hints(), "vuln_files": self.get_vuln_files(), "fuzzer_name": self.get_fuzzer_name(), "cwe_list": self.get_cwe_list(), @@ -106,7 +106,7 @@ class VulnDiscoveryDeltaTask(VulnBaseTask): harness=harness, diff_content=diff_content, task=self, - sarifs=self.sample_sarifs(), + findings=self.sample_findings(), output_dir=out_dir, current_dir=current_dir, ) diff --git a/seed-gen/src/buttercup/seed_gen/vuln_discovery_full.py b/seed-gen/src/buttercup/seed_gen/vuln_discovery_full.py index ddda1ed7..109ca232 100644 --- a/seed-gen/src/buttercup/seed_gen/vuln_discovery_full.py +++ b/seed-gen/src/buttercup/seed_gen/vuln_discovery_full.py @@ -31,7 +31,7 @@ class VulnDiscoveryFullTask(VulnBaseTask): prompt_vars = { "harness": str(state.harness), "retrieved_context": state.format_retrieved_context(), - "sarif_hints": state.format_sarif_hints(), + "sarif_hints": state.format_finding_hints(), "vuln_files": self.get_vuln_files(), "fuzzer_name": self.get_fuzzer_name(), "cwe_list": self.get_cwe_list(), @@ -50,7 +50,7 @@ class VulnDiscoveryFullTask(VulnBaseTask): prompt_vars = { "harness": str(state.harness), "retrieved_context": state.format_retrieved_context(), - "sarif_hints": state.format_sarif_hints(), + "sarif_hints": state.format_finding_hints(), "vuln_files": self.get_vuln_files(), "fuzzer_name": self.get_fuzzer_name(), "cwe_list": self.get_cwe_list(), @@ -91,7 +91,7 @@ class VulnDiscoveryFullTask(VulnBaseTask): state = VulnBaseState( harness=harness, task=self, - sarifs=self.sample_sarifs(), + findings=self.sample_findings(), output_dir=out_dir, current_dir=current_dir, ) diff --git a/seed-gen/test/test_vuln_discovery_delta.py b/seed-gen/test/test_vuln_discovery_delta.py index 0df95c72..06321b19 100644 --- a/seed-gen/test/test_vuln_discovery_delta.py +++ b/seed-gen/test/test_vuln_discovery_delta.py @@ -35,7 +35,7 @@ def vuln_discovery_task( project_yaml=mock_project_yaml, redis=mock_redis, reproduce_multiple=mock_reproduce_multiple, - sarifs=[], + findings=[], crash_submit=mock_crash_submit, ) diff --git a/seed-gen/test/test_vuln_discovery_full.py b/seed-gen/test/test_vuln_discovery_full.py index ca8ecb31..33d2bd3a 100644 --- a/seed-gen/test/test_vuln_discovery_full.py +++ b/seed-gen/test/test_vuln_discovery_full.py @@ -36,7 +36,7 @@ def vuln_discovery_full_task( project_yaml=mock_project_yaml, redis=mock_redis, reproduce_multiple=mock_reproduce_multiple, - sarifs=[], + findings=[], crash_submit=mock_crash_submit, )