backdisco 0.4.0

use std::collections::HashSet;

/// Position-aware hostname structure
/// Splits hostname by both '.' and '-' delimiters to preserve positional context
#[derive(Debug, Clone)]
pub struct HostnameStructure {
    pub segments: Vec<String>,  // All segments split by . and -
    pub base_domain: String,     // Final base domain (e.g., "testcorp.com")
    pub subdomain_segments: Vec<String>,  // Segments before base domain
}

/// Candidate generation context for structured generation
#[derive(Debug, Clone)]
pub struct GenerationContext {
    pub max_depth: usize,
    pub position_expansions: Vec<Vec<String>>,  // Alternatives for each position
    pub base_domain: String,
}

/// Common subdomain prefixes for brute forcing
pub const COMMON_SUBDOMAINS: &[&str] = &[
    "www", "api", "admin", "dev", "staging", "stage", "test", "qa", "uat",
    "prod", "production", "internal", "int", "ext", "external", "private",
    "public", "backend", "back", "origin", "origin-", "cdn", "static",
    "assets", "media", "images", "img", "files", "download", "upload",
    "app", "apps", "web", "mobile", "m", "portal", "dashboard", "panel",
    "console", "manage", "management", "auth", "login", "sso", "oauth",
    "api-v1", "api-v2", "api-v3", "v1", "v2", "v3", "graphql", "rest",
    "ws", "websocket", "socket", "realtime", "live", "stream",
    "mail", "email", "smtp", "imap", "pop", "mx",
    "db", "database", "mysql", "postgres", "redis", "data", "cache",
    "search", "elastic", "elasticsearch", "solr",
    "git", "gitlab", "github", "bitbucket", "svn", "repo", "repos",
    "ci", "cd", "jenkins", "travis", "build", "builds", "deploy",
    "logs", "log", "metrics", "monitor", "monitoring", "status", "health",
    "docs", "doc", "documentation", "help", "support", "kb",
    "blog", "news", "press", "marketing", "promo",
    "shop", "store", "ecommerce", "cart", "checkout", "pay", "payment",
    "billing", "invoice", "account", "accounts", "user", "users", "customer",
    "crm", "erp", "hr", "finance", "sales", "inventory",
    "vpn", "remote", "gateway", "proxy", "lb", "loadbalancer",
    "ns1", "ns2", "dns", "ntp", "time",
    "ftp", "sftp", "ssh", "rdp", "vnc",
    "demo", "sandbox", "lab", "labs", "preview", "beta", "alpha",
    "old", "new", "legacy", "archive", "backup", "bak", "dr",
    "us", "eu", "asia", "ap", "na", "emea", "us-east", "us-west", "eu-west",
    "data", "analytics", "bi", "report", "reports", "reporting",
    "api-internal", "api-external", "api-public", "api-private",
];

/// Extract the base domain from a hostname (e.g., "sub.example.com" -> "example.com")
/// Handles common TLDs including multi-part ones like .co.uk
pub fn extract_base_domain(hostname: &str) -> String {
    let parts: Vec<&str> = hostname.split('.').collect();
    
    if parts.len() <= 2 {
        return hostname.to_string();
    }
    
    // Check for common multi-part TLDs
    let multi_tlds = ["co.uk", "com.au", "co.nz", "co.jp", "com.br", "co.za", "org.uk"];
    
    let last_two = if parts.len() >= 2 {
        format!("{}.{}", parts[parts.len() - 2], parts[parts.len() - 1])
    } else {
        String::new()
    };
    
    if multi_tlds.contains(&last_two.as_str()) && parts.len() >= 3 {
        // Return last 3 parts (e.g., "example.co.uk")
        parts[parts.len() - 3..].join(".")
    } else {
        // Return last 2 parts (e.g., "example.com")
        parts[parts.len() - 2..].join(".")
    }
}

/// Extract subdomain words from a hostname
/// e.g., "blah.dev.api.test.com" -> ["blah", "dev", "api", "test"]
pub fn extract_subdomain_words(hostname: &str) -> Vec<String> {
    let base_domain = extract_base_domain(hostname);
    let prefix = if hostname.ends_with(&base_domain) && hostname.len() > base_domain.len() {
        &hostname[..hostname.len() - base_domain.len() - 1] // -1 for the dot
    } else {
        return Vec::new();
    };
    
    prefix
        .split('.')
        .filter(|s| !s.is_empty())
        .map(|s| s.to_lowercase())
        .collect()
}

/// Count the subdomain depth (number of labels before the base domain)
/// e.g., "blah.dev.api.test.com" -> 3 (blah, dev, api)
pub fn subdomain_depth(hostname: &str) -> usize {
    extract_subdomain_words(hostname).len()
}

/// Extract base domain from a wildcard SAN
/// e.g., "*.example.com" -> "example.com"
/// e.g., "*.cdn.example.com" -> "cdn.example.com"
pub fn wildcard_to_base(wildcard: &str) -> Option<String> {
    if wildcard.starts_with("*.") {
        Some(wildcard[2..].to_string())
    } else if wildcard.starts_with('*') {
        Some(wildcard[1..].trim_start_matches('.').to_string())
    } else {
        None
    }
}

