admin/trailofbits-buttercup
1
0
Fork 0
mirror of https://github.com/trailofbits/buttercup synced 2026-06-21 14:11:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Split SARIF broadcasts into a deduplicated finding pool

Browse Source 
Extract individual findings from SARIF results into a separate Redis
pool with fingerprint-based deduplication, and switch seed-gen vuln
discovery to consume findings instead of raw SARIF broadcasts. Findings
are formatted as compact XML hints in prompts, and sampling probability
now scales with pool size.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Francesco Bertolaccini
2026-04-08 09:10:55 +00:00
parent 946d5ba6d9
commit 5a354888bd
8 changed files with 453 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/src/buttercup/common/sarif_store.py
+170 -62
View File
@@ -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)
common/tests/test_sarif_store.py
+235 -50
View File
@@ -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 == []
seed-gen/src/buttercup/seed_gen/seed_gen_bot.py
+3 -3
View File
@@ -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)
seed-gen/src/buttercup/seed_gen/vuln_base_task.py
+37 -17
View File
@@ -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'<finding tool="{f.tool_name}" rule="{f.rule_id}" level="{f.level}">\n'
f' <location file="{f.file_uri}" lines="{lines}"/>\n'
f" <message>{f.message}</message>\n"
f"</finding>",
)
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"""
seed-gen/src/buttercup/seed_gen/vuln_discovery_delta.py
+3 -3
View File
@@ -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,
)
seed-gen/src/buttercup/seed_gen/vuln_discovery_full.py
+3 -3
View File
@@ -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,
)
seed-gen/test/test_vuln_discovery_delta.py
+1 -1
View File
@@ -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,
)
seed-gen/test/test_vuln_discovery_full.py
+1 -1
View File
@@ -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,
)