Update build.rs

build.rs

@@ -1,96 +1,213 @@
+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_dispatch(
-        "src/modules/exploits",
-        "exploit_dispatch.rs",
-        "crate::modules::exploits"
-    );
-    generate_dispatch(
-        "src/modules/creds",
-        "creds_dispatch.rs",
-        "crate::modules::creds"
-    );
-    generate_dispatch(
-        "src/modules/scanners",
-        "scanner_dispatch.rs",
-        "crate::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);
+        }
+    }
 }
 
-fn generate_dispatch(root: &str, out_file: &str, mod_prefix: &str) {
-    let out_dir = env::var("OUT_DIR").unwrap();
+/// 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 mut file = File::create(&dest_path).unwrap();
-
+    
     let root_path = Path::new(root);
-    let mut mappings = Vec::new();
-    visit_dirs(root_path, "".to_string(), &mut mappings).unwrap();
+    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 {{"
-    ).unwrap();
+    )?;
 
-    for (key, mod_path) in &mappings {
-        writeln!(
-            file,
-            r#"        "{k}" => {{ {p}::{m}::run(target).await? }},"#,
-            k = key,
-            m = mod_path.replace("/", "::"),
-            p = mod_prefix
-        ).unwrap();
+    // 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),"#
-    ).unwrap();
+        r#"        _ => anyhow::bail!("{} module '{{}}' not found.", module_name),"#,
+        category_name
+    )?;
 
-    writeln!(file, "    }}\n    Ok(())\n}}").unwrap();
+    writeln!(file, "    }}\n    Ok(())\n}}")?;
+
+    println!("✅ Generated {} with {} modules", out_file, sorted_mappings.len());
+    Ok(())
 }
 
-fn visit_dirs(dir: &Path, prefix: String, mappings: &mut Vec<(String, String)>) -> std::io::Result<()> {
-    let sig_re = Regex::new(r"pub\s+async\s+fn\s+run\s*\(\s*[_a-zA-Z]+\s*:\s*&str\s*\)").unwrap();
+/// 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() {
-        for entry in fs::read_dir(dir)? {
-            let entry = entry?;
-            let path = entry.path();
+    if !dir.is_dir() {
+        return Ok(());
+    }
 
-            if path.is_dir() {
-                let sub_prefix = format!("{}/{}", prefix, entry.file_name().to_string_lossy());
-                visit_dirs(&path, sub_prefix, mappings)?;
-            } else if path.extension().map_or(false, |e| e == "rs") {
-                let file_name = path.file_stem().unwrap().to_string_lossy().to_string();
-                if file_name == "mod" {
-                    continue;
-                }
+    let mut entries: Vec<_> = fs::read_dir(dir)?
+        .collect::<Result<Vec<_>, _>>()?;
+    
+    // Sort entries for deterministic processing
+    entries.sort_by_key(|e| e.file_name());
 
-                let mod_path = format!("{}/{}", prefix, file_name)
-                    .trim_start_matches('/')
-                    .to_string();
-                let key = mod_path.clone();
+    for entry in entries {
+        let path = entry.path();
+        let file_name = entry.file_name();
 
-                let mut source = String::new();
-                fs::File::open(&path)?.read_to_string(&mut source)?;
+        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()))?;
 
-                if sig_re.is_match(&source) {
-                    mappings.push((key.clone(), mod_path));
-                    println!("✅ Registered module: {}/{}", prefix, file_name);
+            // 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 {
-                    println!("⚠️  Skipping '{}': no matching 'pub async fn run(...)'", path.display());
+                    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(())
 }