feat: backport steer notes for analysis-only mode
Cherry-pick of upstream Shannon PR #329. Adds per-mode output format builders in queue-schemas.ts so the notes field description steers LLM output toward defensive context when exploit is disabled. Updates agent-execution to pass the exploit flag through to getOutputFormat. Co-Authored-By: Paperclip <noreply@paperclip.ing>
This commit is contained in:
@@ -17,15 +17,27 @@ import type { AgentName } from '../types/agents.js';
|
|||||||
|
|
||||||
// === Common Fields ===
|
// === Common Fields ===
|
||||||
|
|
||||||
const baseVulnerability = z.object({
|
const ANALYSIS_NOTES_DESCRIPTION =
|
||||||
ID: z.string(),
|
'Plain context for defenders (caveats, scope, what is at risk). Not attack steps.';
|
||||||
vulnerability_type: z.string(),
|
|
||||||
externally_exploitable: z.boolean(),
|
|
||||||
confidence: z.string(),
|
|
||||||
notes: z.string().optional(),
|
|
||||||
});
|
|
||||||
|
|
||||||
// === Per-Vuln-Type Schemas ===
|
function notesField(exploit: boolean) {
|
||||||
|
const f = z.string().optional();
|
||||||
|
return exploit ? f : f.describe(ANALYSIS_NOTES_DESCRIPTION);
|
||||||
|
}
|
||||||
|
|
||||||
|
function makeBase(exploit: boolean) {
|
||||||
|
return z.object({
|
||||||
|
ID: z.string(),
|
||||||
|
vulnerability_type: z.string(),
|
||||||
|
externally_exploitable: z.boolean(),
|
||||||
|
confidence: z.string(),
|
||||||
|
notes: notesField(exploit),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// === Per-Vuln-Type Schemas (used for type inference; notes description is mode-agnostic for types) ===
|
||||||
|
|
||||||
|
const baseVulnerability = makeBase(true);
|
||||||
|
|
||||||
const InjectionVulnerability = baseVulnerability.extend({
|
const InjectionVulnerability = baseVulnerability.extend({
|
||||||
source: z.string().optional(),
|
source: z.string().optional(),
|
||||||
@@ -87,14 +99,6 @@ export type AuthFinding = z.infer<typeof AuthVulnerability>;
|
|||||||
export type SsrfFinding = z.infer<typeof SsrfVulnerability>;
|
export type SsrfFinding = z.infer<typeof SsrfVulnerability>;
|
||||||
export type AuthzFinding = z.infer<typeof AuthzVulnerability>;
|
export type AuthzFinding = z.infer<typeof AuthzVulnerability>;
|
||||||
|
|
||||||
// === Queue Wrapper Schemas ===
|
|
||||||
|
|
||||||
const InjectionQueueSchema = z.object({ vulnerabilities: z.array(InjectionVulnerability) });
|
|
||||||
const XssQueueSchema = z.object({ vulnerabilities: z.array(XssVulnerability) });
|
|
||||||
const AuthQueueSchema = z.object({ vulnerabilities: z.array(AuthVulnerability) });
|
|
||||||
const SsrfQueueSchema = z.object({ vulnerabilities: z.array(SsrfVulnerability) });
|
|
||||||
const AuthzQueueSchema = z.object({ vulnerabilities: z.array(AuthzVulnerability) });
|
|
||||||
|
|
||||||
// === Convert to JSON Schema for SDK ===
|
// === Convert to JSON Schema for SDK ===
|
||||||
|
|
||||||
// NOTE: The SDK's AJV validator expects draft-07. Zod defaults to draft-2020-12 which
|
// NOTE: The SDK's AJV validator expects draft-07. Zod defaults to draft-2020-12 which
|
||||||
@@ -103,15 +107,65 @@ function toOutputFormat(zodSchema: z.ZodType): JsonSchemaOutputFormat {
|
|||||||
return { type: 'json_schema', schema: z.toJSONSchema(zodSchema, { target: 'draft-07' }) as Record<string, unknown> };
|
return { type: 'json_schema', schema: z.toJSONSchema(zodSchema, { target: 'draft-07' }) as Record<string, unknown> };
|
||||||
}
|
}
|
||||||
|
|
||||||
// === Lookup Maps ===
|
// === Per-Mode Output Format Builders ===
|
||||||
|
// Two maps cached at module load; the only per-mode difference is the
|
||||||
|
// description on the `notes` field, which steers the LLM's writing.
|
||||||
|
|
||||||
const VULN_AGENT_OUTPUT_FORMAT: Partial<Record<AgentName, JsonSchemaOutputFormat>> = {
|
function buildOutputFormats(exploit: boolean): Partial<Record<AgentName, JsonSchemaOutputFormat>> {
|
||||||
'injection-vuln': toOutputFormat(InjectionQueueSchema),
|
const base = makeBase(exploit);
|
||||||
'xss-vuln': toOutputFormat(XssQueueSchema),
|
return {
|
||||||
'auth-vuln': toOutputFormat(AuthQueueSchema),
|
'injection-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({
|
||||||
'ssrf-vuln': toOutputFormat(SsrfQueueSchema),
|
source: z.string().optional(),
|
||||||
'authz-vuln': toOutputFormat(AuthzQueueSchema),
|
combined_sources: z.string().optional(),
|
||||||
};
|
path: z.string().optional(),
|
||||||
|
sink_call: z.string().optional(),
|
||||||
|
slot_type: z.string().optional(),
|
||||||
|
sanitization_observed: z.string().optional(),
|
||||||
|
concat_occurrences: z.string().optional(),
|
||||||
|
verdict: z.string().optional(),
|
||||||
|
mismatch_reason: z.string().optional(),
|
||||||
|
witness_payload: z.string().optional(),
|
||||||
|
})) })),
|
||||||
|
'xss-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({
|
||||||
|
source: z.string().optional(),
|
||||||
|
source_detail: z.string().optional(),
|
||||||
|
path: z.string().optional(),
|
||||||
|
sink_function: z.string().optional(),
|
||||||
|
render_context: z.string().optional(),
|
||||||
|
encoding_observed: z.string().optional(),
|
||||||
|
verdict: z.string().optional(),
|
||||||
|
mismatch_reason: z.string().optional(),
|
||||||
|
witness_payload: z.string().optional(),
|
||||||
|
})) })),
|
||||||
|
'auth-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({
|
||||||
|
source_endpoint: z.string().optional(),
|
||||||
|
vulnerable_code_location: z.string().optional(),
|
||||||
|
missing_defense: z.string().optional(),
|
||||||
|
exploitation_hypothesis: z.string().optional(),
|
||||||
|
suggested_exploit_technique: z.string().optional(),
|
||||||
|
})) })),
|
||||||
|
'ssrf-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({
|
||||||
|
source_endpoint: z.string().optional(),
|
||||||
|
vulnerable_parameter: z.string().optional(),
|
||||||
|
vulnerable_code_location: z.string().optional(),
|
||||||
|
missing_defense: z.string().optional(),
|
||||||
|
exploitation_hypothesis: z.string().optional(),
|
||||||
|
suggested_exploit_technique: z.string().optional(),
|
||||||
|
})) })),
|
||||||
|
'authz-vuln': toOutputFormat(z.object({ vulnerabilities: z.array(base.extend({
|
||||||
|
endpoint: z.string().optional(),
|
||||||
|
vulnerable_code_location: z.string().optional(),
|
||||||
|
role_context: z.string().optional(),
|
||||||
|
guard_evidence: z.string().optional(),
|
||||||
|
side_effect: z.string().optional(),
|
||||||
|
reason: z.string().optional(),
|
||||||
|
minimal_witness: z.string().optional(),
|
||||||
|
})) })),
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
const OUTPUT_FORMATS_EXPLOIT = buildOutputFormats(true);
|
||||||
|
const OUTPUT_FORMATS_ANALYSIS = buildOutputFormats(false);
|
||||||
|
|
||||||
const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
|
const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
|
||||||
'injection-vuln': 'injection_exploitation_queue.json',
|
'injection-vuln': 'injection_exploitation_queue.json',
|
||||||
@@ -122,8 +176,8 @@ const VULN_AGENT_QUEUE_FILENAMES: Partial<Record<AgentName, string>> = {
|
|||||||
};
|
};
|
||||||
|
|
||||||
/** Returns the structured output format for a vuln agent, or undefined for non-vuln agents. */
|
/** Returns the structured output format for a vuln agent, or undefined for non-vuln agents. */
|
||||||
export function getOutputFormat(agentName: AgentName): JsonSchemaOutputFormat | undefined {
|
export function getOutputFormat(agentName: AgentName, exploit = true): JsonSchemaOutputFormat | undefined {
|
||||||
return VULN_AGENT_OUTPUT_FORMAT[agentName];
|
return (exploit ? OUTPUT_FORMATS_EXPLOIT : OUTPUT_FORMATS_ANALYSIS)[agentName];
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Returns the queue filename for a vuln agent, or undefined for non-vuln agents. */
|
/** Returns the queue filename for a vuln agent, or undefined for non-vuln agents. */
|
||||||
|
|||||||
@@ -161,7 +161,7 @@ export class AgentExecutionService {
|
|||||||
await auditSession.startAgent(agentName, prompt, attemptNumber);
|
await auditSession.startAgent(agentName, prompt, attemptNumber);
|
||||||
|
|
||||||
// 5. Execute agent
|
// 5. Execute agent
|
||||||
const outputFormat = getOutputFormat(agentName);
|
const outputFormat = getOutputFormat(agentName, distributedConfig?.exploit ?? true);
|
||||||
const result: ClaudePromptResult = await runClaudePrompt(
|
const result: ClaudePromptResult = await runClaudePrompt(
|
||||||
prompt,
|
prompt,
|
||||||
repoPath,
|
repoPath,
|
||||||
|
|||||||
Reference in New Issue
Block a user