diff --git a/build.rs b/build.rs index d8931cc..5a9871e 100644 --- a/build.rs +++ b/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> { + 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> { + // 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::, _>>()?; + + // 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(()) }