/// Generate common subdomain variations from a wildcard
/// e.g., "*.example.com" -> ["www.example.com", "api.example.com", ...]
pub fn expand_wildcard(wildcard: &str, max_subdomains: usize) -> Vec<String> {
    let base = match wildcard_to_base(wildcard) {
        Some(b) => b,
        None => return Vec::new(),
    };
    
    let mut results = Vec::new();
    
    // Add the base domain itself
    results.push(base.clone());
    
    // Add common single-level subdomains
    for prefix in COMMON_SUBDOMAINS.iter() {
        results.push(format!("{}.{}", prefix, base));
    }
    
    // If max_subdomains > 1, add some two-level combinations
    if max_subdomains >= 2 {
        let depth_prefixes = ["dev", "staging", "prod", "internal", "api", "v1", "v2"];
        let second_level = ["api", "app", "web", "backend", "origin", "admin"];
        
        for d in depth_prefixes.iter() {
            for s in second_level.iter() {
                if d != s {
                    results.push(format!("{}.{}.{}", d, s, base));
                }
            }
        }
    }
    
    results
}

/// Generate brute force candidates by combining seed words with a base domain
pub fn generate_brute_candidates(
    seed_words: &[String],
    base_domain: &str,
    max_depth: usize,
) -> Vec<String> {
    let mut candidates = HashSet::new();
    
    // Single word subdomains
    for word in seed_words {
        candidates.insert(format!("{}.{}", word, base_domain));
    }
    
    // Two-word combinations if depth >= 2
    if max_depth >= 2 {
        for word1 in seed_words {
            for word2 in seed_words {
                if word1 != word2 {
                    candidates.insert(format!("{}.{}.{}", word1, word2, base_domain));
                }
            }
            // Also combine with common subdomains
            for common in COMMON_SUBDOMAINS.iter().take(20) { // Limit to avoid explosion
                candidates.insert(format!("{}.{}.{}", word1, common, base_domain));
                candidates.insert(format!("{}.{}.{}", common, word1, base_domain));
            }
        }
    }
    
    // Three-word combinations if depth >= 3
    if max_depth >= 3 {
        let key_prefixes = ["dev", "staging", "prod", "internal", "test"];
        for prefix in key_prefixes {
            for word in seed_words {
                for suffix in ["api", "app", "web", "backend"].iter() {
                    candidates.insert(format!("{}.{}.{}.{}", prefix, word, suffix, base_domain));
                }
            }
        }
    }
    
    let mut result: Vec<String> = candidates.into_iter().collect();
    result.sort();
    result
}

/// Parse hostname into position-aware structure
/// Splits by both '.' and '-' to preserve positional context
/// Example: "service-dev.corp.testcorp.com" -> segments: ["service", "dev", "corp", "testcorp", "com"], base: "testcorp.com"
pub fn parse_hostname_structure(hostname: &str) -> HostnameStructure {
    let base_domain = extract_base_domain(hostname);
    
    // Get the prefix part (everything before base domain)
    let prefix = if hostname.ends_with(&base_domain) && hostname.len() > base_domain.len() {
        &hostname[..hostname.len() - base_domain.len() - 1] // -1 for the dot
    } else {
        return HostnameStructure {
            segments: vec![base_domain.clone()],
            base_domain,
            subdomain_segments: Vec::new(),
        };
    };
    
    // Split prefix by both '.' and '-' to get all segments
    let mut subdomain_segments = Vec::new();
    for part in prefix.split('.') {
        // Further split each part by '-'
        for segment in part.split('-') {
            if !segment.is_empty() {
                subdomain_segments.push(segment.to_lowercase());
            }
        }
    }
    
    // Build full segments list (subdomain + base domain parts)
    let mut all_segments = subdomain_segments.clone();
    for part in base_domain.split('.') {
        all_segments.push(part.to_string());
    }
    
    HostnameStructure {
        segments: all_segments,
        base_domain,
        subdomain_segments,
    }
}

/// Generate cartesian product of multiple vectors
/// Returns all combinations of elements from each vector
fn cartesian_product(vectors: &[Vec<String>]) -> Vec<Vec<String>> {
    if vectors.is_empty() {
        return vec![Vec::new()];
    }
    
    let mut result = Vec::new();
    let first = &vectors[0];
    let rest = cartesian_product(&vectors[1..]);
    
    for item in first {
        for combo in &rest {
            let mut new_combo = vec![item.clone()];
            new_combo.extend_from_slice(combo);
            result.push(new_combo);
        }
    }
    
    result
}

/// Generate structured candidates using position-aware combinations
/// Generates candidates at all depth levels from max_depth down to 1
pub fn generate_structured_candidates(context: &GenerationContext) -> Vec<String> {
    let mut candidates = HashSet::new();
    
    // Generate candidates at each depth level (max_depth down to 1)
    for depth in (1..=context.max_depth.min(context.position_expansions.len())).rev() {
        if depth == 0 || depth > context.position_expansions.len() {
            continue;
        }
        
        // Get expansions for this depth
        let expansions = &context.position_expansions[..depth];
        
        // Generate all combinations at this depth
        let combinations = cartesian_product(expansions);
        
        for combo in combinations {
            // Join segments with dots to form subdomain
            let subdomain = combo.join(".");
            let candidate = format!("{}.{}", subdomain, context.base_domain);
            candidates.insert(candidate);
        }
    }
    
    let mut result: Vec<String> = candidates.into_iter().collect();
    result.sort();
    result
}

