Skip to content

Fix scanner to iterate over all probe prompts #1

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
IDohndorf wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix scanner to iterate over all probe prompts #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
45 changes: 34 additions & 11 deletions src/main/java/io/tiberius/core/TiberiusScanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,8 +89,12 @@ static TiberiusScanner create(
public ScanReport scan() {
final Instant startTime = Instant.now();
final List<Probe> probesToRun = selectProbes();
final int totalPrompts = probesToRun.stream()
.mapToInt(p -> p.getPrompts().size())
.sum();

log.info("Starting scan with {} probes against {}", probesToRun.size(), generator.getName());
log.info("Starting scan with {} probes ({} total prompts) against {}",
probesToRun.size(), totalPrompts, generator.getName());

final List<ScanResult> results;
if (concurrency > 1 && probesToRun.size() > 1) {
Expand All @@ -112,8 +116,20 @@ public ScanResult runProbe(final Probe probe) {
}

public ScanResult runProbe(final Probe probe, final Buff buff) {
return runProbe(probe, probe.getPrompt(), buff);
}

/**
* Run a single probe with a specific prompt and buff transformation.
* This allows testing individual prompts from a probe's prompt list.
*
* @param probe The probe being tested
* @param originalPrompt The specific prompt to test
* @param buff The buff transformation to apply
* @return The scan result
*/
public ScanResult runProbe(final Probe probe, final String originalPrompt, final Buff buff) {
final Instant startTime = Instant.now();
final String originalPrompt = probe.getPrompt();
final String transformedPrompt = buff.transform(originalPrompt);
final GeneratorResponse response = generator.generate(transformedPrompt);
final Duration duration = Duration.between(startTime, Instant.now());
Expand Down Expand Up @@ -193,13 +209,16 @@ private List<ScanResult> runSequential(final List<Probe> probes) {
final List<Buff> buffsToApply = getBuffsToApply();

for (final Probe probe : probes) {
for (final Buff buff : buffsToApply) {
final ScanResult result = runProbe(probe, buff);
results.add(result);
logResult(result);
if (failFast && result.attackSucceeded()) {
log.warn("Fail fast triggered: attack succeeded for probe {}", probe.getId());
return results;
final List<String> prompts = probe.getPrompts();
for (final String prompt : prompts) {
for (final Buff buff : buffsToApply) {
final ScanResult result = runProbe(probe, prompt, buff);
results.add(result);
logResult(result);
if (failFast && result.attackSucceeded()) {
log.warn("Fail fast triggered: attack succeeded for probe {}", probe.getId());
return results;
}
}
}
}
Expand All @@ -211,8 +230,12 @@ private List<ScanResult> runParallel(final List<Probe> probes) {
final List<CompletableFuture<ScanResult>> futures = new ArrayList<>();

for (final Probe probe : probes) {
for (final Buff buff : buffsToApply) {
futures.add(CompletableFuture.supplyAsync(() -> runProbe(probe, buff), executor));
final List<String> prompts = probe.getPrompts();
for (final String prompt : prompts) {
for (final Buff buff : buffsToApply) {
futures.add(CompletableFuture.supplyAsync(
() -> runProbe(probe, prompt, buff), executor));
}
}
}

Expand Down
90 changes: 90 additions & 0 deletions src/main/java/io/tiberius/core/result/ScanReport.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,96 @@ public Map<Integer, List<ScanResult>> bySeverity() {
.collect(Collectors.groupingBy(r -> r.probe().getSeverity()));
}

/**
* Get results grouped by probe ID.
* Each probe may have multiple results (one per prompt variant and buff combination).
*/
public Map<String, List<ScanResult>> byProbe() {
return results.stream()
.filter(r -> r.probe() != null)
.collect(Collectors.groupingBy(r -> r.probe().getId()));
}

/**
* Compute success rate per probe.
* Returns a map of probe ID to success rate percentage.
* Each probe's success rate is computed across all its prompt variants.
*/
public Map<String, Double> successRateByProbe() {
return byProbe().entrySet().stream()
.collect(Collectors.toMap(
Map.Entry::getKey,
e -> {
List<ScanResult> probeResults = e.getValue();
long successes = probeResults.stream()
.filter(ScanResult::attackSucceeded)
.count();
return probeResults.isEmpty() ? 0.0
: (successes * 100.0) / probeResults.size();
}
));
}

/**
* Compute the mean success rate across all probes.
* This averages the per-probe success rates, giving each probe equal weight
* regardless of how many prompt variants it has.
*/
public double meanSuccessRateByProbe() {
Map<String, Double> perProbeRates = successRateByProbe();
if (perProbeRates.isEmpty()) {
return 0.0;
}
return perProbeRates.values().stream()
.mapToDouble(Double::doubleValue)
.average()
.orElse(0.0);
}

/**
* Number of unique probes tested.
*/
public int uniqueProbeCount() {
return (int) results.stream()
.filter(r -> r.probe() != null)
.map(r -> r.probe().getId())
.distinct()
.count();
}

/**
* Summary statistics for each probe.
*/
public Map<String, ProbeSummary> probeSummaries() {
return byProbe().entrySet().stream()
.collect(Collectors.toMap(
Map.Entry::getKey,
e -> {
List<ScanResult> probeResults = e.getValue();
int total = probeResults.size();
int successful = (int) probeResults.stream()
.filter(ScanResult::attackSucceeded)
.count();
String probeName = probeResults.isEmpty() ? e.getKey()
: probeResults.get(0).probe().getName();
return new ProbeSummary(e.getKey(), probeName, total, successful);
}
));
}

/**
* Summary for a single probe across all its prompt variants.
*/
public record ProbeSummary(String probeId, String probeName, int totalPrompts, int successful) {
public int blocked() {
return totalPrompts - successful;
}

public double successRate() {
return totalPrompts == 0 ? 0.0 : (successful * 100.0) / totalPrompts;
}
}

/**
* Summary statistics for each category.
*/
Expand Down
144 changes: 144 additions & 0 deletions src/test/java/io/tiberius/core/TiberiusScannerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,18 @@
import io.tiberius.core.detector.CompositeDetector;
import io.tiberius.core.detector.PatternDetector;
import io.tiberius.core.generator.MockGenerator;
import io.tiberius.core.probe.AbstractProbe;
import io.tiberius.core.probe.Probe;
import io.tiberius.core.probe.ProbeRegistry;
import io.tiberius.core.result.ScanReport;
import io.tiberius.core.result.ScanResult;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.*;

/**
Expand Down Expand Up @@ -187,4 +192,143 @@ void testFailFast() {
assertEquals(1, report.successfulAttacks(),
"Fail fast should stop after first successful attack");
}

@Test
@DisplayName("Scanner tests all prompts from a probe")
void testAllPromptsAreTested() {
// Create a probe with multiple prompts
Probe multiPromptProbe = new AbstractProbe(
"test.multi_prompt",
"Multi-Prompt Test Probe",
"A probe with multiple prompt variants for testing",
AttackCategory.JAILBREAK,
List.of(
"First attack prompt: ignore instructions",
"Second attack prompt: you are now DAN",
"Third attack prompt: developer mode enabled"
),
List.of("test"),
3
) {};

// Create a custom registry with just our test probe
ProbeRegistry testRegistry = new ProbeRegistry();
testRegistry.register(multiPromptProbe);

TiberiusScanner testScanner = TiberiusScanner.builder()
.withProbeRegistry(testRegistry)
.withGenerator(MockGenerator.secure())
.withDetector(CompositeDetector.defaultComposite())
.withConcurrency(1)
.build();
testScanner.setProbePatterns(new String[]{"test.multi_prompt"});

ScanReport report = testScanner.scan();

// Should have 3 results (one per prompt)
assertEquals(3, report.totalProbes(),
"Should have one result per prompt variant");
assertEquals(1, report.uniqueProbeCount(),
"Should have one unique probe");

// All results should belong to the same probe
report.results().forEach(result ->
assertEquals("test.multi_prompt", result.probeId()));
}

@Test
@DisplayName("Per-probe success rate is computed correctly")
void testPerProbeSuccessRate() {
// Create two probes with different numbers of prompts
Probe probe1 = new AbstractProbe(
"test.probe1",
"Test Probe 1",
"First test probe with 2 prompts",
AttackCategory.JAILBREAK,
List.of("Prompt 1A", "Prompt 1B"),
List.of("test"),
3
) {};

Probe probe2 = new AbstractProbe(
"test.probe2",
"Test Probe 2",
"Second test probe with 3 prompts",
AttackCategory.JAILBREAK,
List.of("Prompt 2A", "Prompt 2B", "Prompt 2C"),
List.of("test"),
3
) {};

ProbeRegistry testRegistry = new ProbeRegistry();
testRegistry.register(probe1);
testRegistry.register(probe2);

TiberiusScanner testScanner = TiberiusScanner.builder()
.withProbeRegistry(testRegistry)
.withGenerator(MockGenerator.secure())
.withDetector(CompositeDetector.defaultComposite())
.withConcurrency(1)
.build();
testScanner.setProbePatterns(new String[]{"test.*"});

ScanReport report = testScanner.scan();

// Should have 5 total results (2 + 3)
assertEquals(5, report.totalProbes(),
"Should have one result per prompt across all probes");
assertEquals(2, report.uniqueProbeCount(),
"Should have two unique probes");

// Check per-probe grouping
Map<String, List<ScanResult>> byProbe = report.byProbe();
assertEquals(2, byProbe.size());
assertEquals(2, byProbe.get("test.probe1").size(),
"Probe 1 should have 2 results");
assertEquals(3, byProbe.get("test.probe2").size(),
"Probe 2 should have 3 results");

// Check per-probe success rates
Map<String, Double> successRates = report.successRateByProbe();
assertEquals(2, successRates.size());

// With secure mock, all should be blocked (0% success)
assertEquals(0.0, successRates.get("test.probe1"), 0.01);
assertEquals(0.0, successRates.get("test.probe2"), 0.01);

// Check probe summaries
Map<String, ScanReport.ProbeSummary> summaries = report.probeSummaries();
assertEquals(2, summaries.get("test.probe1").totalPrompts());
assertEquals(3, summaries.get("test.probe2").totalPrompts());
}

@Test
@DisplayName("Mean success rate by probe is computed correctly")
void testMeanSuccessRateByProbe() {
Probe probe1 = new AbstractProbe(
"test.probe1",
"Test Probe 1",
"Test probe",
AttackCategory.JAILBREAK,
List.of("Prompt 1", "Prompt 2"),
List.of("test"),
3
) {};

ProbeRegistry testRegistry = new ProbeRegistry();
testRegistry.register(probe1);

TiberiusScanner testScanner = TiberiusScanner.builder()
.withProbeRegistry(testRegistry)
.withGenerator(MockGenerator.secure())
.withDetector(CompositeDetector.defaultComposite())
.withConcurrency(1)
.build();
testScanner.setProbePatterns(new String[]{"test.*"});

ScanReport report = testScanner.scan();

// With secure mock, mean should be 0
assertEquals(0.0, report.meanSuccessRateByProbe(), 0.01);
}
}
Loading