feat: adding new seeds to a running AFL++ campaign

src/afl/coverage.rs

@@ -520,7 +520,7 @@ impl CoverageCollector {
             .par_chunks(1000)
             .enumerate()
             .map(|(i, chunk)| {
-                let temp_output = temp_dir.path().join(format!("temp_merged_{}.profdata", i));
+                let temp_output = temp_dir.path().join(format!("temp_merged_{i}.profdata"));
 
                 let output = Command::new("llvm-profdata")
                     .arg("merge")

src/argument_aggregator.rs

@@ -1,4 +1,4 @@
-use crate::cli::{ArgMerge, Args, CovArgs, GenArgs, RunArgs};
+use crate::cli::{AddSeedArgs, ArgMerge, Args, CovArgs, GenArgs, RunArgs};
 use anyhow::{bail, Context, Result};
 use std::{env, fs, path::PathBuf};
 
@@ -92,4 +92,15 @@ impl ArgumentAggregator {
             .as_ref()
             .map_or_else(|| args.clone(), |config| args.merge_with_config(config)))
     }
+
+    /// Merge the provided adding seeds arguments with the config
+    ///
+    /// # Errors
+    /// * If the config cannot be merged
+    pub fn merge_add_seed_args(&self, args: &AddSeedArgs) -> Result<AddSeedArgs> {
+        Ok(self
+            .config
+            .as_ref()
+            .map_or_else(|| args.clone(), |config| args.merge_with_config(config)))
+    }
 }

src/cli/add_seed.rs

@@ -0,0 +1,25 @@
+use clap::Args;
+use std::path::PathBuf;
+
+#[derive(Args, Clone, Debug, Default)]
+pub struct AddSeedArgs {
+    /// Target binary to fuzz
+    #[arg(short, long, help = "Instrumented target binary to fuzz")]
+    pub target: Option<PathBuf>,
+    /// Target binary arguments
+    #[arg(help = "Target binary arguments, including @@ if needed", raw = true)]
+    pub target_args: Option<Vec<String>>,
+    /// Output directory
+    #[arg(
+        short = 'o',
+        long,
+        help = "Solution/Crash output directory of the running campaign"
+    )]
+    pub output_dir: Option<PathBuf>,
+    /// Path to a TOML config file
+    #[arg(long, help = "Path to TOML config file")]
+    pub config: Option<PathBuf>,
+    /// Seed(s) to add to the corpus
+    #[arg(long, help = "Seed(s) to add to the corpus", value_name = "SEED(S)")]
+    pub seed: PathBuf,
+}

src/cli/mod.rs

@@ -1,6 +1,7 @@
 use clap::{Parser, Subcommand};
 use serde::Deserialize;
 
+mod add_seed;
 mod afl;
 pub mod constants;
 mod cov;
@@ -14,6 +15,7 @@ mod target;
 mod tui;
 mod utils;
 
+pub use add_seed::AddSeedArgs;
 pub use afl::AflArgs;
 use constants::{AFL_CORPUS, AFL_OUTPUT};
 pub use cov::CovArgs;
@@ -51,6 +53,8 @@ pub enum Commands {
     Tui(TuiArgs),
     /// Kills a running session and all spawned processes inside
     Kill(KillArgs),
+    /// Allows adding new seeds to a running campaign
+    AddSeed(AddSeedArgs),
 }
 
 #[derive(Deserialize, Default, Debug, Clone)]
@@ -193,6 +197,30 @@ impl ArgMerge<Self> for CovArgs {
     }
 }
 