/// Merge all wordlist sources into a single seed word list
pub fn build_seed_wordlist(
    backend_words: &[String],
    additional_words: &[String],
) -> Vec<String> {
    let mut words: HashSet<String> = HashSet::new();
    
    // Add extracted backend words
    for word in backend_words {
        words.insert(word.to_lowercase());
    }
    
    // Add additional/LLM-expanded words
    for word in additional_words {
        words.insert(word.to_lowercase());
    }
    
    // Add common subdomains
    for word in COMMON_SUBDOMAINS.iter() {
        words.insert(word.to_string());
    }
    
    let mut result: Vec<String> = words.into_iter().collect();
    result.sort();
    result
}

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

    #[test]
    fn test_extract_base_domain() {
        assert_eq!(extract_base_domain("example.com"), "example.com");
        assert_eq!(extract_base_domain("sub.example.com"), "example.com");
        assert_eq!(extract_base_domain("a.b.c.example.com"), "example.com");
        assert_eq!(extract_base_domain("example.co.uk"), "example.co.uk");
        assert_eq!(extract_base_domain("sub.example.co.uk"), "example.co.uk");
    }

    #[test]
    fn test_extract_subdomain_words() {
        assert_eq!(
            extract_subdomain_words("blah.dev.api.test.com"),
            vec!["blah", "dev", "api"]
        );
        assert_eq!(
            extract_subdomain_words("www.example.com"),
            vec!["www"]
        );
        assert_eq!(
            extract_subdomain_words("example.com"),
            Vec::<String>::new()
        );
    }

    #[test]
    fn test_subdomain_depth() {
        assert_eq!(subdomain_depth("example.com"), 0);
        assert_eq!(subdomain_depth("www.example.com"), 1);
        assert_eq!(subdomain_depth("blah.dev.api.example.com"), 3);
    }

    #[test]
    fn test_wildcard_to_base() {
        assert_eq!(wildcard_to_base("*.example.com"), Some("example.com".to_string()));
        assert_eq!(wildcard_to_base("*.cdn.example.com"), Some("cdn.example.com".to_string()));
        assert_eq!(wildcard_to_base("example.com"), None);
    }

    #[test]
    fn test_expand_wildcard() {
        let expanded = expand_wildcard("*.example.com", 1);
        assert!(expanded.contains(&"example.com".to_string()));
        assert!(expanded.contains(&"www.example.com".to_string()));
        assert!(expanded.contains(&"api.example.com".to_string()));
    }

    #[test]
    fn test_parse_hostname_structure() {
        let structure = parse_hostname_structure("service-dev.corp.testcorp.com");
        assert_eq!(structure.subdomain_segments, vec!["service", "dev", "corp"]);
        assert_eq!(structure.base_domain, "testcorp.com");
        
        let structure2 = parse_hostname_structure("api.example.com");
        assert_eq!(structure2.subdomain_segments, vec!["api"]);
        assert_eq!(structure2.base_domain, "example.com");
        
        let structure3 = parse_hostname_structure("service-ol-dev-pbape.static-hosting-dev.backend-dev.example.cc");
        assert!(structure3.subdomain_segments.len() > 5);
        assert_eq!(structure3.base_domain, "example.cc");
    }

    #[test]
    fn test_cartesian_product() {
        let vectors = vec![
            vec!["a".to_string(), "b".to_string()],
            vec!["1".to_string(), "2".to_string()],
        ];
        let result = cartesian_product(&vectors);
        assert_eq!(result.len(), 4);
        assert!(result.contains(&vec!["a".to_string(), "1".to_string()]));
        assert!(result.contains(&vec!["a".to_string(), "2".to_string()]));
        assert!(result.contains(&vec!["b".to_string(), "1".to_string()]));
        assert!(result.contains(&vec!["b".to_string(), "2".to_string()]));
    }

    #[test]
    fn test_generate_structured_candidates() {
        let context = GenerationContext {
            max_depth: 2,
            position_expansions: vec![
                vec!["api".to_string(), "app".to_string()],
                vec!["dev".to_string(), "prod".to_string()],
            ],
            base_domain: "example.com".to_string(),
        };
        
        let candidates = generate_structured_candidates(&context);
        assert!(candidates.contains(&"api.dev.example.com".to_string()));
        assert!(candidates.contains(&"api.prod.example.com".to_string()));
        assert!(candidates.contains(&"app.dev.example.com".to_string()));
        assert!(candidates.contains(&"app.prod.example.com".to_string()));
        assert!(candidates.contains(&"api.example.com".to_string()));
        assert!(candidates.contains(&"app.example.com".to_string()));
    }
}