use std::thread;
use crate::attacks::http::HttpClient;
use crate::klr::ast::{Attack, Expectation, Value};
use crate::klr::interpreter::{Interpreter, RunConfig};
use crate::results::AttackOutcome;
pub fn expand_all(attacks: &[Attack]) -> Vec<Attack> {
attacks.iter().flat_map(expand_attack).collect()
}
pub fn expand_attack(attack: &Attack) -> Vec<Attack> {
let mut base = attack.clone();
for c in &attack.checks {
base.expectations.extend(check_expectations(c));
}
base.checks.clear();
if base.mutations.is_empty() {
return vec![base];
}
let mutations = std::mem::take(&mut base.mutations);
let mut variants = Vec::new();
for m in &mutations {
for generator in &m.generators {
let values = crate::fuzz::generate(generator);
let multi = values.len() > 1;
for (i, val) in values.into_iter().enumerate() {
let mut v = base.clone();
set_send_field(&mut v.send, &m.field, &val);
let label = if multi {
format!("{generator}#{}", i + 1)
} else {
generator.clone()
};
v.name = format!("{} [{}={}]", base.name, m.field, label);
variants.push(v);
}
}
}
if variants.is_empty() {
vec![base]
} else {
variants
}
}
pub fn run_all<C: HttpClient + Sync>(
attacks: &[Attack],
client: &C,
config: &RunConfig,
workers: usize,
) -> Vec<AttackOutcome> {
let expanded = expand_all(attacks);
if workers <= 1 || expanded.len() <= 1 {
let interp = Interpreter::new(client, config.clone());
return interp.run(&expanded);
}
let worker_count = workers.min(expanded.len());
let chunks = split_indexed(&expanded, worker_count);
let mut indexed: Vec<(usize, AttackOutcome)> = thread::scope(|scope| {
let handles: Vec<_> = chunks
.into_iter()
.map(|chunk| {
scope.spawn(move || {
let interp = Interpreter::new(client, config.clone());
chunk
.into_iter()
.map(|(idx, attack)| {
(idx, interp.run(std::slice::from_ref(attack)).remove(0))
})
.collect::<Vec<_>>()
})
})
.collect();
handles
.into_iter()
.flat_map(|h| h.join().expect("worker thread panicked"))
.collect()
});
indexed.sort_by_key(|(i, _)| *i);
indexed.into_iter().map(|(_, o)| o).collect()
}
fn split_indexed(items: &[Attack], n: usize) -> Vec<Vec<(usize, &Attack)>> {
let n = n.max(1);
let mut chunks: Vec<Vec<(usize, &Attack)>> = vec![Vec::new(); n];
for (i, item) in items.iter().enumerate() {
chunks[i % n].push((i, item));
}
chunks
}
fn check_expectations(name: &str) -> Vec<Expectation> {
match name.to_ascii_lowercase().as_str() {
"authentication" | "auth" => vec![Expectation::Named {
name: "requires_auth".to_string(),
expected: true,
}],
"rate_limit" | "rate_limiting" | "ratelimit" => vec![Expectation::BlockedAfter(50)],
"injection" | "sql_injection" | "sqli" => vec![Expectation::Named {
name: "no_sql_error".to_string(),
expected: true,
}],
other => vec![Expectation::Named {
name: other.to_string(),
expected: true,
}],
}
}
fn set_send_field(send: &mut Vec<(String, Value)>, field: &str, value: &str) {
let v = Value::Str(value.to_string());
if let Some(entry) = send.iter_mut().find(|(k, _)| k == field) {
entry.1 = v;
} else {
send.push((field.to_string(), v));
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::klr::parser::parse;
#[test]
fn mutation_expands_into_variants() {
let src = r#"
attack duplicate_payment {
request POST "/payment"
send { amount = 100 }
mutate amount {
negative_numbers
huge_values
decimals
}
}
"#;
let program = parse(src).unwrap();
let expanded = expand_all(&program.attacks);
assert_eq!(expanded.len(), 6);
for v in &expanded {
let amount = v.send.iter().find(|(k, _)| k == "amount").unwrap();
assert!(matches!(amount.1, Value::Str(_)));
assert!(v.name.contains("amount="));
}
}
#[test]
fn check_becomes_expectation() {
let src = r#"
test login_security {
endpoint "/login"
check authentication
}
"#;
let program = parse(src).unwrap();
let expanded = expand_all(&program.attacks);
assert_eq!(expanded.len(), 1);
assert!(expanded[0].checks.is_empty());
assert_eq!(
expanded[0].expectations,
vec![Expectation::Named {
name: "requires_auth".to_string(),
expected: true
}]
);
}
#[test]
fn repeat_block_multiplies_request_count() {
let src = r#"
repeat 100 {
attack login {
target "/login"
}
}
"#;
let program = parse(src).unwrap();
assert_eq!(program.attacks.len(), 1);
assert_eq!(program.attacks[0].repeat, Some(100));
}
#[test]
fn unknown_check_expands_to_a_named_expectation() {
let src = r#"
test csrf_protection {
endpoint "/transfer"
check csrf
}
"#;
let program = parse(src).unwrap();
let expanded = expand_all(&program.attacks);
assert_eq!(
expanded[0].expectations,
vec![Expectation::Named {
name: "csrf".to_string(),
expected: true
}]
);
}
#[test]
fn split_covers_all_items_once() {
let attacks: Vec<Attack> = (0..10).map(|i| Attack::empty(format!("a{i}"), 1)).collect();
let chunks = split_indexed(&attacks, 3);
let mut seen: Vec<usize> = chunks
.iter()
.flat_map(|c| c.iter().map(|(i, _)| *i))
.collect();
seen.sort();
assert_eq!(seen, (0..10).collect::<Vec<_>>());
}
}