+impl ArgMerge<Self> for AddSeedArgs {
+    fn merge_with_config(&self, args: &Args) -> Self {
+        let merge_path = |opt: Option<std::path::PathBuf>, cfg_str: Option<String>| {
+            opt.or_else(|| {
+                cfg_str
+                    .filter(|p| !p.is_empty())
+                    .map(std::path::PathBuf::from)
+            })
+        };
+
+        Self {
+            target: merge_path(self.target.clone(), args.target.path.clone()),
+            target_args: self
+                .target_args
+                .clone()
+                .or_else(|| args.target.args.clone().filter(|args| !args.is_empty())),
+            output_dir: merge_path(self.output_dir.clone(), args.afl_cfg.solution_dir.clone())
+                .or_else(|| Some(std::path::PathBuf::from(AFL_OUTPUT))),
+            config: self.config.clone(),
+            seed: self.seed.clone(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/commands/add_seed.rs

@@ -0,0 +1,103 @@
+use anyhow::{Context, Result};
+use std::path::Path;
+use std::path::PathBuf;
+use std::process;
+use tempfile::TempDir;
+
+use crate::{argument_aggregator::ArgumentAggregator, cli::AddSeedArgs, commands::Command};
+
+pub struct AddSeedCommand<'a> {
+    args: &'a AddSeedArgs,
+    arg_aggregator: &'a ArgumentAggregator,
+}
+
+impl<'a> AddSeedCommand<'a> {
+    pub fn new(args: &'a AddSeedArgs, arg_aggregator: &'a ArgumentAggregator) -> Self {
+        Self {
+            args,
+            arg_aggregator,
+        }
+    }
+
+    fn execute_add_seed_afl(
+        seed: &Path,
+        corpus_dir: &Path,
+        target: &Path,
+        target_args: &[String],
+    ) -> Result<()> {
+        let target = if target.starts_with("./") || target.starts_with("/") {
+            target.to_path_buf()
+        } else {
+            Path::new("./").join(target)
+        };
+
+        let status = process::Command::new("afl-fuzz")
+            .env("AFL_BENCH_JUST_ONE", "1")
+            .env("AFL_FAST_CAL", "1")
+            .env("AFL_IGNORE_SEED_PROBLEMS", "1")
+            .env("AFL_AUTORESUME", "1")
+            .arg("-i")
+            .arg(seed)
+            .arg("-o")
+            .arg(corpus_dir)
+            .arg("-S")
+            .arg(&uuid::Uuid::new_v4().simple().to_string()[..8])
+            .arg("--")
+            .arg(target)
+            .args(target_args)
+            .stdout(process::Stdio::null())
+            .stderr(process::Stdio::null())
+            .status()
+            .context("Failed to execute afl-fuzz")?;
+
+        if !status.success() {
+            return Err(anyhow::anyhow!(
+                "afl-fuzz failed with exit code: {}",
+                status.code().unwrap_or(-1),
+            ));
+        }
+
+        Ok(())
+    }
+
+    fn add_seed(
+        seed: &PathBuf,
+        target: &Path,
+        target_args: &[String],
+        output_dir: &Path,
+    ) -> Result<()> {
+        if !seed.exists() {
+            return Err(anyhow::anyhow!("Seed file does not exist: {:?}", seed));
+        }
+
+        if seed.is_file() {
+            let tmpdir = TempDir::new().context("Failed to create temporary directory")?;
+            let new_seed_dir = tmpdir.path();
+            std::fs::copy(seed, new_seed_dir.join(seed.file_name().unwrap()))?;
+
+            Self::execute_add_seed_afl(new_seed_dir, output_dir, target, target_args)
+        } else {
+            Self::execute_add_seed_afl(seed, output_dir, target, target_args)
+        }
+    }
+}
+
+impl Command for AddSeedCommand<'_> {
+    fn execute(&self) -> Result<()> {
+        let merged_args = self.arg_aggregator.merge_add_seed_args(self.args)?;
+
+        let target = merged_args.target.as_ref().context("Target is required")?;
+        let target_args = merged_args
+            .target_args
+            .as_ref()
+            .context("Target arguments are required")?;
+        let output_dir = merged_args
+            .output_dir
+            .as_ref()
+            .context("Output directory is required")?;
+
+        Self::add_seed(&merged_args.seed, target, target_args, output_dir)?;
+        println!("[+] Seeds added successfully");
+        Ok(())
+    }
+}

src/commands/mod.rs

@@ -1,3 +1,4 @@
+pub mod add_seed;
 pub mod cov;
 pub mod gen;
 pub mod kill;

src/commands/run.rs

@@ -64,7 +64,7 @@ impl<'a> RunCommand<'a> {
         } else {
             session.run()?;
             if !args.detached {
-                session.attach()?
+                session.attach()?;
             }
         }
         Ok(())

src/main.rs

@@ -12,8 +12,8 @@ pub mod utils;
 use argument_aggregator::ArgumentAggregator;
 use cli::{Cli, Commands};
 use commands::{
-    cov::CovCommand, gen::GenCommand, kill::KillCommand, render_tui::RenderCommand,
-    run::RunCommand, Command,
+    add_seed::AddSeedCommand, cov::CovCommand, gen::GenCommand, kill::KillCommand,
+    render_tui::RenderCommand, run::RunCommand, Command,
 };
 
 fn main() -> Result<()> {
@@ -25,6 +25,7 @@ fn main() -> Result<()> {
         Commands::Gen(args) => arg_aggregator.load(args.config.as_ref()),
         Commands::Run(args) => arg_aggregator.load(args.gen_args.config.as_ref()),
         Commands::Cov(args) => arg_aggregator.load(args.config.as_ref()),
+        Commands::AddSeed(args) => arg_aggregator.load(args.config.as_ref()),
         _ => Ok(()),
     }?;
 
@@ -35,10 +36,11 @@ fn main() -> Result<()> {
         Commands::Cov(args) => CovCommand::new(args, &arg_aggregator).execute(),
         Commands::Tui(args) => RenderCommand::new(args).execute(),
         Commands::Kill(args) => KillCommand::new(args).execute(),
+        Commands::AddSeed(args) => AddSeedCommand::new(args, &arg_aggregator).execute(),
     };
 
     if let Err(e) = result {
-        eprintln!("{}", e);
+        eprintln!("{e}");
         std::process::exit(1);
     }