diff --git a/src/main/java/io/tiberius/core/TiberiusScanner.java b/src/main/java/io/tiberius/core/TiberiusScanner.java index 7129cbf..a71f0bd 100644 --- a/src/main/java/io/tiberius/core/TiberiusScanner.java +++ b/src/main/java/io/tiberius/core/TiberiusScanner.java @@ -89,8 +89,12 @@ static TiberiusScanner create( public ScanReport scan() { final Instant startTime = Instant.now(); final List 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 results; if (concurrency > 1 && probesToRun.size() > 1) { @@ -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()); @@ -193,13 +209,16 @@ private List runSequential(final List probes) { final List 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 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; + } } } } @@ -211,8 +230,12 @@ private List runParallel(final List probes) { final List> futures = new ArrayList<>(); for (final Probe probe : probes) { - for (final Buff buff : buffsToApply) { - futures.add(CompletableFuture.supplyAsync(() -> runProbe(probe, buff), executor)); + final List prompts = probe.getPrompts(); + for (final String prompt : prompts) { + for (final Buff buff : buffsToApply) { + futures.add(CompletableFuture.supplyAsync( + () -> runProbe(probe, prompt, buff), executor)); + } } } diff --git a/src/main/java/io/tiberius/core/result/ScanReport.java b/src/main/java/io/tiberius/core/result/ScanReport.java index c3bf21e..9c88c93 100644 --- a/src/main/java/io/tiberius/core/result/ScanReport.java +++ b/src/main/java/io/tiberius/core/result/ScanReport.java @@ -105,6 +105,96 @@ public Map> 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> 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 successRateByProbe() { + return byProbe().entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + e -> { + List 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 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 probeSummaries() { + return byProbe().entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + e -> { + List 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. */ diff --git a/src/test/java/io/tiberius/core/TiberiusScannerTest.java b/src/test/java/io/tiberius/core/TiberiusScannerTest.java index 1344172..389c985 100644 --- a/src/test/java/io/tiberius/core/TiberiusScannerTest.java +++ b/src/test/java/io/tiberius/core/TiberiusScannerTest.java @@ -5,6 +5,8 @@ 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; @@ -12,6 +14,9 @@ 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.*; /** @@ -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> 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 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 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); + } }