s-b-repo-rustsploit/build.rs

use std::collections::HashSet;
use std::env;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::Path;
use regex::Regex;

/// Build script that generates module dispatchers for exploits, scanners, and creds.
/// 
/// This script:
/// - Scans `src/modules/{category}/` directories recursively
/// - Finds all `.rs` files (excluding `mod.rs`) that export `pub async fn run(target: &str)`
/// - Generates dispatch functions that support both short names and full paths
/// - Creates deterministic, sorted output for better maintainability

fn main() {
    // Tell Cargo to rerun this build script if module directories change
    println!("cargo:rerun-if-changed=src/modules/exploits");
    println!("cargo:rerun-if-changed=src/modules/creds");
    println!("cargo:rerun-if-changed=src/modules/scanners");

    // Generate dispatchers for each module category
    let categories = vec![
        ("src/modules/exploits", "exploit_dispatch.rs", "crate::modules::exploits", "Exploit"),
        ("src/modules/creds", "creds_dispatch.rs", "crate::modules::creds", "Cred"),
        ("src/modules/scanners", "scanner_dispatch.rs", "crate::modules::scanners", "Scanner"),
    ];

    for (root, out_file, mod_prefix, category_name) in categories {
        if let Err(e) = generate_dispatch(root, out_file, mod_prefix, category_name) {
            eprintln!("❌ Error generating {} dispatcher: {}", category_name, e);
            std::process::exit(1);
        }
    }
}

/// Generates a dispatch function for a module category.
/// 
/// # Arguments
/// * `root` - Root directory to scan (e.g., "src/modules/exploits")
/// * `out_file` - Output filename (e.g., "exploit_dispatch.rs")
/// * `mod_prefix` - Module path prefix (e.g., "crate::modules::exploits")
/// * `category_name` - Category name for error messages (e.g., "Exploit")
fn generate_dispatch(
    root: &str,
    out_file: &str,
    mod_prefix: &str,
    category_name: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let out_dir = env::var("OUT_DIR")
        .map_err(|_| "OUT_DIR environment variable not set")?;
    let dest_path = Path::new(&out_dir).join(out_file);
    
    let root_path = Path::new(root);
    if !root_path.exists() {
        return Err(format!("Module directory '{}' does not exist", root).into());
    }

    // Collect all module mappings (using HashSet to avoid duplicates)
    let mut mappings = HashSet::new();
    visit_dirs(root_path, "".to_string(), &mut mappings)?;

    if mappings.is_empty() {
        eprintln!("⚠️  Warning: No modules found in {}", root);
    }

    // Sort mappings for deterministic output
    let mut sorted_mappings: Vec<_> = mappings.iter().collect();
    sorted_mappings.sort_by_key(|(key, _)| key);

    // Generate the dispatch function
    let mut file = File::create(&dest_path)
        .map_err(|e| format!("Failed to create {}: {}", dest_path.display(), e))?;

    writeln!(
        file,
        "// Auto-generated by build.rs - DO NOT EDIT MANUALLY\n"
    )?;

    writeln!(
        file,
        "/// Dispatches to the appropriate {} module based on module name.\n\
         /// Supports both short names (e.g., 'port_scanner') and full paths (e.g., 'scanners/port_scanner').",
        category_name.to_lowercase()
    )?;

    writeln!(
        file,
        "pub async fn dispatch(module_name: &str, target: &str) -> anyhow::Result<()> {{\n    match module_name {{"
    )?;

    // Generate match arms for each module (supporting both short and full names)
    for (key, mod_path) in &sorted_mappings {
        let short_key = key.rsplit('/').next().unwrap_or(key);
        let mod_code_path = mod_path.replace("/", "::");

        // Support both short name and full path
        if short_key == *key {
            // No subdirectory, only short name
            writeln!(
                file,
                r#"        "{k}" => {{ {p}::{m}::run(target).await? }},"#,
                k = key,
                m = mod_code_path,
                p = mod_prefix
            )?;
        } else {
            // Has subdirectory, support both short and full
            writeln!(
                file,
                r#"        "{short}" | "{full}" => {{ {p}::{m}::run(target).await? }},"#,
                short = short_key,
                full = key,
                m = mod_code_path,
                p = mod_prefix
            )?;
        }
    }

    writeln!(
        file,
        r#"        _ => anyhow::bail!("{} module '{{}}' not found.", module_name),"#,
        category_name
    )?;

    writeln!(file, "    }}\n    Ok(())\n}}")?;

    println!("✅ Generated {} with {} modules", out_file, sorted_mappings.len());
    Ok(())
}

/// Recursively visits directories to find all module files.
/// 
/// # Arguments
/// * `dir` - Directory to scan
/// * `prefix` - Current path prefix (e.g., "generic" or "camera/acti")
/// * `mappings` - Set to store (full_path, module_path) tuples
fn visit_dirs(
    dir: &Path,
    prefix: String,
    mappings: &mut HashSet<(String, String)>,
) -> Result<(), Box<dyn std::error::Error>> {
    // Compile regex once for better performance
    // Matches: pub async fn run(target: &str) or pub async fn run(_target: &str)
    let sig_re = Regex::new(r"pub\s+async\s+fn\s+run\s*\(\s*[^)]*:\s*&str\s*\)")
        .map_err(|e| format!("Failed to compile regex: {}", e))?;

    if !dir.is_dir() {
        return Ok(());
    }

    let mut entries: Vec<_> = fs::read_dir(dir)?
        .collect::<Result<Vec<_>, _>>()?;
    
    // Sort entries for deterministic processing
    entries.sort_by_key(|e| e.file_name());

    for entry in entries {
        let path = entry.path();
        let file_name = entry.file_name();

        if path.is_dir() {
            // Recursively visit subdirectories
            let sub_prefix = if prefix.is_empty() {
                file_name.to_string_lossy().to_string()
            } else {
                format!("{}/{}", prefix, file_name.to_string_lossy())
            };
            visit_dirs(&path, sub_prefix, mappings)?;
        } else if path.extension().map_or(false, |e| e == "rs") {
            // Process Rust files
            let file_stem = path.file_stem()
                .and_then(|s| s.to_str())
                .ok_or_else(|| format!("Invalid file name: {}", path.display()))?;

            // Skip mod.rs files
            if file_stem == "mod" {
                continue;
            }

            // Build module path
            let mod_path = if prefix.is_empty() {
                file_stem.to_string()
            } else {
                format!("{}/{}", prefix, file_stem)
            };

            // Full key includes the category prefix (will be added in generate_dispatch)
            let key = mod_path.clone();

            // Read and check for the run function signature
            let mut source = String::new();
            File::open(&path)?.read_to_string(&mut source)?;

            if sig_re.is_match(&source) {
                mappings.insert((key.clone(), mod_path.clone()));
                let display_path = if prefix.is_empty() {
                    file_stem.to_string()
                } else {
                    format!("{}/{}", prefix, file_stem)
                };
                println!("  ✅ Registered module: {}", display_path);
            } else {
                // Only warn in verbose mode to reduce noise
                if env::var("RUSTSPLOIT_VERBOSE_BUILD").is_ok() {
                    println!("  ⚠️  Skipping '{}': no matching 'pub async fn run(target: &str)'", path.display());
                }
            }
        }
    }

    Ok(())
}