diff --git a/Cargo.lock b/Cargo.lock index ab67213..b798a40 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,7 +4,7 @@ version = 3 [[package]] name = "afl_runner" -version = "0.3.0" +version = "0.3.1" dependencies = [ "anyhow", "clap", diff --git a/Cargo.toml b/Cargo.toml index 364b3ec..4930aa3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ name = "afl_runner" authors = ["0x434b , + /// Input directory for AFL + pub input_dir: PathBuf, + /// Output directory for AFL + pub output_dir: PathBuf, + /// Miscellaneous AFL flags + pub misc_afl_flags: Vec, + /// Path to the target binary + pub target_binary: PathBuf, + /// Arguments for the target binary + pub target_args: Option, +} + +impl AflCmd { + /// Creates a new `AflCmd` instance + pub fn new(afl_binary: PathBuf, target_binary: PathBuf) -> Self { + Self { + afl_binary, + env: Vec::new(), + input_dir: PathBuf::new(), + output_dir: PathBuf::new(), + misc_afl_flags: Vec::new(), + target_binary, + target_args: None, + } + } + + /// Sets the environment variables for the AFL command + pub fn set_env(&mut self, env: Vec) { + self.env = env; + } + + /// Sets the input directory for AFL + pub fn set_input_dir(&mut self, input_dir: PathBuf) { + self.input_dir = input_dir; + } + + /// Sets the output directory for AFL + pub fn set_output_dir(&mut self, output_dir: PathBuf) { + self.output_dir = output_dir; + } + + /// Sets the miscellaneous AFL flags + pub fn set_misc_afl_flags(&mut self, misc_afl_flags: Vec) { + self.misc_afl_flags = misc_afl_flags; + } + + /// Sets the arguments for the target binary + pub fn set_target_args(&mut self, target_args: Option) { + self.target_args = target_args; + } + + /// Assembles the AFL command into a string + pub fn assemble(&self) -> String { + let mut cmd_parts = Vec::new(); + + cmd_parts.extend(self.env.iter().cloned()); + cmd_parts.push(self.afl_binary.display().to_string()); + cmd_parts.push(format!("-i {}", self.input_dir.display())); + cmd_parts.push(format!("-o {}", self.output_dir.display())); + cmd_parts.extend(self.misc_afl_flags.iter().cloned()); + cmd_parts.push(format!("-- {}", self.target_binary.display())); + + if let Some(target_args) = &self.target_args { + cmd_parts.push(target_args.clone()); + } + + cmd_parts.join(" ").trim().replace(" ", " ") + } +} + +/// Retrieves the amount of free memory in the system fn get_free_mem() -> u64 { let s = System::new_all(); s.free_memory() } +/// Applies a flag to a percentage of AFL configurations fn apply_flags(configs: &mut [AFLEnv], flag_accessor: F, percentage: f64, rng: &mut impl Rng) where F: Fn(&mut AFLEnv) -> &mut bool, @@ -29,76 +107,64 @@ where } } -fn apply_constrained_args(strings: &mut [String], args: &[(&str, f64)], rng: &mut impl Rng) { - let n = strings.len(); +/// Applies constrained arguments to a percentage of AFL commands +fn apply_constrained_args(cmds: &mut [AflCmd], args: &[(&str, f64)], rng: &mut impl Rng) { + let n = cmds.len(); for &(arg, percentage) in args { let count = (n as f64 * percentage) as usize; let mut available_indices: Vec = (0..n) - .filter(|i| !strings[*i].contains(arg.split_whitespace().next().unwrap())) + .filter(|i| !cmds[*i].misc_afl_flags.iter().any(|f| f.contains(arg))) .collect(); available_indices.shuffle(rng); for &index in available_indices.iter().take(count) { - strings[index].push_str(&format!(" {arg}")); + cmds[index].misc_afl_flags.push(arg.to_string()); } } } -fn apply_args(strings: &mut [String], arg: &str, percentage: f64, rng: &mut impl Rng) { - let count = (strings.len() as f64 * percentage) as usize; +/// Applies an argument to a percentage of AFL commands +fn apply_args(cmds: &mut [AflCmd], arg: &str, percentage: f64, rng: &mut impl Rng) { + let count = (cmds.len() as f64 * percentage) as usize; let mut indices = HashSet::new(); while indices.len() < count { - indices.insert(rng.gen_range(0..strings.len())); + indices.insert(rng.gen_range(0..cmds.len())); } for index in indices { - strings[index].push_str(&format!(" {arg}")); + cmds[index].misc_afl_flags.push(arg.to_string()); } } +/// Generates AFL commands based on the provided configuration pub struct AFLCmdGenerator { - /// Path to afl-fuzz - pub afl_binary: PathBuf, - /// Harness holding the binaries and arguments + /// The harness configuration pub harness: Harness, - /// Corpus directory + /// Input directory for AFL pub input_dir: PathBuf, - /// Output directory + /// Output directory for AFL pub output_dir: PathBuf, - /// Amount of runners + /// Number of AFL runners pub runners: u32, - /// Dictionary + /// Path to the dictionary file pub dictionary: Option, - /// Other raw AFL++ flags + /// Raw AFL flags pub raw_afl_flags: Option, -} - -impl Default for AFLCmdGenerator { - fn default() -> Self { - Self { - afl_binary: PathBuf::new(), - harness: Harness::new(PathBuf::new(), None, None, None, None), - input_dir: PathBuf::new(), - output_dir: PathBuf::new(), - runners: 1, - dictionary: None, - raw_afl_flags: None, - } - } + /// Path to the AFL binary + pub afl_binary: Option, } impl AFLCmdGenerator { + /// Creates a new `AFLCmdGenerator` instance pub fn new( harness: Harness, runners: u32, - afl_binary: Option, input_dir: PathBuf, output_dir: PathBuf, dictionary: Option, raw_afl_flags: Option, + afl_binary: Option, ) -> Self { - let afl_binary = Self::get_afl_fuzz(afl_binary).expect("Could not find afl-fuzz binary"); - let dict = dictionary.and_then(|d| { if d.exists() && d.is_file() { d.to_str().map(String::from) @@ -108,16 +174,17 @@ impl AFLCmdGenerator { }); Self { - afl_binary, harness, input_dir, output_dir, runners, dictionary: dict, raw_afl_flags, + afl_binary, } } + /// Retrieves AFL environment variables fn get_afl_env_vars() -> Vec { std::env::vars() .filter(|(k, _)| k.starts_with("AFL_")) @@ -125,11 +192,29 @@ impl AFLCmdGenerator { .collect::>() } - fn get_afl_fuzz + AsRef>(afl_fuzz: Option

) -> Option { - afl_fuzz - .map(Into::into) - .or_else(Self::get_afl_fuzz_from_path) - .or_else(Self::get_afl_fuzz_from_env) + /// Retrieves the path to the AFL binary + fn get_afl_fuzz(&self) -> Result { + self.afl_binary + .as_ref() + .map(PathBuf::from) + .or_else(|| std::env::var("AFL_PATH").map(PathBuf::from).ok()) + .or_else(|| { + let output = Command::new("which") + .arg("afl-fuzz") + .output() + .context("Failed to execute 'which'") + .ok()?; + if output.status.success() { + let afl_path = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if !afl_path.is_empty() { + Some(PathBuf::from(afl_path)) + } else { + None + } + } else { + None + } + }) .and_then(|path| { if path.exists() && path.is_file() && path.ends_with("afl-fuzz") { Some(path) @@ -137,68 +222,42 @@ impl AFLCmdGenerator { None } }) + .context("Could not find afl-fuzz binary") } - fn get_afl_fuzz_from_env() -> Option { - std::env::var("AFL_PATH").ok().and_then(|path| { - let afl_path = PathBuf::from(path); - if afl_path.exists() && afl_path.is_file() && afl_path.ends_with("afl-fuzz") { - Some(afl_path) - } else { - None - } - }) - } - - fn get_afl_fuzz_from_path() -> Option { - let output = Command::new("which") - .arg("afl-fuzz") - .output() - .expect("Failed to execute 'which'"); - if output.status.success() { - let afl_path = String::from_utf8_lossy(&output.stdout).trim().to_string(); - if afl_path.is_empty() { - None - } else { - Some(PathBuf::from(afl_path)) - } - } else { - None - } - } - - pub fn generate_afl_commands(&self) -> Vec { + /// Generates AFL commands based on the configuration + pub fn generate_afl_commands(&self) -> Result> { let mut rng = rand::thread_rng(); let configs = self.initialize_configs(&mut rng); - let mut strings = self.create_initial_strings(&configs); - - Self::apply_mutation_strategies(&mut strings, &mut rng); - Self::apply_queue_selection(&mut strings, &mut rng); - Self::apply_power_schedules(&mut strings); - self.apply_directory(&mut strings); - self.apply_fuzzer_roles(&mut strings); - self.apply_dictionary(&mut strings); - self.apply_sanitizer_or_target_binary(&mut strings); - self.apply_cmplog(&mut strings, &mut rng); - self.apply_target_args(&mut strings); - self.apply_cmpcov(&mut strings, &mut rng); + let mut cmds = self.create_initial_cmds(&configs)?; + + Self::apply_mutation_strategies(&mut cmds, &mut rng); + Self::apply_queue_selection(&mut cmds, &mut rng); + Self::apply_power_schedules(&mut cmds); + self.apply_directory(&mut cmds); + self.apply_fuzzer_roles(&mut cmds); + self.apply_dictionary(&mut cmds)?; + self.apply_sanitizer_or_target_binary(&mut cmds); + self.apply_cmplog(&mut cmds, &mut rng); + self.apply_target_args(&mut cmds); + self.apply_cmpcov(&mut cmds, &mut rng); // Inherit global AFL environment variables that are not already set let afl_env_vars: Vec = Self::get_afl_env_vars(); - let to_apply = afl_env_vars - .iter() - .filter(|env| !strings[0].contains(env.split('=').next().unwrap())) - .map(std::string::ToString::to_string) - .collect::>() - .join(" "); - - for s in &mut strings { - s.insert_str(0, &format!("{to_apply} ")); + for cmd in &mut cmds { + let to_apply = afl_env_vars + .iter() + .filter(|env| !cmd.env.iter().any(|e| e.starts_with(*env))) + .cloned() + .collect::>(); + cmd.set_env(to_apply); } - strings + let cmd_strings = cmds.into_iter().map(|cmd| cmd.assemble()).collect(); + Ok(cmd_strings) } + /// Initializes AFL configurations fn initialize_configs(&self, rng: &mut impl Rng) -> Vec { let mut configs = vec![AFLEnv::new(); self.runners as usize]; configs.last_mut().unwrap().final_sync = true; @@ -220,53 +279,64 @@ impl AFLCmdGenerator { configs } - fn create_initial_strings(&self, configs: &[AFLEnv]) -> Vec { - configs + /// Creates initial AFL commands + fn create_initial_cmds(&self, configs: &[AFLEnv]) -> Result> { + let afl_binary = self.get_afl_fuzz()?; + let target_binary = self.harness.target_binary.clone(); + + let cmds = configs .iter() .map(|config| { - format!( - "{} {} {}", - config.generate_afl_env_cmd(), - self.afl_binary.display(), - self.raw_afl_flags.clone().unwrap_or_default() - ) + let mut cmd = AflCmd::new(afl_binary.clone(), target_binary.clone()); + cmd.set_env(config.generate_afl_env_cmd()); + if let Some(raw_afl_flags) = &self.raw_afl_flags { + cmd.set_misc_afl_flags( + raw_afl_flags + .split_whitespace() + .map(str::to_string) + .collect(), + ); + } + cmd }) - .collect() + .collect(); + + Ok(cmds) } - fn apply_mutation_strategies(strings: &mut [String], rng: &mut impl Rng) { - // Apply different mutation strategies to a percentage of the configs + /// Applies mutation strategies to AFL commands + fn apply_mutation_strategies(cmds: &mut [AflCmd], rng: &mut impl Rng) { let mode_args = [("-P explore", 0.4), ("-P exploit", 0.2)]; - apply_constrained_args(strings, &mode_args, rng); + apply_constrained_args(cmds, &mode_args, rng); let format_args = [("-a binary", 0.3), ("-a text", 0.3)]; - apply_constrained_args(strings, &format_args, rng); - apply_args(strings, "-L 0", 0.1, rng); + apply_constrained_args(cmds, &format_args, rng); + apply_args(cmds, "-L 0", 0.1, rng); } - fn apply_queue_selection(strings: &mut [String], rng: &mut impl Rng) { - // Apply sequential queue selection to 20% of the configs - apply_args(strings, "-Z", 0.2, rng); + /// Applies queue selection to AFL commands + fn apply_queue_selection(cmds: &mut [AflCmd], rng: &mut impl Rng) { + apply_args(cmds, "-Z", 0.2, rng); } - fn apply_power_schedules(strings: &mut [String]) { - // Cycle through the different power schedules for the available runners + /// Applies power schedules to AFL commands + fn apply_power_schedules(cmds: &mut [AflCmd]) { let pscheds = ["fast", "explore", "coe", "lin", "quad", "exploit", "rare"]; - strings.iter_mut().enumerate().for_each(|(i, s)| { - s.push_str(&format!(" -p {}", pscheds[i % pscheds.len()])); + cmds.iter_mut().enumerate().for_each(|(i, cmd)| { + cmd.misc_afl_flags + .push(format!("-p {}", pscheds[i % pscheds.len()])); }); } - fn apply_directory(&self, strings: &mut Vec) { - for s in strings { - s.push_str(&format!( - " -i {} -o {}", - self.input_dir.display(), - self.output_dir.display() - )); + /// Applies input and output directories to AFL commands + fn apply_directory(&self, cmds: &mut [AflCmd]) { + for cmd in cmds { + cmd.set_input_dir(self.input_dir.clone()); + cmd.set_output_dir(self.output_dir.clone()); } } - fn apply_fuzzer_roles(&self, strings: &mut [String]) { + /// Applies fuzzer roles to AFL commands + fn apply_fuzzer_roles(&self, cmds: &mut [AflCmd]) { let target_fname = self .harness .target_binary @@ -275,159 +345,104 @@ impl AFLCmdGenerator { .to_str() .unwrap() .replace('.', "_"); - // Set one main fuzzer - strings[0].push_str(&format!(" -M main_{target_fname}")); - // Set the rest to be slaves - for (i, s) in strings[1..].iter_mut().enumerate() { - s.push_str(&format!(" -S secondary_{i}_{target_fname}")); + cmds[0] + .misc_afl_flags + .push(format!("-M main_{target_fname}")); + for (i, cmd) in cmds[1..].iter_mut().enumerate() { + cmd.misc_afl_flags + .push(format!("-S sub_{i}_{target_fname}")); } } - fn apply_dictionary(&self, strings: &mut Vec) { - // If a dictionary is provided, set it for all configs + /// Applies dictionary to AFL commands + fn apply_dictionary(&self, cmds: &mut [AflCmd]) -> Result<()> { if let Some(dict) = self.dictionary.as_ref() { - let dict_path = fs::canonicalize(dict).expect("Failed to resolve dictionary path"); - for s in strings { - s.push_str(&format!(" -x {}", dict_path.display())); + let dict_path = fs::canonicalize(dict).context("Failed to resolve dictionary path")?; + for cmd in cmds { + cmd.misc_afl_flags + .push(format!("-x {}", dict_path.display())); } } + Ok(()) } - fn apply_sanitizer_or_target_binary(&self, strings: &mut [String]) { - // Set the first one to be a sanitizer_binary if available, otherwise the target_binary + /// Applies sanitizer or target binary to AFL commands + fn apply_sanitizer_or_target_binary(&self, cmds: &mut [AflCmd]) { let binary = self .harness .sanitizer_binary .as_ref() .unwrap_or(&self.harness.target_binary); - strings[0].push_str(&format!(" -- {}", binary.display())); + cmds[0].target_binary.clone_from(binary); } - fn apply_cmplog(&self, strings: &mut [String], rng: &mut impl Rng) { + /// Applies CMPLOG instrumentation to AFL commands + fn apply_cmplog(&self, cmds: &mut [AflCmd], rng: &mut impl Rng) { if let Some(cmplog_binary) = self.harness.cmplog_binary.as_ref() { let num_cmplog_cfgs = (f64::from(self.runners) * 0.3) as usize; match num_cmplog_cfgs { - 0 => { - // We have a cmplog binary but not enough config slots to use it - } + 0 => {} 1 => { - // We have exactly one runner available for cmplog so we use `-l 2` - strings[1].push_str( - format!( - " -l 2 -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); + cmds[1] + .misc_afl_flags + .push(format!("-l 2 -c {}", cmplog_binary.display())); } 2 => { - // We have exactly two runners available for cmplog so we use `-l 2` and `-l 2AT` - strings[1].push_str( - format!( - " -l 2 -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); - strings[2].push_str( - format!( - " -l 2AT -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); + cmds[1] + .misc_afl_flags + .push(format!("-l 2 -c {}", cmplog_binary.display())); + cmds[2] + .misc_afl_flags + .push(format!("-l 2AT -c {}", cmplog_binary.display())); } 3 => { - // We can now use all three modes - strings[1].push_str( - format!( - " -l 2 -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); - strings[2].push_str( - format!( - " -l 2AT -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); - strings[3].push_str( - format!( - " -l 3 -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); - } - _ => { - // We have more than 3 runners available for cmplog so we use all three modes with - // the following distribution: - // - 70% for -l 2 - // - 10% for -l 3 - // - 20% for -l 2AT. - self.apply_cmplog_instrumentation(strings, num_cmplog_cfgs, cmplog_binary, rng); + cmds[1] + .misc_afl_flags + .push(format!("-l 2 -c {}", cmplog_binary.display())); + cmds[2] + .misc_afl_flags + .push(format!("-l 2AT -c {}", cmplog_binary.display())); + cmds[3] + .misc_afl_flags + .push(format!("-l 3 -c {}", cmplog_binary.display())); } + _ => Self::apply_cmplog_instrumentation_many( + cmds, + num_cmplog_cfgs, + cmplog_binary, + rng, + ), } - self.apply_normal_instrumentation(strings, num_cmplog_cfgs); - } else { - self.apply_normal_instrumentation(strings, 0); } } - fn apply_cmplog_instrumentation( - &self, - strings: &mut [String], + /// Applies CMPLOG instrumentation to >4 AFL commands + fn apply_cmplog_instrumentation_many( + cmds: &mut [AflCmd], num_cmplog_cfgs: usize, cmplog_binary: &Path, rng: &mut impl Rng, ) { - let cmplog_args = [("-l 2 ", 0.7), ("-l 3", 0.1), ("-l 2AT", 0.2)]; - apply_constrained_args(&mut strings[1..=num_cmplog_cfgs], &cmplog_args, rng); - for s in &mut strings[1..=num_cmplog_cfgs] { - s.push_str( - format!( - " -c {} -- {}", - cmplog_binary.display(), - self.harness.target_binary.display() - ) - .as_str(), - ); - } - } - - fn apply_normal_instrumentation(&self, strings: &mut [String], start_index: usize) { - for s in &mut strings[start_index + 1..] { - s.push_str(format!(" -- {}", self.harness.target_binary.display()).as_str()); + let cmplog_args = [("-l 2", 0.7), ("-l 3", 0.1), ("-l 2AT", 0.2)]; + apply_constrained_args(&mut cmds[1..=num_cmplog_cfgs], &cmplog_args, rng); + for cmd in &mut cmds[1..=num_cmplog_cfgs] { + cmd.misc_afl_flags + .push(format!("-c {}", cmplog_binary.display())); } } - fn apply_target_args(&self, strings: &mut [String]) { - // Appends target arguments to the command string + /// Applies target arguments to AFL commands + fn apply_target_args(&self, cmds: &mut [AflCmd]) { if let Some(target_args) = self.harness.target_args.as_ref() { - for s in strings { - s.push_str(format!(" {target_args}").as_str()); + for cmd in cmds { + cmd.set_target_args(Some(target_args.clone())); } } } - fn apply_cmpcov(&self, strings: &mut [String], rng: &mut impl Rng) { - // Use 1-3 CMPCOV instances if available. - // We want the following distribution: - // - 1 instance when >= 3 && < 8 runners - // - 2 instances when >= 8 && < 16 runners - // - 3 instances when >= 16 runners - // Unlike CMPLOG we need to replace the target binary with the cmpcov binary - // We never want to replace instace[0] as that one houses a SAN binary - // It is unclear if we want to pair CMPLOG with CMPCOV but let's attempt to avoid it - if self.harness.cmpcov_binary.as_ref().is_some() { + /// Applies CMPCOV instrumentation to AFL commands + fn apply_cmpcov(&self, cmds: &mut [AflCmd], rng: &mut impl Rng) { + if let Some(cmpcov_binary) = self.harness.cmpcov_binary.as_ref() { let max_cmpcov_instances = if self.runners >= 16 { 3 } else if self.runners >= 8 { @@ -437,27 +452,14 @@ impl AFLCmdGenerator { } else { 0 }; - // Find instances that don't have CMPLOG (-c) and replace the target binary with CMPLOG) - // Also skip instance[0] as that one houses a SAN binary - let mut cmpcov_indices = (1..strings.len()) - .filter(|i| !strings[*i].contains("-c")) + + let mut cmpcov_indices = (1..cmds.len()) + .filter(|i| !cmds[*i].misc_afl_flags.iter().any(|f| f.contains("-c"))) .collect::>(); - // Shuffle the indices so we don't always replace the same instances - // Stylepoints only cmpcov_indices.shuffle(rng); - // Find and replace the target binary string after the -- with the cmpcov binary - for i in cmpcov_indices.iter().take(max_cmpcov_instances) { - let target_binary = self.harness.target_binary.display().to_string(); - let cmpcov_binary = self - .harness - .cmpcov_binary - .as_ref() - .unwrap() - .display() - .to_string(); - if let Some(pos) = strings[*i].find(&target_binary) { - strings[*i].replace_range(pos..pos + target_binary.len(), &cmpcov_binary); - } + + for i in cmpcov_indices.into_iter().take(max_cmpcov_instances) { + cmds[i].target_binary.clone_from(cmpcov_binary); } } } diff --git a/src/afl_env.rs b/src/afl_env.rs index c35039c..f81e829 100644 --- a/src/afl_env.rs +++ b/src/afl_env.rs @@ -2,34 +2,30 @@ // AFLPlusPlus flags // Based on: https://aflplus.plus/docs/env_variables/ // ----------------------------------------- + +/// Struct representing the environment variables for `AFLPlusPlus` #[derive(Debug, Clone)] pub struct AFLEnv { /// `AFL_AUTORESUME` will resume a fuzz run (same as providing -i -) for an existing out folder, even if a different -i was provided. - /// Without this setting, afl-fuzz will refuse execution for a long-fuzzed out dir + /// Without this setting, afl-fuzz will refuse execution for a long-fuzzed out dir. pub autoresume: bool, /// `AFL_FINAL_SYNC` will cause the fuzzer to perform a final import of test cases when terminating. /// This is beneficial for -M main fuzzers to ensure it has all unique test cases and hence you only need to afl-cmin this single queue. pub final_sync: bool, /// Setting `AFL_DISABLE_TRIM` tells afl-fuzz not to trim test cases. pub disable_trim: bool, - + /// Setting `AFL_KEEP_TIMEOUTS` will keep longer running inputs if they reach new coverage. + pub keep_timeouts: bool, /// Setting `AFL_EXPAND_HAVOC_NOW` will start in the extended havoc mode that includes costly mutations. /// afl-fuzz automatically enables this mode when deemed useful otherwise. - pub keep_timeouts: bool, - - /// Setting `AFL_KEEP_TIMEOUTS` will keep longer running inputs if they reach new coverage pub expand_havoc_now: bool, - /// `AFL_IGNORE_SEED_PROBLEMS` will skip over crashes and timeouts in the seeds instead of exiting. pub ignore_seed_problems: bool, - - /// When setting `AFL_IMPORT_FIRST` tests cases from other fuzzers in the campaign are loaded first. - /// Note: This can slow down the start of the first fuzz - /// by quite a lot of you have many fuzzers and/or many seeds. + /// When setting `AFL_IMPORT_FIRST`, test cases from other fuzzers in the campaign are loaded first. + /// Note: This can slow down the start of the first fuzz by quite a lot if you have many fuzzers and/or many seeds. pub import_first: bool, - /// `AFL_TESTCACHE_SIZE` sets caching of test cases in MB (default: 50). - /// If enough RAM is available it is recommended to target values between 50-500MB. + /// If enough RAM is available, it is recommended to target values between 50-500MB. pub testcache_size: u32, } @@ -49,37 +45,32 @@ impl Default for AFLEnv { } impl AFLEnv { + /// Creates a new `AFLEnv` instance with default values pub fn new() -> Self { Self::default() } - // Generates a AFLPlusPlus environment variable string for the current settings - pub fn generate_afl_env_cmd(&self) -> String { - let mut command = String::new(); + /// Generates an `AFLPlusPlus` environment variable string for the current settings + pub fn generate_afl_env_cmd(&self) -> Vec { + let mut command = Vec::new(); - command.push_str(&format!("AFL_AUTORESUME={} ", u8::from(self.autoresume))); - command.push_str(&format!("AFL_FINAL_SYNC={} ", u8::from(self.final_sync))); - command.push_str(&format!( - "AFL_DISABLE_TRIM={} ", - u8::from(self.disable_trim) - )); - command.push_str(&format!( + command.push(format!("AFL_AUTORESUME={} ", u8::from(self.autoresume))); + command.push(format!("AFL_FINAL_SYNC={} ", u8::from(self.final_sync))); + command.push(format!("AFL_DISABLE_TRIM={} ", u8::from(self.disable_trim))); + command.push(format!( "AFL_KEEP_TIMEOUTS={} ", u8::from(self.keep_timeouts) )); - command.push_str(&format!( + command.push(format!( "AFL_EXPAND_HAVOC_NOW={} ", u8::from(self.expand_havoc_now) )); - command.push_str(&format!( + command.push(format!( "AFL_IGNORE_SEED_PROBLEMS={} ", u8::from(self.ignore_seed_problems) )); - command.push_str(&format!( - "AFL_IMPORT_FIRST={} ", - u8::from(self.import_first) - )); - command.push_str(&format!("AFL_TESTCACHE_SIZE={} ", self.testcache_size)); + command.push(format!("AFL_IMPORT_FIRST={} ", u8::from(self.import_first))); + command.push(format!("AFL_TESTCACHE_SIZE={} ", self.testcache_size)); command } diff --git a/src/cli.rs b/src/cli.rs index e00a250..3f31bd2 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -8,15 +8,18 @@ pub const AFL_CORPUS: &str = "/tmp/afl_input"; /// Default output directory const AFL_OUTPUT: &str = "/tmp/afl_output"; +/// Command-line interface for the Parallelized `AFLPlusPlus` Campaign Runner #[derive(Parser, Debug, Clone)] #[command(name = "Parallelized AFLPlusPlus Campaign Runner")] #[command(author = "C.K. ")] -#[command(version = "0.3.0")] +#[command(version = "0.3.1")] pub struct Cli { + /// Subcommand to execute #[command(subcommand)] pub cmd: Commands, } +/// Available subcommands #[derive(Subcommand, Clone, Debug)] pub enum Commands { /// Only generate the commands, don't run them @@ -27,6 +30,7 @@ pub enum Commands { Tui(TuiArgs), } +/// Arguments for the `tui` subcommand #[derive(Args, Clone, Debug, Default)] pub struct TuiArgs { /// Path to a `AFLPlusPlus` campaign directory, e.g. `afl_output` @@ -37,6 +41,7 @@ pub struct TuiArgs { pub afl_output: PathBuf, } +/// Arguments for the `gen` subcommand #[derive(Args, Clone, Debug)] pub struct GenArgs { /// Target binary to fuzz @@ -129,6 +134,7 @@ pub struct GenArgs { } impl GenArgs { + /// Merge the command-line arguments with the configuration pub fn merge(&self, config: &Config) -> Self { Self { target: self.target.clone().or_else(|| { @@ -215,8 +221,10 @@ impl GenArgs { } } +/// Arguments for the `run` subcommand #[derive(Args, Clone, Debug)] pub struct RunArgs { + /// Arguments for generating the commands #[command(flatten)] pub gen_args: GenArgs, /// Only show the generated commands, don't run them @@ -235,6 +243,7 @@ pub struct RunArgs { } impl RunArgs { + /// Merge the command-line arguments with the configuration pub fn merge(&self, config: &Config) -> Self { let gen_args = self.gen_args.merge(config); Self { @@ -253,40 +262,63 @@ impl RunArgs { } } +/// Configuration for the Parallelized `AFLPlusPlus` Campaign Runner #[derive(Deserialize, Default, Debug, Clone)] pub struct Config { + /// Target configuration pub target: TargetConfig, + /// AFL configuration pub afl_cfg: AflConfig, + /// Tmux configuration pub tmux: TmuxConfig, + /// Miscellaneous configuration pub misc: MiscConfig, } +/// Configuration for the target binary #[derive(Deserialize, Default, Debug, Clone)] pub struct TargetConfig { + /// Path to the target binary pub path: Option, + /// Path to the sanitizer binary pub san_path: Option, + /// Path to the CMPLOG binary pub cmpl_path: Option, + /// Path to the CMPCOV binary pub cmpc_path: Option, + /// Arguments for the target binary pub args: Option>, } +/// Configuration for AFL #[derive(Deserialize, Default, Debug, Clone)] pub struct AflConfig { + /// Number of AFL runners pub runners: Option, + /// Path to the AFL binary pub afl_binary: Option, + /// Path to the seed directory pub seed_dir: Option, + /// Path to the solution directory pub solution_dir: Option, + /// Path to the dictionary pub dictionary: Option, + /// Additional AFL flags pub afl_flags: Option, } +/// Configuration for tmux #[derive(Deserialize, Default, Debug, Clone)] pub struct TmuxConfig { + /// Dry run mode pub dry_run: Option, + /// Tmux session name pub session_name: Option, } +/// Miscellaneous configuration options #[derive(Deserialize, Default, Debug, Clone)] pub struct MiscConfig { + /// Enable TUI mode pub tui: Option, } diff --git a/src/data_collection.rs b/src/data_collection.rs index 447a842..d109f7a 100644 --- a/src/data_collection.rs +++ b/src/data_collection.rs @@ -1,322 +1,362 @@ -use crate::session::{CrashInfoDetails, SessionData}; +use crate::session::{CampaignData, CrashInfoDetails}; use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -pub fn collect_session_data(output_dir: &PathBuf) -> SessionData { - let mut session_data = SessionData::new(); - - let mut fuzzers_alive = 0; - - for entry in fs::read_dir(output_dir).unwrap() { - let entry = entry.unwrap(); - let path = entry.path(); - - if path.is_dir() { - let fuzzer_stats_path = path.join("fuzzer_stats"); - - if fuzzer_stats_path.exists() { - let content = fs::read_to_string(fuzzer_stats_path).unwrap(); - let lines: Vec<&str> = content.lines().collect(); - - for line in lines { - let parts: Vec<&str> = line.split(':').map(str::trim).collect(); - - if parts.len() == 2 { - let key = parts[0]; - let value = parts[1]; - - match key { - "start_time" => { - let start_time = UNIX_EPOCH - + Duration::from_secs(value.parse::().unwrap_or(0)); - let current_time = SystemTime::now(); - let duration = current_time.duration_since(start_time).unwrap(); - session_data.total_run_time = duration; - } - "execs_per_sec" => { - let exec_ps = value.parse::().unwrap_or(0.0); - if exec_ps > session_data.executions.ps_max { - session_data.executions.ps_max = exec_ps; - } else if exec_ps < session_data.executions.ps_min - || session_data.executions.ps_min == 0.0 - { - session_data.executions.ps_min = exec_ps; - } - session_data.executions.ps_cum += exec_ps; - } - "execs_done" => { - let execs_done = value.parse::().unwrap_or(0); - if execs_done > session_data.executions.max { - session_data.executions.max = execs_done; - } else if execs_done < session_data.executions.min - || session_data.executions.min == 0 - { - session_data.executions.min = execs_done; - } - session_data.executions.cum += execs_done; - } - "pending_favs" => { - let pending_favs = value.parse::().unwrap_or(0); - if pending_favs > session_data.pending.favorites_max { - session_data.pending.favorites_max = pending_favs; - } else if pending_favs < session_data.pending.favorites_min - || session_data.pending.favorites_min == 0 - { - session_data.pending.favorites_min = pending_favs; - } - session_data.pending.favorites_cum += pending_favs; - } - "pending_total" => { - let pending_total = value.parse::().unwrap_or(0); - if pending_total > session_data.pending.total_max { - session_data.pending.total_max = pending_total; - } else if pending_total < session_data.pending.total_min - || session_data.pending.total_min == 0 - { - session_data.pending.total_min = pending_total; - } - session_data.pending.total_cum += pending_total; - } - "stability" => { - let stability = - value.trim_end_matches('%').parse::().unwrap_or(0.0); - if stability > session_data.stability.max { - session_data.stability.max = stability; - } else if stability < session_data.stability.min - || session_data.stability.min == 0.0 - { - session_data.stability.min = stability; - } - } - "corpus_count" => { - let corpus_count = value.parse::().unwrap_or(0); - if corpus_count > session_data.corpus.count_max { - session_data.corpus.count_max = corpus_count; - } else if corpus_count < session_data.corpus.count_min - || session_data.corpus.count_min == 0 - { - session_data.corpus.count_min = corpus_count; - } - session_data.corpus.count_cum += corpus_count; - } - "bitmap_cvg" => { - let cvg = value.trim_end_matches('%').parse::().unwrap_or(0.0); - if cvg < session_data.coverage.bitmap_min - || session_data.coverage.bitmap_min == 0.0 - { - session_data.coverage.bitmap_min = cvg; - } else if cvg > session_data.coverage.bitmap_max { - session_data.coverage.bitmap_max = cvg; - } - } - "max_depth" => { - let levels = value.parse::().unwrap_or(0); - if levels > session_data.levels.max { - session_data.levels.max = levels; - } else if levels < session_data.levels.min - || session_data.levels.min == 0 - { - session_data.levels.min = levels; - } - } - "saved_crashes" => { - let saved_crashes = value.parse::().unwrap_or(0); - if saved_crashes > session_data.crashes.saved_max { - session_data.crashes.saved_max = saved_crashes; - } else if saved_crashes < session_data.crashes.saved_min - || session_data.crashes.saved_min == 0 - { - session_data.crashes.saved_min = saved_crashes; - } - session_data.crashes.saved_cum += saved_crashes; - } - "saved_hangs" => { - let saved_hangs = value.parse::().unwrap_or(0); - if saved_hangs > session_data.hangs.saved_max { - session_data.hangs.saved_max = saved_hangs; - } else if saved_hangs < session_data.hangs.saved_min - || session_data.hangs.saved_min == 0 - { - session_data.hangs.saved_min = saved_hangs; - } - session_data.hangs.saved_cum += saved_hangs; - } - "last_find" => { - let last_find = value.parse::().unwrap_or(0); - let last_find = UNIX_EPOCH + Duration::from_secs(last_find); - let current_time = SystemTime::now(); - let duration = current_time.duration_since(last_find).unwrap(); - if duration > session_data.time_without_finds { - session_data.time_without_finds = duration; - } - } - "afl_banner" => { - session_data.misc.afl_banner = value.to_string(); - } - "afl_version" => { - session_data.misc.afl_version = value.to_string(); - } - "cycles_done" => { - let cycles_done = value.parse::().unwrap_or(0); - if cycles_done > session_data.cycles.done_max { - session_data.cycles.done_max = cycles_done; - } else if cycles_done < session_data.cycles.done_min - || session_data.cycles.done_min == 0 - { - session_data.cycles.done_min = cycles_done; - } - } - "cycles_wo_finds" => { - let cycles_wo_finds = value.parse::().unwrap_or(0); - if cycles_wo_finds > session_data.cycles.wo_finds_max { - session_data.cycles.wo_finds_max = cycles_wo_finds; - } else if cycles_wo_finds < session_data.cycles.wo_finds_min - || session_data.cycles.wo_finds_min == 0 - { - session_data.cycles.wo_finds_min = cycles_wo_finds; - } - } - _ => {} +/// Data fetcher that collects session data based on the `AFLPlusPlus` `fuzzer_stats` file +#[derive(Debug, Clone)] +pub struct DataFetcher { + /// Output directory of the `AFLPlusPlus` fuzzing run + pub output_dir: PathBuf, +} + +impl DataFetcher { + /// Create a new `DataFetcher` instance + pub fn new(output_dir: &Path) -> Self { + Self { + output_dir: output_dir.to_path_buf(), + } + } + + /// Collects session data from the specified output directory + pub fn collect_session_data(&self) -> CampaignData { + let mut session_data = CampaignData::new(); + let mut fuzzers_alive = 0; + + fs::read_dir(&self.output_dir) + .unwrap() + .flatten() + .for_each(|entry| { + let path = entry.path(); + if path.is_dir() { + let fuzzer_stats_path = path.join("fuzzer_stats"); + if fuzzer_stats_path.exists() { + if let Ok(content) = fs::read_to_string(fuzzer_stats_path) { + Self::process_fuzzer_stats(&content, &mut session_data); + fuzzers_alive += 1; } } } + }); + + session_data.fuzzers_alive = fuzzers_alive; + Self::calculate_averages(&mut session_data, fuzzers_alive); + + let (last_crashes, last_hangs) = self.collect_session_crashes_hangs(10); + session_data.last_crashes = last_crashes; + session_data.last_hangs = last_hangs; + + session_data + } + + fn process_fuzzer_stats(content: &str, session_data: &mut CampaignData) { + let lines: Vec<&str> = content.lines().collect(); - fuzzers_alive += 1; + for line in lines { + let parts: Vec<&str> = line.split(':').map(str::trim).collect(); + + if parts.len() == 2 { + let key = parts[0]; + let value = parts[1]; + + match key { + "start_time" => Self::process_start_time(value, session_data), + "execs_per_sec" => Self::process_execs_per_sec(value, session_data), + "execs_done" => Self::process_execs_done(value, session_data), + "pending_favs" => Self::process_pending_favs(value, session_data), + "pending_total" => Self::process_pending_total(value, session_data), + "stability" => Self::process_stability(value, session_data), + "corpus_count" => Self::process_corpus_count(value, session_data), + "bitmap_cvg" => Self::process_bitmap_cvg(value, session_data), + "max_depth" => Self::process_max_depth(value, session_data), + "saved_crashes" => Self::process_saved_crashes(value, session_data), + "saved_hangs" => Self::process_saved_hangs(value, session_data), + "last_find" => Self::process_last_find(value, session_data), + "afl_banner" => session_data.misc.afl_banner = value.to_string(), + "afl_version" => session_data.misc.afl_version = value.to_string(), + "cycles_done" => Self::process_cycles_done(value, session_data), + "cycles_wo_finds" => Self::process_cycles_wo_finds(value, session_data), + _ => {} + } } } } - session_data.fuzzers_alive = fuzzers_alive; - - session_data.executions.ps_avg = if fuzzers_alive > 0 { - session_data.executions.ps_cum / fuzzers_alive as f64 - } else { - 0.0 - }; - session_data.executions.avg = if fuzzers_alive > 0 { - session_data.executions.cum / fuzzers_alive - } else { - 0 - }; - session_data.pending.favorites_avg = if fuzzers_alive > 0 { - session_data.pending.favorites_cum / fuzzers_alive - } else { - 0 - }; - session_data.pending.total_avg = if fuzzers_alive > 0 { - session_data.pending.total_cum / fuzzers_alive - } else { - 0 - }; - session_data.corpus.count_avg = if fuzzers_alive > 0 { - session_data.corpus.count_cum / fuzzers_alive - } else { - 0 - }; - session_data.crashes.saved_avg = if fuzzers_alive > 0 { - session_data.crashes.saved_cum / fuzzers_alive - } else { - 0 - }; - session_data.hangs.saved_avg = if fuzzers_alive > 0 { - session_data.hangs.saved_cum / fuzzers_alive - } else { - 0 - }; - session_data.coverage.bitmap_avg = - (session_data.coverage.bitmap_min + session_data.coverage.bitmap_max) / 2.0; - session_data.stability.avg = (session_data.stability.min + session_data.stability.max) / 2.0; - session_data.cycles.done_avg = - (session_data.cycles.done_min + session_data.cycles.done_max) / 2; - session_data.cycles.wo_finds_avg = - (session_data.cycles.wo_finds_min + session_data.cycles.wo_finds_max) / 2; - session_data.levels.avg = (session_data.levels.min + session_data.levels.max) / 2; - - let output_dir = output_dir.clone().into_os_string().into_string().unwrap(); - let (last_crashes, last_hangs) = collect_session_crashes_hangs(&output_dir, 10); - session_data.last_crashes = last_crashes; - session_data.last_hangs = last_hangs; - - session_data -} -fn collect_session_crashes_hangs( - output_dir: &str, - num_latest: usize, -) -> (Vec, Vec) { - let mut crashes = Vec::new(); - let mut hangs = Vec::new(); + fn process_start_time(value: &str, session_data: &mut CampaignData) { + let start_time = UNIX_EPOCH + Duration::from_secs(value.parse::().unwrap_or(0)); + let current_time = SystemTime::now(); + let duration = current_time.duration_since(start_time).unwrap(); + session_data.total_run_time = duration; + } - for entry in fs::read_dir(output_dir).unwrap() { - let entry = entry.unwrap(); - let subdir = entry.path(); + fn process_execs_per_sec(value: &str, session_data: &mut CampaignData) { + let exec_ps = value.parse::().unwrap_or(0.0); + if exec_ps > session_data.executions.ps_max { + session_data.executions.ps_max = exec_ps; + } else if exec_ps < session_data.executions.ps_min || session_data.executions.ps_min == 0.0 + { + session_data.executions.ps_min = exec_ps; + } + session_data.executions.ps_cum += exec_ps; + } - if subdir.is_dir() { - let fuzzer_name = subdir.file_name().unwrap().to_str().unwrap().to_string(); + fn process_execs_done(value: &str, session_data: &mut CampaignData) { + let execs_done = value.parse::().unwrap_or(0); + if execs_done > session_data.executions.max { + session_data.executions.max = execs_done; + } else if execs_done < session_data.executions.min || session_data.executions.min == 0 { + session_data.executions.min = execs_done; + } + session_data.executions.cum += execs_done; + } - let crashes_dir = subdir.join("crashes"); - if crashes_dir.is_dir() { - process_files(&crashes_dir, &fuzzer_name, &mut crashes); - } + fn process_pending_favs(value: &str, session_data: &mut CampaignData) { + let pending_favs = value.parse::().unwrap_or(0); + if pending_favs > session_data.pending.favorites_max { + session_data.pending.favorites_max = pending_favs; + } else if pending_favs < session_data.pending.favorites_min + || session_data.pending.favorites_min == 0 + { + session_data.pending.favorites_min = pending_favs; + } + session_data.pending.favorites_cum += pending_favs; + } - let hangs_dir = subdir.join("hangs"); - if hangs_dir.is_dir() { - process_files(&hangs_dir, &fuzzer_name, &mut hangs); - } + fn process_pending_total(value: &str, session_data: &mut CampaignData) { + let pending_total = value.parse::().unwrap_or(0); + if pending_total > session_data.pending.total_max { + session_data.pending.total_max = pending_total; + } else if pending_total < session_data.pending.total_min + || session_data.pending.total_min == 0 + { + session_data.pending.total_min = pending_total; + } + session_data.pending.total_cum += pending_total; + } + + fn process_stability(value: &str, session_data: &mut CampaignData) { + let stability = value.trim_end_matches('%').parse::().unwrap_or(0.0); + if stability > session_data.stability.max { + session_data.stability.max = stability; + } else if stability < session_data.stability.min || session_data.stability.min == 0.0 { + session_data.stability.min = stability; } } - crashes.sort_by(|a, b| b.time.cmp(&a.time)); - hangs.sort_by(|a, b| b.time.cmp(&a.time)); + fn process_corpus_count(value: &str, session_data: &mut CampaignData) { + let corpus_count = value.parse::().unwrap_or(0); + if corpus_count > session_data.corpus.count_max { + session_data.corpus.count_max = corpus_count; + } else if corpus_count < session_data.corpus.count_min || session_data.corpus.count_min == 0 + { + session_data.corpus.count_min = corpus_count; + } + session_data.corpus.count_cum += corpus_count; + } - ( - crashes.into_iter().take(num_latest).collect(), - hangs.into_iter().take(num_latest).collect(), - ) -} + fn process_bitmap_cvg(value: &str, session_data: &mut CampaignData) { + let cvg = value.trim_end_matches('%').parse::().unwrap_or(0.0); + if cvg < session_data.coverage.bitmap_min || session_data.coverage.bitmap_min == 0.0 { + session_data.coverage.bitmap_min = cvg; + } else if cvg > session_data.coverage.bitmap_max { + session_data.coverage.bitmap_max = cvg; + } + } + + fn process_max_depth(value: &str, session_data: &mut CampaignData) { + let levels = value.parse::().unwrap_or(0); + if levels > session_data.levels.max { + session_data.levels.max = levels; + } else if levels < session_data.levels.min || session_data.levels.min == 0 { + session_data.levels.min = levels; + } + } + + fn process_saved_crashes(value: &str, session_data: &mut CampaignData) { + let saved_crashes = value.parse::().unwrap_or(0); + if saved_crashes > session_data.crashes.saved_max { + session_data.crashes.saved_max = saved_crashes; + } else if saved_crashes < session_data.crashes.saved_min + || session_data.crashes.saved_min == 0 + { + session_data.crashes.saved_min = saved_crashes; + } + session_data.crashes.saved_cum += saved_crashes; + } + + fn process_saved_hangs(value: &str, session_data: &mut CampaignData) { + let saved_hangs = value.parse::().unwrap_or(0); + if saved_hangs > session_data.hangs.saved_max { + session_data.hangs.saved_max = saved_hangs; + } else if saved_hangs < session_data.hangs.saved_min || session_data.hangs.saved_min == 0 { + session_data.hangs.saved_min = saved_hangs; + } + session_data.hangs.saved_cum += saved_hangs; + } -fn process_files(dir: &PathBuf, fuzzer_name: &str, file_infos: &mut Vec) { - fs::read_dir(dir).unwrap().flatten().for_each(|file_entry| { - let file = file_entry.path(); - if file.is_file() { - let filename = file.file_name().unwrap().to_str().unwrap(); - if let Some(file_info) = parse_filename(filename) { - let file_info = CrashInfoDetails { - fuzzer_name: fuzzer_name.to_string(), - file_path: file, - id: file_info.0, - sig: file_info.1, - src: file_info.2, - time: file_info.3, - execs: file_info.4, - op: file_info.5, - rep: file_info.6, - }; - file_infos.push(file_info); + fn process_last_find(value: &str, session_data: &mut CampaignData) { + let last_find = value.parse::().unwrap_or(0); + let last_find = UNIX_EPOCH + Duration::from_secs(last_find); + let current_time = SystemTime::now(); + let duration = current_time.duration_since(last_find).unwrap(); + if duration > session_data.time_without_finds { + session_data.time_without_finds = duration; + } + } + + fn process_cycles_done(value: &str, session_data: &mut CampaignData) { + let cycles_done = value.parse::().unwrap_or(0); + if cycles_done > session_data.cycles.done_max { + session_data.cycles.done_max = cycles_done; + } else if cycles_done < session_data.cycles.done_min || session_data.cycles.done_min == 0 { + session_data.cycles.done_min = cycles_done; + } + } + + fn process_cycles_wo_finds(value: &str, session_data: &mut CampaignData) { + let cycles_wo_finds = value.parse::().unwrap_or(0); + if cycles_wo_finds > session_data.cycles.wo_finds_max { + session_data.cycles.wo_finds_max = cycles_wo_finds; + } else if cycles_wo_finds < session_data.cycles.wo_finds_min + || session_data.cycles.wo_finds_min == 0 + { + session_data.cycles.wo_finds_min = cycles_wo_finds; + } + } + + fn calculate_averages(session_data: &mut CampaignData, fuzzers_alive: usize) { + session_data.executions.ps_avg = if fuzzers_alive > 0 { + session_data.executions.ps_cum / fuzzers_alive as f64 + } else { + 0.0 + }; + session_data.executions.avg = if fuzzers_alive > 0 { + session_data.executions.cum / fuzzers_alive + } else { + 0 + }; + session_data.pending.favorites_avg = if fuzzers_alive > 0 { + session_data.pending.favorites_cum / fuzzers_alive + } else { + 0 + }; + session_data.pending.total_avg = if fuzzers_alive > 0 { + session_data.pending.total_cum / fuzzers_alive + } else { + 0 + }; + session_data.corpus.count_avg = if fuzzers_alive > 0 { + session_data.corpus.count_cum / fuzzers_alive + } else { + 0 + }; + session_data.crashes.saved_avg = if fuzzers_alive > 0 { + session_data.crashes.saved_cum / fuzzers_alive + } else { + 0 + }; + session_data.hangs.saved_avg = if fuzzers_alive > 0 { + session_data.hangs.saved_cum / fuzzers_alive + } else { + 0 + }; + session_data.coverage.bitmap_avg = + (session_data.coverage.bitmap_min + session_data.coverage.bitmap_max) / 2.0; + session_data.stability.avg = + (session_data.stability.min + session_data.stability.max) / 2.0; + session_data.cycles.done_avg = + (session_data.cycles.done_min + session_data.cycles.done_max) / 2; + session_data.cycles.wo_finds_avg = + (session_data.cycles.wo_finds_min + session_data.cycles.wo_finds_max) / 2; + session_data.levels.avg = (session_data.levels.min + session_data.levels.max) / 2; + } + + /// Collects information about the latest crashes and hangs from the output directory + fn collect_session_crashes_hangs( + &self, + num_latest: usize, + ) -> (Vec, Vec) { + let out_dir_str = self + .output_dir + .clone() + .into_os_string() + .into_string() + .unwrap(); + let mut crashes = Vec::new(); + let mut hangs = Vec::new(); + + for entry in fs::read_dir(out_dir_str).unwrap() { + let entry = entry.unwrap(); + let subdir = entry.path(); + + if subdir.is_dir() { + let fuzzer_name = subdir.file_name().unwrap().to_str().unwrap().to_string(); + + let crashes_dir = subdir.join("crashes"); + if crashes_dir.is_dir() { + Self::process_files(&crashes_dir, &fuzzer_name, &mut crashes); + } + + let hangs_dir = subdir.join("hangs"); + if hangs_dir.is_dir() { + Self::process_files(&hangs_dir, &fuzzer_name, &mut hangs); + } } } - }); -} -fn parse_filename( - filename: &str, -) -> Option<(String, Option, String, u64, u64, String, u64)> { - let parts: Vec<&str> = filename.split(',').collect(); - if parts.len() == 6 || parts.len() == 7 { - let id = parts[0].split(':').nth(1)?.to_string(); - let sig = if parts.len() == 7 { - Some(parts[1].split(':').nth(1)?.to_string()) + crashes.sort_by(|a, b| b.time.cmp(&a.time)); + hangs.sort_by(|a, b| b.time.cmp(&a.time)); + + ( + crashes.into_iter().take(num_latest).collect(), + hangs.into_iter().take(num_latest).collect(), + ) + } + + /// Processes files in a directory and extracts crash/hang information + fn process_files(dir: &PathBuf, fuzzer_name: &str, file_infos: &mut Vec) { + fs::read_dir(dir).unwrap().flatten().for_each(|file_entry| { + let file = file_entry.path(); + if file.is_file() { + let filename = file.file_name().unwrap().to_str().unwrap(); + if let Some(mut file_info) = Self::parse_filename(filename) { + file_info.fuzzer_name = fuzzer_name.to_string(); + file_info.file_path = file; + file_infos.push(file_info); + } + } + }); + } + + /// Parses a filename and extracts crash/hang information + fn parse_filename(filename: &str) -> Option { + let parts: Vec<&str> = filename.split(',').collect(); + if parts.len() == 6 || parts.len() == 7 { + let id = parts[0].split(':').nth(1)?.to_string(); + let sig = if parts.len() == 7 { + Some(parts[1].split(':').nth(1)?.to_string()) + } else { + None + }; + let src_index = if sig.is_some() { 2 } else { 1 }; + let src = parts[src_index].split(':').nth(1)?.to_string(); + let time = parts[src_index + 1].split(':').nth(1)?.parse().ok()?; + let execs = parts[src_index + 2].split(':').nth(1)?.parse().ok()?; + let op = parts[src_index + 3].split(':').nth(1)?.to_string(); + let rep = parts[src_index + 4].split(':').nth(1)?.parse().ok()?; + Some(CrashInfoDetails { + fuzzer_name: String::new(), + file_path: PathBuf::new(), + id, + sig, + src, + time, + execs, + op, + rep, + }) } else { None - }; - let src_index = if sig.is_some() { 2 } else { 1 }; - let src = parts[src_index].split(':').nth(1)?.to_string(); - let time = parts[src_index + 1].split(':').nth(1)?.parse().ok()?; - let execs = parts[src_index + 2].split(':').nth(1)?.parse().ok()?; - let op = parts[src_index + 3].split(':').nth(1)?.to_string(); - let rep = parts[src_index + 4].split(':').nth(1)?.parse().ok()?; - Some((id, sig, src, time, execs, op, rep)) - } else { - None + } } } diff --git a/src/harness.rs b/src/harness.rs index 84f02d7..d31447d 100644 --- a/src/harness.rs +++ b/src/harness.rs @@ -1,21 +1,36 @@ use std::fs; use std::path::PathBuf; +/// Represents a harness configuration +#[derive(Debug)] pub struct Harness { - // Instrumented and maybe AFL_HARDEN=1 + /// Instrumented and maybe `AFL_HARDEN=1` pub target_binary: PathBuf, - // AFL_USE_*SAN=1 + /// `AFL_USE_*SAN=1` pub sanitizer_binary: Option, - // AFL_LLVM_CMPLOG=1 + /// `AFL_LLVM_CMPLOG=1` pub cmplog_binary: Option, - // AFL_LLVM_LAF_ALL=1 + /// `AFL_LLVM_LAF_ALL=1` pub cmpcov_binary: Option, - // Additional arguments for the harness - // If the harness reads from stdin, use @@ as placeholder + /// Additional arguments for the harness + /// If the harness reads from stdin, use @@ as placeholder pub target_args: Option, } impl Harness { + /// Creates a new `Harness` instance + /// + /// # Arguments + /// + /// * `target_binary` - Path to the target binary + /// * `sanitizer_binary` - Optional path to the sanitizer binary + /// * `cmplog_binary` - Optional path to the CMPLOG binary + /// * `cmpcov_binary` - Optional path to the CMPCOV binary + /// * `target_args` - Optional additional arguments for the harness + /// + /// # Panics + /// + /// Panics if the target binary is not found pub fn new( target_binary: PathBuf, sanitizer_binary: Option, @@ -23,12 +38,12 @@ impl Harness { cmpcov_binary: Option, target_args: Option, ) -> Self { - let target_binary = Self::_get_target_binary(target_binary); + let target_binary = Self::get_target_binary(target_binary); assert!(target_binary.is_some(), "Could not find target binary"); - let sanitizer_binary = Self::_get_binary(sanitizer_binary); - let cmplog_binary = Self::_get_binary(cmplog_binary); - let cmpcov_binary = Self::_get_binary(cmpcov_binary); + let sanitizer_binary = Self::get_binary(sanitizer_binary); + let cmplog_binary = Self::get_binary(cmplog_binary); + let cmpcov_binary = Self::get_binary(cmpcov_binary); Self { target_binary: target_binary.unwrap(), sanitizer_binary, @@ -38,27 +53,42 @@ impl Harness { } } - fn _is_path_binary + std::convert::AsRef>(path: &P) -> bool { + /// Checks if the given path is a binary file + /// + /// # Arguments + /// + /// * `path` - The path to check + fn is_path_binary + std::convert::AsRef>(path: &P) -> bool { let path: PathBuf = path.into(); path.exists() && path.is_file() } - fn _get_binary + std::convert::AsRef>( + /// Resolves the path to a binary file + /// + /// # Arguments + /// + /// * `binary` - Optional path to the binary file + fn get_binary + std::convert::AsRef>( binary: Option

, ) -> Option { let binary = binary.map_or_else(PathBuf::new, std::convert::Into::into); - if Self::_is_path_binary(&binary) { + if Self::is_path_binary(&binary) { let resolved_bin = fs::canonicalize(binary).expect("Failed to resolve path"); return Some(resolved_bin); } None } - fn _get_target_binary + std::convert::AsRef>( + /// Resolves the path to the target binary + /// + /// # Arguments + /// + /// * `target_binary` - Path to the target binary + fn get_target_binary + std::convert::AsRef>( target_binary: P, ) -> Option { let target_binary = target_binary.into(); - if Self::_is_path_binary(&target_binary) { + if Self::is_path_binary(&target_binary) { let resolved_tbin = fs::canonicalize(target_binary).expect("Failed to resolve path"); return Some(resolved_tbin); } diff --git a/src/main.rs b/src/main.rs index 7aab8ec..1a69990 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,7 @@ +use anyhow::{bail, Result}; use std::path::Path; +use tmux::Session; +use tui::Tui; use clap::Parser; use cli::{Cli, Commands}; @@ -15,38 +18,34 @@ mod utils; use utils::{create_afl_runner, create_harness, generate_tmux_name, load_config}; -fn main() { +fn main() -> Result<()> { let cli_args = Cli::parse(); match cli_args.cmd { - Commands::Gen(gen_args) => handle_gen_command(&gen_args), - Commands::Run(run_args) => handle_run_command(&run_args), - Commands::Tui(tui_args) => handle_tui_command(&tui_args), + Commands::Gen(gen_args) => handle_gen_command(&gen_args)?, + Commands::Run(run_args) => handle_run_command(&run_args)?, + Commands::Tui(tui_args) => handle_tui_command(&tui_args)?, } + Ok(()) } -fn handle_gen_command(gen_args: &cli::GenArgs) { - let config_args = load_config(gen_args.config.as_ref()); +fn handle_gen_command(gen_args: &cli::GenArgs) -> Result<()> { + let config_args = load_config(gen_args.config.as_ref())?; let raw_afl_flags = config_args.afl_cfg.afl_flags.clone(); let merged_args = gen_args.merge(&config_args); - let harness = create_harness(&merged_args).unwrap_or_else(|e| { - eprintln!("Error creating harness: {e}"); - std::process::exit(1); - }); + let harness = create_harness(&merged_args)?; let afl_runner = create_afl_runner(&merged_args, harness, raw_afl_flags); - let cmds = afl_runner.generate_afl_commands(); + let cmds = afl_runner.generate_afl_commands()?; utils::print_generated_commands(&cmds); + Ok(()) } -fn handle_run_command(run_args: &cli::RunArgs) { - let config_args = load_config(run_args.gen_args.config.as_ref()); +fn handle_run_command(run_args: &cli::RunArgs) -> Result<()> { + let config_args = load_config(run_args.gen_args.config.as_ref())?; let raw_afl_flags = config_args.afl_cfg.afl_flags.clone(); let merged_args = run_args.merge(&config_args); - let harness = create_harness(&merged_args.gen_args).unwrap_or_else(|e| { - eprintln!("Error creating harness: {e}"); - std::process::exit(1); - }); + let harness = create_harness(&merged_args.gen_args)?; let afl_runner = create_afl_runner(&merged_args.gen_args, harness, raw_afl_flags); - let cmds = afl_runner.generate_afl_commands(); + let cmds = afl_runner.generate_afl_commands()?; let target_args = merged_args .gen_args .target_args @@ -55,34 +54,34 @@ fn handle_run_command(run_args: &cli::RunArgs) { .join(" "); let tmux_name = generate_tmux_name(&merged_args, &target_args); if merged_args.tui { - tmux::run_tmux_session_with_tui(&tmux_name, &cmds, &merged_args.gen_args); + Session::new(&tmux_name, &cmds).run_tmux_session_with_tui(&merged_args.gen_args); } else { - tmux::run_tmux_session(&tmux_name, &cmds); + Session::new(&tmux_name, &cmds).run_tmux_session(); } + Ok(()) } -fn handle_tui_command(tui_args: &cli::TuiArgs) { +fn handle_tui_command(tui_args: &cli::TuiArgs) -> Result<()> { if !tui_args.afl_output.exists() { - eprintln!("Output directory is required for TUI mode"); - std::process::exit(1); + bail!("Output directory is required for TUI mode"); } - validate_tui_output_dir(&tui_args.afl_output); - tui::run_tui_standalone(&tui_args.afl_output); + validate_tui_output_dir(&tui_args.afl_output)?; + Tui::run(&tui_args.afl_output); + Ok(()) } -fn validate_tui_output_dir(output_dir: &Path) { - output_dir.read_dir().unwrap().for_each(|entry| { - let entry = entry.unwrap(); - let path = entry.path(); +fn validate_tui_output_dir(output_dir: &Path) -> Result<()> { + for entry in output_dir.read_dir()? { + let path = entry?.path(); if path.is_dir() { let fuzzer_stats = path.join("fuzzer_stats"); if !fuzzer_stats.exists() { - eprintln!( + bail!( "Invalid output directory: {} is missing 'fuzzer_stats' file", path.display() ); - std::process::exit(1); } } - }); + } + Ok(()) } diff --git a/src/session.rs b/src/session.rs index 6674f73..3fec531 100644 --- a/src/session.rs +++ b/src/session.rs @@ -1,6 +1,8 @@ use std::path::PathBuf; use std::time::Duration; +#[allow(dead_code)] +#[derive(Default, Debug)] pub struct CrashInfoDetails { pub fuzzer_name: String, pub file_path: PathBuf, @@ -13,7 +15,8 @@ pub struct CrashInfoDetails { pub rep: u64, } -pub struct SessionData { +#[derive(Debug)] +pub struct CampaignData { pub fuzzers_alive: usize, pub total_run_time: Duration, pub executions: ExecutionsInfo, @@ -31,7 +34,7 @@ pub struct SessionData { pub misc: Misc, } -impl Default for SessionData { +impl Default for CampaignData { fn default() -> Self { Self { fuzzers_alive: 0, @@ -53,20 +56,20 @@ impl Default for SessionData { } } -impl SessionData { +impl CampaignData { pub fn new() -> Self { Self::default() } } -#[derive(Default)] +#[derive(Default, Debug)] pub struct Levels { pub avg: usize, pub min: usize, pub max: usize, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct CrashInfo { pub saved_cum: usize, pub saved_avg: usize, @@ -74,14 +77,14 @@ pub struct CrashInfo { pub saved_max: usize, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct StabilityInfo { pub avg: f64, pub min: f64, pub max: f64, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct Cycles { pub done_avg: usize, pub done_min: usize, @@ -91,7 +94,7 @@ pub struct Cycles { pub wo_finds_max: usize, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct ExecutionsInfo { pub avg: usize, pub min: usize, @@ -103,14 +106,14 @@ pub struct ExecutionsInfo { pub ps_cum: f64, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct CoverageInfo { pub bitmap_avg: f64, pub bitmap_min: f64, pub bitmap_max: f64, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct PendingInfo { pub favorites_avg: usize, pub favorites_cum: usize, @@ -122,7 +125,7 @@ pub struct PendingInfo { pub total_max: usize, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct CorpusInfo { pub count_avg: usize, pub count_cum: usize, @@ -130,7 +133,7 @@ pub struct CorpusInfo { pub count_max: usize, } -#[derive(Default)] +#[derive(Default, Debug)] pub struct Misc { pub afl_version: String, pub afl_banner: String, diff --git a/src/tmux.rs b/src/tmux.rs index 8dc67e6..961fb4c 100644 --- a/src/tmux.rs +++ b/src/tmux.rs @@ -6,15 +6,26 @@ use std::process::{Command, Stdio}; use std::{env, thread, time::Duration}; use crate::cli::GenArgs; +use crate::tui::Tui; +/// Represents a tmux session #[derive(Debug, Clone)] pub struct Session { + /// The name of the tmux session pub name: String, + /// The commands to be executed in the tmux session pub commands: Vec, + /// The path to the log file for the tmux session pub log_file: PathBuf, } impl Session { + /// Creates a new `Session` instance + /// + /// # Arguments + /// + /// * `session_name` - The name of the tmux session + /// * `cmds` - The commands to be executed in the tmux session pub fn new(session_name: &str, cmds: &[String]) -> Self { let commands = cmds .iter() @@ -31,10 +42,12 @@ impl Session { } } + /// Checks if the current process is running inside a tmux session fn in_tmux() -> bool { env::var("TMUX").is_ok() } + /// Creates a bash script for running the tmux session fn create_bash_script(&self) -> Result { let mut engine = upon::Engine::new(); engine.add_template("afl_fuzz", include_str!("templates/tmux.txt"))?; @@ -49,6 +62,7 @@ impl Session { .map_err(|e| anyhow::anyhow!("Error creating bash script: {}", e)) } + /// Kills the tmux session pub fn kill_session(&self) -> Result<()> { let mut cmd = Command::new("tmux"); cmd.arg("kill-session") @@ -61,6 +75,7 @@ impl Session { Ok(()) } + /// Finds the ID of the first window in the tmux session fn find_first_window_id(&self) -> Result { let output = Command::new("tmux") .args(["list-windows", "-t", &self.name]) @@ -80,6 +95,7 @@ impl Session { } } + /// Attaches to the tmux session pub fn attach(&self) -> Result<()> { let get_first_window_id = self.find_first_window_id()?; let target = format!("{}:{}", &self.name, get_first_window_id); @@ -92,6 +108,7 @@ impl Session { Ok(()) } + /// Helper function for creating directories fn mkdir_helper(dir: &Path, check_empty: bool) -> Result<()> { if dir.is_file() { bail!("{} is a file", dir.display()); @@ -114,6 +131,7 @@ impl Session { Ok(()) } + /// Runs the tmux session pub fn run(&self) -> Result<()> { if Self::in_tmux() { bail!("Already in tmux session. Nested tmux sessions are not supported."); @@ -164,8 +182,42 @@ impl Session { } Ok(()) } + + /// Runs the tmux session and attaches to it + pub fn run_tmux_session(&self) { + if let Err(e) = self.run() { + let _ = self.kill_session(); + eprintln!("Error running tmux session: {e}"); + } else { + self.attach().unwrap(); + } + } + + /// Runs the tmux session with a TUI + pub fn run_tmux_session_with_tui(&self, args: &GenArgs) { + if let Err(e) = self.run_tmux_session_detached() { + eprintln!("Error running TUI: {e}"); + return; + } + let output_dir = args.output_dir.clone().unwrap(); + + thread::sleep(Duration::from_secs(1)); // Wait for tmux session to start + + Tui::run(&output_dir); + } + + /// Runs the tmux session in detached mode + pub fn run_tmux_session_detached(&self) -> Result<()> { + if let Err(e) = self.run() { + let _ = self.kill_session(); + return Err(e); + } + println!("Session {} started in detached mode", self.name); + Ok(()) + } } +/// Gets user input from stdin fn get_user_input() -> String { std::io::stdin() .bytes() @@ -182,35 +234,3 @@ fn get_user_input() -> String { } }) } - -pub fn run_tmux_session(tmux_name: &str, cmds: &[String]) { - let tmux = Session::new(tmux_name, cmds); - if let Err(e) = tmux.run() { - let _ = tmux.kill_session(); - eprintln!("Error running tmux session: {e}"); - } else { - tmux.attach().unwrap(); - } -} - -pub fn run_tmux_session_with_tui(tmux_name: &str, cmds: &[String], args: &GenArgs) { - if let Err(e) = run_tmux_session_detached(tmux_name, cmds) { - eprintln!("Error running TUI: {e}"); - return; - } - let output_dir = args.output_dir.clone().unwrap(); - - thread::sleep(Duration::from_secs(1)); // Wait for tmux session to start - - crate::tui::run_tui_standalone(&output_dir); -} - -pub fn run_tmux_session_detached(tmux_name: &str, cmds: &[String]) -> Result<()> { - let tmux = Session::new(tmux_name, cmds); - if let Err(e) = tmux.run() { - let _ = tmux.kill_session(); - return Err(e); - } - println!("Session {tmux_name} started in detached mode"); - Ok(()) -} diff --git a/src/tui.rs b/src/tui.rs index 54f8c59..1c16d5a 100644 --- a/src/tui.rs +++ b/src/tui.rs @@ -4,8 +4,8 @@ use std::sync::mpsc; use std::thread; use std::time::Duration; -use crate::data_collection; -use crate::session::{CrashInfoDetails, SessionData}; +use crate::data_collection::DataFetcher; +use crate::session::{CampaignData, CrashInfoDetails}; use crossterm::terminal::{ disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen, }; @@ -18,313 +18,351 @@ use ratatui::{ Terminal, }; -pub fn run_tui_standalone(output_dir: &Path) { - let output_dir = output_dir.to_path_buf(); - let (session_data_tx, session_data_rx) = mpsc::channel(); - thread::spawn(move || loop { - let session_data = data_collection::collect_session_data(&output_dir); - if let Err(e) = session_data_tx.send(session_data) { - eprintln!("Error sending session data: {e}"); - break; - } - thread::sleep(Duration::from_secs(1)); - }); - - if let Err(e) = run(&session_data_rx) { - eprintln!("Error running TUI: {e}"); - } +/// Represents the TUI (Text User Interface) +pub struct Tui { + /// The terminal instance + terminal: Terminal>, } -pub fn run(session_data_rx: &mpsc::Receiver) -> io::Result<()> { - let stdout = io::stdout(); - let backend = CrosstermBackend::new(stdout); - let mut terminal = Terminal::new(backend)?; - - terminal.clear()?; - enable_raw_mode()?; - crossterm::execute!(terminal.backend_mut(), EnterAlternateScreen)?; - - loop { - if let Ok(session_data) = session_data_rx.recv_timeout(Duration::from_millis(500)) { - terminal.draw(|f| { - let chunks = create_layout(f.size()); - render_title(f, &session_data, chunks[0]); - render_process_timings_and_overall_results(f, &session_data, chunks[1]); - render_stage_progress_and_nerd_stats(f, &session_data, chunks[2]); - render_crash_solutions(f, &session_data, chunks[3]); - render_hang_solutions(f, &session_data, chunks[4]); - })?; - } +impl Tui { + /// Creates a new `Tui` instance + pub fn new() -> io::Result { + let stdout = io::stdout(); + let backend = CrosstermBackend::new(stdout); + let terminal = Terminal::new(backend)?; + Ok(Self { terminal }) + } - if crossterm::event::poll(Duration::from_millis(200))? { - if let crossterm::event::Event::Key(_) = crossterm::event::read()? { + /// Runs the TUI standalone with the specified output directory + pub fn run(output_dir: &Path) { + let output_dir = output_dir.to_path_buf(); + let (session_data_tx, session_data_rx) = mpsc::channel(); + thread::spawn(move || loop { + let session_data = DataFetcher::new(&output_dir).collect_session_data(); + if let Err(e) = session_data_tx.send(session_data) { + eprintln!("Error sending session data: {e}"); break; } + thread::sleep(Duration::from_millis(500)); + }); + + if let Err(e) = Self::new().and_then(|mut tui| tui.run_internal(&session_data_rx)) { + eprintln!("Error running TUI: {e}"); } } - disable_raw_mode()?; - crossterm::execute!(terminal.backend_mut(), LeaveAlternateScreen)?; - terminal.clear()?; - terminal.show_cursor()?; + /// Runs the TUI with the specified session data receiver + fn run_internal(&mut self, session_data_rx: &mpsc::Receiver) -> io::Result<()> { + self.terminal.clear()?; + enable_raw_mode()?; + crossterm::execute!(self.terminal.backend_mut(), EnterAlternateScreen)?; - Ok(()) -} + loop { + if let Ok(session_data) = session_data_rx.recv_timeout(Duration::from_millis(500)) { + self.draw(&session_data)?; + } -fn create_layout(size: Rect) -> Vec { - let main_layout = Layout::default() - .direction(Direction::Vertical) - .margin(1) - .constraints([Constraint::Length(1), Constraint::Min(0)].as_ref()) - .split(size); - - let inner_layout = Layout::default() - .direction(Direction::Vertical) - .margin(1) - .constraints( - [ - Constraint::Percentage(15), - Constraint::Percentage(10), - Constraint::Percentage(30), - Constraint::Percentage(30), - ] - .as_ref(), - ) - .split(main_layout[1]); - - [ - main_layout[0], - inner_layout[0], - inner_layout[1], - inner_layout[2], - inner_layout[3], - ] - .to_vec() -} + if crossterm::event::poll(Duration::from_millis(200))? { + if let crossterm::event::Event::Key(_) = crossterm::event::read()? { + break; + } + } + } -fn render_title(f: &mut Frame, session_data: &SessionData, area: Rect) { - let title = Paragraph::new(format!( - "AFL {} - {} - Fuzzing campaign runner by @0xricksanchez", - session_data.misc.afl_version, session_data.misc.afl_banner - )) - .alignment(Alignment::Center) - .style( - Style::default() - .fg(Color::Yellow) - .add_modifier(Modifier::BOLD), - ); - - f.render_widget(title, area); -} + disable_raw_mode()?; + crossterm::execute!(self.terminal.backend_mut(), LeaveAlternateScreen)?; + self.terminal.clear()?; + self.terminal.show_cursor()?; -fn render_process_timings_and_overall_results( - f: &mut Frame, - session_data: &SessionData, - area: Rect, -) { - let hor_layout = Layout::default() - .direction(Direction::Horizontal) - .constraints([Constraint::Ratio(1, 2), Constraint::Ratio(1, 2)]) - .split(area); - - let p_proc_timings = create_process_timings_paragraph(session_data); - let p_overall_res = create_overall_results_paragraph(session_data); - - f.render_widget(p_proc_timings, hor_layout[0]); - f.render_widget(p_overall_res, hor_layout[1]); -} + Ok(()) + } -fn render_stage_progress_and_nerd_stats(f: &mut Frame, session_data: &SessionData, area: Rect) { - let hor_layout = Layout::default() - .direction(Direction::Horizontal) - .constraints([Constraint::Ratio(1, 2), Constraint::Ratio(1, 2)]) - .split(area); + /// Draws the TUI with the specified session data + fn draw(&mut self, session_data: &CampaignData) -> io::Result<()> { + self.terminal.draw(|f| { + let chunks = Self::create_layout(f.size()); + Self::render_title(f, session_data, chunks[0]); + Self::render_process_timings_and_overall_results(f, session_data, chunks[1]); + Self::render_stage_progress_and_nerd_stats(f, session_data, chunks[2]); + Self::render_crash_solutions(f, session_data, chunks[3]); + Self::render_hang_solutions(f, session_data, chunks[4]); + })?; + Ok(()) + } - let p_stage_prog = create_stage_progress_paragraph(session_data); - let p_nerd_stats = create_nerd_stats_paragraph(session_data); + /// Creates the layout for the TUI + fn create_layout(size: Rect) -> Vec { + let main_layout = Layout::default() + .direction(Direction::Vertical) + .margin(1) + .constraints([Constraint::Length(1), Constraint::Min(0)].as_ref()) + .split(size); + + let inner_layout = Layout::default() + .direction(Direction::Vertical) + .margin(1) + .constraints( + [ + Constraint::Percentage(15), + Constraint::Percentage(10), + Constraint::Percentage(30), + Constraint::Percentage(30), + ] + .as_ref(), + ) + .split(main_layout[1]); + + [ + main_layout[0], + inner_layout[0], + inner_layout[1], + inner_layout[2], + inner_layout[3], + ] + .to_vec() + } - f.render_widget(p_stage_prog, hor_layout[0]); - f.render_widget(p_nerd_stats, hor_layout[1]); -} + /// Renders the title section of the TUI + fn render_title(f: &mut Frame, session_data: &CampaignData, area: Rect) { + let title = Paragraph::new(format!( + "AFL {} - {} - Fuzzing campaign runner by @0xricksanchez", + session_data.misc.afl_version, session_data.misc.afl_banner + )) + .alignment(Alignment::Center) + .style( + Style::default() + .fg(Color::Yellow) + .add_modifier(Modifier::BOLD), + ); -fn render_crash_solutions(f: &mut Frame, session_data: &SessionData, area: Rect) { - let p_crash_solutions = Paragraph::new(format_solutions(&session_data.last_crashes)) - .block( - Block::default() - .title("10 Latest Crashes") - .borders(Borders::ALL) - .border_style(Style::default()), - ) - .style(Style::default()); - - f.render_widget(p_crash_solutions, area); -} + f.render_widget(title, area); + } -fn render_hang_solutions(f: &mut Frame, session_data: &SessionData, area: Rect) { - let p_hang_solutions = Paragraph::new(format_solutions(&session_data.last_hangs)) - .block( - Block::default() - .title("10 Latest Hangs") - .borders(Borders::ALL) - .border_style(Style::default()), - ) - .style(Style::default()); - - f.render_widget(p_hang_solutions, area); -} + /// Renders the process timings and overall results section of the TUI + fn render_process_timings_and_overall_results( + f: &mut Frame, + session_data: &CampaignData, + area: Rect, + ) { + let hor_layout = Layout::default() + .direction(Direction::Horizontal) + .constraints([Constraint::Ratio(1, 2), Constraint::Ratio(1, 2)]) + .split(area); + + let p_proc_timings = Self::create_process_timings_paragraph(session_data); + let p_overall_res = Self::create_overall_results_paragraph(session_data); + + f.render_widget(p_proc_timings, hor_layout[0]); + f.render_widget(p_overall_res, hor_layout[1]); + } + + /// Renders the stage progress and nerd stats section of the TUI + fn render_stage_progress_and_nerd_stats( + f: &mut Frame, + session_data: &CampaignData, + area: Rect, + ) { + let hor_layout = Layout::default() + .direction(Direction::Horizontal) + .constraints([Constraint::Ratio(1, 2), Constraint::Ratio(1, 2)]) + .split(area); + + let p_stage_prog = Self::create_stage_progress_paragraph(session_data); + let p_nerd_stats = Self::create_nerd_stats_paragraph(session_data); + + f.render_widget(p_stage_prog, hor_layout[0]); + f.render_widget(p_nerd_stats, hor_layout[1]); + } + + /// Renders the crash solutions section of the TUI + fn render_crash_solutions(f: &mut Frame, session_data: &CampaignData, area: Rect) { + let p_crash_solutions = Paragraph::new(Self::format_solutions(&session_data.last_crashes)) + .block( + Block::default() + .title("10 Latest Crashes") + .borders(Borders::ALL) + .border_style(Style::default()), + ) + .style(Style::default()); + + f.render_widget(p_crash_solutions, area); + } -fn create_process_timings_paragraph(session_data: &SessionData) -> Paragraph { - let last_seen_crash = - format_last_event(&session_data.last_crashes, &session_data.total_run_time); - let last_seen_hang = format_last_event(&session_data.last_hangs, &session_data.total_run_time); + /// Renders the hang solutions section of the TUI + fn render_hang_solutions(f: &mut Frame, session_data: &CampaignData, area: Rect) { + let p_hang_solutions = Paragraph::new(Self::format_solutions(&session_data.last_hangs)) + .block( + Block::default() + .title("10 Latest Hangs") + .borders(Borders::ALL) + .border_style(Style::default()), + ) + .style(Style::default()); - // TODO: Check if fuzzers are actually alive and report back when >1 was lost (warning) - // TODO: If all fuzzers are dead display an error and stop `total_run_time` - let content = format!( - "Fuzzers alive: {} + f.render_widget(p_hang_solutions, area); + } + + /// Creates the process timings paragraph + fn create_process_timings_paragraph(session_data: &CampaignData) -> Paragraph { + let last_seen_crash = + Self::format_last_event(&session_data.last_crashes, &session_data.total_run_time); + let last_seen_hang = + Self::format_last_event(&session_data.last_hangs, &session_data.total_run_time); + + // TODO: Check if fuzzers are actually alive and report back when >1 was lost (warning) + // TODO: If all fuzzers are dead display an error and stop `total_run_time` + let content = format!( + "Fuzzers alive: {} Total run time: {} Time without finds: {} Last saved crash: {} Last saved hang: {}", - session_data.fuzzers_alive, - format_duration(session_data.total_run_time), - format_duration(session_data.time_without_finds), - last_seen_crash, - last_seen_hang - ); - - Paragraph::new(content) - .block( - Block::default() - .title("Process timing") - .borders(Borders::ALL) - .add_modifier(Modifier::BOLD), - ) - .style(Style::default()) -} + session_data.fuzzers_alive, + Self::format_duration(session_data.total_run_time), + Self::format_duration(session_data.time_without_finds), + last_seen_crash, + last_seen_hang + ); + + Paragraph::new(content) + .block( + Block::default() + .title("Process timing") + .borders(Borders::ALL) + .add_modifier(Modifier::BOLD), + ) + .style(Style::default()) + } -fn create_overall_results_paragraph(session_data: &SessionData) -> Paragraph { - let content = format!( - "Cycles done: {} ({}/{}) + /// Creates the overall results paragraph + fn create_overall_results_paragraph(session_data: &CampaignData) -> Paragraph { + let content = format!( + "Cycles done: {} ({}/{}) Crashes saved: {} ({}->{}<-{}) Hangs saved: {} ({}->{}<-{}) Corpus count: {} ({}->{}<-{})", - session_data.cycles.done_avg, - session_data.cycles.done_min, - session_data.cycles.done_max, - session_data.crashes.saved_cum, - session_data.crashes.saved_min, - session_data.crashes.saved_avg, - session_data.crashes.saved_max, - session_data.hangs.saved_cum, - session_data.hangs.saved_min, - session_data.hangs.saved_avg, - session_data.hangs.saved_max, - session_data.corpus.count_cum, - session_data.corpus.count_min, - session_data.corpus.count_avg, - session_data.corpus.count_max - ); - - Paragraph::new(content) - .block( - Block::default() - .title("Overall results") - .borders(Borders::ALL) - .add_modifier(Modifier::BOLD), - ) - .style(Style::default()) -} + session_data.cycles.done_avg, + session_data.cycles.done_min, + session_data.cycles.done_max, + session_data.crashes.saved_cum, + session_data.crashes.saved_min, + session_data.crashes.saved_avg, + session_data.crashes.saved_max, + session_data.hangs.saved_cum, + session_data.hangs.saved_min, + session_data.hangs.saved_avg, + session_data.hangs.saved_max, + session_data.corpus.count_cum, + session_data.corpus.count_min, + session_data.corpus.count_avg, + session_data.corpus.count_max + ); + + Paragraph::new(content) + .block( + Block::default() + .title("Overall results") + .borders(Borders::ALL) + .add_modifier(Modifier::BOLD), + ) + .style(Style::default()) + } -fn create_stage_progress_paragraph(session_data: &SessionData) -> Paragraph { - let content = format!( - "Execs: {} ({}->{}<-{}) + /// Creates the stage progress paragraph + fn create_stage_progress_paragraph(session_data: &CampaignData) -> Paragraph { + let content = format!( + "Execs: {} ({}->{}<-{}) Execs/s: {:.2} ({:.2}->{:.2}<-{:.2}), Coverage: {:.2}% ({:.2}%/{:.2}%)", - session_data.executions.cum, - session_data.executions.min, - session_data.executions.avg, - session_data.executions.max, - session_data.executions.ps_cum, - session_data.executions.ps_min, - session_data.executions.ps_avg, - session_data.executions.ps_max, - session_data.coverage.bitmap_avg, - session_data.coverage.bitmap_min, - session_data.coverage.bitmap_max, - ); - - Paragraph::new(content) - .block( - Block::default() - .title("Stage Progress") - .borders(Borders::ALL) - .add_modifier(Modifier::BOLD), - ) - .style(Style::default()) -} + session_data.executions.cum, + session_data.executions.min, + session_data.executions.avg, + session_data.executions.max, + session_data.executions.ps_cum, + session_data.executions.ps_min, + session_data.executions.ps_avg, + session_data.executions.ps_max, + session_data.coverage.bitmap_avg, + session_data.coverage.bitmap_min, + session_data.coverage.bitmap_max, + ); + + Paragraph::new(content) + .block( + Block::default() + .title("Stage Progress") + .borders(Borders::ALL) + .add_modifier(Modifier::BOLD), + ) + .style(Style::default()) + } -fn create_nerd_stats_paragraph(session_data: &SessionData) -> Paragraph { - let content = format!( - "Levels: {} ({}/{}) + /// Creates the nerd stats paragraph + fn create_nerd_stats_paragraph(session_data: &CampaignData) -> Paragraph { + let content = format!( + "Levels: {} ({}/{}) Pending favorites: {} ({}->{}<-{}) Pending total: {} ({}->{}<-{}), Cycles without finds: {} ({}/{})", - session_data.levels.avg, - session_data.levels.min, - session_data.levels.max, - session_data.pending.favorites_cum, - session_data.pending.favorites_min, - session_data.pending.favorites_avg, - session_data.pending.favorites_max, - session_data.pending.total_cum, - session_data.pending.total_min, - session_data.pending.total_avg, - session_data.pending.total_max, - session_data.cycles.wo_finds_avg, - session_data.cycles.wo_finds_min, - session_data.cycles.wo_finds_max - ); - - Paragraph::new(content) - .block(Block::default().title("Nerd Stats").borders(Borders::ALL)) - .style(Style::default()) -} + session_data.levels.avg, + session_data.levels.min, + session_data.levels.max, + session_data.pending.favorites_cum, + session_data.pending.favorites_min, + session_data.pending.favorites_avg, + session_data.pending.favorites_max, + session_data.pending.total_cum, + session_data.pending.total_min, + session_data.pending.total_avg, + session_data.pending.total_max, + session_data.cycles.wo_finds_avg, + session_data.cycles.wo_finds_min, + session_data.cycles.wo_finds_max + ); + + Paragraph::new(content) + .block(Block::default().title("Nerd Stats").borders(Borders::ALL)) + .style(Style::default()) + } -fn format_last_event(events: &[CrashInfoDetails], total_run_time: &Duration) -> String { - if events.is_empty() { - "N/A".to_string() - } else { - let event_time = *total_run_time - Duration::from_millis(events[0].time); - format_duration(event_time) + /// Formats the last event duration + fn format_last_event(events: &[CrashInfoDetails], total_run_time: &Duration) -> String { + if events.is_empty() { + "N/A".to_string() + } else { + let event_time = *total_run_time - Duration::from_millis(events[0].time); + Self::format_duration(event_time) + } } -} -fn format_solutions(solutions: &[CrashInfoDetails]) -> String { - solutions - .iter() - .map(|s| { - format!( - "{} | SIG: {} | TIME: {} | EXEC: {} | SRC: {} | OP: {} | REP: {}", - s.fuzzer_name, - s.sig.clone().unwrap_or_else(|| "-".to_string()), - s.time, - s.execs, - s.src, - s.op, - s.rep - ) - }) - .collect::>() - .join("\n") -} + /// Formats the solutions into a string + fn format_solutions(solutions: &[CrashInfoDetails]) -> String { + solutions + .iter() + .map(|s| { + format!( + "{} | SIG: {} | TIME: {} | EXEC: {} | SRC: {} | OP: {} | REP: {}", + s.fuzzer_name, + s.sig.clone().unwrap_or_else(|| "-".to_string()), + s.time, + s.execs, + s.src, + s.op, + s.rep + ) + }) + .collect::>() + .join("\n") + } -fn format_duration(duration: Duration) -> String { - let secs = duration.as_secs(); - let days = secs / 86400; - let hours = (secs % 86400) / 3600; - let mins = (secs % 3600) / 60; - let secs = secs % 60; + /// Formats a duration into a string + fn format_duration(duration: Duration) -> String { + let secs = duration.as_secs(); + let days = secs / 86400; + let hours = (secs % 86400) / 3600; + let mins = (secs % 3600) / 60; + let secs = secs % 60; - format!("{days} days, {hours:02}:{mins:02}:{secs:02}") + format!("{days} days, {hours:02}:{mins:02}:{secs:02}") + } } diff --git a/src/utils.rs b/src/utils.rs index 74015c3..14282c7 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -4,6 +4,7 @@ use std::hash::{DefaultHasher, Hasher}; use std::path::PathBuf; use anyhow::bail; +use anyhow::Context; use anyhow::Result; use crate::afl_cmd_gen::AFLCmdGenerator; @@ -13,6 +14,11 @@ use crate::cli::RunArgs; use crate::cli::AFL_CORPUS; use crate::harness::Harness; +/// Creates a new `Harness` instance based on the provided `GenArgs`. +/// +/// # Errors +/// +/// Returns an error if the target binary is not specified. pub fn create_harness(args: &GenArgs) -> Result { if args.target.is_none() { bail!("Target binary is required"); @@ -26,6 +32,10 @@ pub fn create_harness(args: &GenArgs) -> Result { )) } +/// Creates a new `AFLCmdGenerator` instance based on the provided `GenArgs` and `Harness`. +/// +/// If the input directory is not specified, it defaults to `AFL_CORPUS`. +/// If the output directory is not specified, it defaults to `/tmp/afl_output`. pub fn create_afl_runner( args: &GenArgs, harness: Harness, @@ -34,7 +44,6 @@ pub fn create_afl_runner( AFLCmdGenerator::new( harness, args.runners.unwrap_or(1), - args.afl_binary.clone(), args.input_dir .clone() .unwrap_or_else(|| PathBuf::from(AFL_CORPUS)), @@ -43,9 +52,14 @@ pub fn create_afl_runner( .unwrap_or_else(|| PathBuf::from("/tmp/afl_output")), args.dictionary.clone(), raw_afl_flags, + args.afl_binary.clone(), ) } +/// Generates a unique tmux session name based on the provided `RunArgs` and `target_args`. +/// +/// If the `tmux_session_name` is not specified in `RunArgs`, the function generates a unique name +/// by combining the target binary name, input directory name, and a hash of the `target_args`. pub fn generate_tmux_name(args: &RunArgs, target_args: &str) -> String { args.tmux_session_name.as_ref().map_or_else( || { @@ -75,26 +89,44 @@ pub fn generate_tmux_name(args: &RunArgs, target_args: &str) -> String { ) } -pub fn load_config(config_path: Option<&PathBuf>) -> Config { - config_path.map_or_else( - || { - let cwd = env::current_dir().unwrap(); - let default_config_path = cwd.join("aflr_cfg.toml"); - if default_config_path.exists() { - let config_content = fs::read_to_string(&default_config_path).unwrap(); - toml::from_str(&config_content).unwrap() - } else { - Config::default() - } - }, - |config_path| { - // TODO: Add error handling for when config_path does not exist - let config_content = fs::read_to_string(config_path).unwrap(); - toml::from_str(&config_content).unwrap() - }, - ) +/// Loads the configuration from the specified `config_path` or the default configuration file. +/// +/// If `config_path` is `None`, the function looks for a default configuration file named `aflr_cfg.toml` +/// in the current working directory. If the default configuration file is not found, an empty `Config` +/// instance is returned. +/// +/// # Errors +/// +/// Returns an error if the configuration file cannot be read or parsed. +pub fn load_config(config_path: Option<&PathBuf>) -> Result { + if let Some(path) = config_path { + let config_content = fs::read_to_string(path) + .with_context(|| format!("Failed to read config file: {}", path.display()))?; + toml::from_str(&config_content) + .with_context(|| format!("Failed to parse config file: {}", path.display())) + } else { + let cwd = env::current_dir().context("Failed to get current directory")?; + let default_config_path = cwd.join("aflr_cfg.toml"); + if default_config_path.exists() { + let config_content = fs::read_to_string(&default_config_path).with_context(|| { + format!( + "Failed to read default config file: {}", + default_config_path.display() + ) + })?; + toml::from_str(&config_content).with_context(|| { + format!( + "Failed to parse default config file: {}", + default_config_path.display() + ) + }) + } else { + Ok(Config::default()) + } + } } +/// Prints the generated commands to the console. pub fn print_generated_commands(cmds: &[String]) { println!("Generated commands:"); for (i, cmd) in cmds.iter().enumerate() {