[WARP] Add example of batch processing using the headless Rust API

5 files changed, 243 insertions, 40 deletions

diff --git a/plugins/warp/Cargo.toml b/plugins/warp/Cargo.toml
index 1e4ecb5e..7a7df157 100644
--- a/plugins/warp/Cargo.toml
+++ b/plugins/warp/Cargo.toml
@@ -18,7 +18,7 @@ log = "0.4"
 itertools = "0.14.0"
 dashmap = { version = "6.1", features = ["rayon"]}
 rayon = "1.10"
-arboard = "3.4"
+arboard = {version = "3.4", default-features = false}
 walkdir = "2.5"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
diff --git a/plugins/warp/README.md b/plugins/warp/README.md
index 6506a7c3..d1193a9e 100644
--- a/plugins/warp/README.md
+++ b/plugins/warp/README.md
@@ -2,43 +2,6 @@
 
 Provides integration with [WARP](https://github.com/Vector35/warp), more specifically the function identification and associated type information.
 
-## Batch Processing
+## Documentation
 
-`sigem` (see code at `src/bin/sigem.rs`) allows for batch processing binaries for signature generation.
-
-The only real "magic" here is the cross-linking of symbols in archives, this means that static libraries will have much finer constraint matching.
-
-Duplicate functions will be removed to save space, any input will always produce a single output signature file, if you want separate signature files, invoke individually.
-
-### Building
-
-> Assuming you have the following:
-> - A compatible Binary Ninja with headless usage (see [this documentation](https://docs.binary.ninja/dev/batch.html#batch-processing-and-other-automation-tips) for more information)
-> - Clang
-> - Rust (currently tested for 1.83.0)
-> - Set `BINARYNINJADIR` env variable to your install directory (see [here](https://docs.binary.ninja/guide/#binary-path) for more details)
-
-1. Clone this repository (`git clone https://github.com/Vector35/binaryninja-api/tree/dev`)
-2. Navigate to this plugin (`cd plugins/warp`)
-3. Build in release (`cargo build --release`)
-
-If compilation fails because it could not link against binaryninjacore than you should double-check you set `BINARYNINJADIR` correctly.
-
-Once it finishes you now will have a `sigem` binary in `target/release` for use.
-
-### Usage
-
-> Assuming you already have the `sigem` binary and a valid headless compatible Binary Ninja license.
-
-To create a signature file simply pass the input as the first positional argument to `sigem`:
-
-- A regular binary
-- An archive of binaries (`.a`, `.lib`, `.rlib`)
-- A directory of binaries
-- A BNDB
-
-Example: `./sigem mylibrary.a` or `./sigem ./all-libs/`
-
-Once its finished you should see a `.sbin` file next to the input file, this can be moved into the corresponding signature folder (see the [user docs](https://docs.binary.ninja/dev/annotation.html?h=install+path#signature-library) for more info)
-
-If you encounter malloc errors or instability try and adjust the number of parallel threads using `RAYON_NUM_THREADS` environment variable (ex. `RAYON_NUM_THREADS=1 ./sigem mylib.a`)
\ No newline at end of file
+Documentation is available in the [User Guide](../../docs/guide/warp.md).
\ No newline at end of file
diff --git a/plugins/warp/examples/headless/Cargo.toml b/plugins/warp/examples/headless/Cargo.toml
new file mode 100644
index 00000000..b28b089d
--- /dev/null
+++ b/plugins/warp/examples/headless/Cargo.toml
@@ -0,0 +1,14 @@
+[package]
+name = "warp_headless"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+clap = { version = "4.5", features = ["derive"] }
+indicatif = "0.18.2"
+warp_ninja.workspace = true
+binaryninjacore-sys.workspace = true
+binaryninja.workspace = true
+log = "0.4"
+env_logger = "0.11"
+ctrlc = "3.5"
\ No newline at end of file
diff --git a/plugins/warp/examples/headless/README.md b/plugins/warp/examples/headless/README.md
new file mode 100644
index 00000000..be4278b5
--- /dev/null
+++ b/plugins/warp/examples/headless/README.md
@@ -0,0 +1,44 @@
+# Headless WARP Processor
+
+Provides headless support for generating WARP signatures using Binary Ninja.
+
+### Building
+
+> Assuming you have the following:
+> - A compatible Binary Ninja with headless usage (see [this documentation](https://docs.binary.ninja/dev/batch.html#batch-processing-and-other-automation-tips) for more information)
+> - Clang
+> - Rust (currently tested for 1.83.0)
+> - Set `BINARYNINJADIR` env variable to your installation directory (see [here](https://docs.binary.ninja/guide/#binary-path) for more details)
+>   - If this is not set, the -sys crate will try and locate using the default installation path and last run location.
+
+1. Clone this repository (`git clone https://github.com/Vector35/binaryninja-api/tree/dev`)
+2. Build in release (`cargo build --release`)
+
+If compilation fails because it could not link against binaryninjacore than you should double-check you set `BINARYNINJADIR` correctly.
+
+Once it finishes you now will have a `warp_headless` binary in `target/release` for use.
+
+### Usage
+
+> Assuming you already have the `warp_headless` binary and a valid headless compatible Binary Ninja license.
+
+To create a signature file simply pass the input as the first positional argument to `warp_headless`, as well as the output file using `-o` or `--output`:
+
+- A regular binary
+- An archive of binaries (`.a`, `.lib`, `.rlib`)
+- A directory of binaries
+- A BNDB
+- A WARP file
+
+Examples: 
+
+- `./warp_headless -o ./results mylibrary.a`
+  - Places a `mylibrary.warp` file in the `results` directory
+- `./warp_headless -o ./signatures.warp ./all-libs/`
+  - Provided files can be merged into a single output file
+- `./warp_headless -o ./signatures.warp ./all-libs/ mylibrary.bndb`
+  - Multiple inputs can be provided in any order
+
+Once its finished you should see a `.warp` file in the provided output location, this can be moved into the corresponding signature folder (see the [user docs](https://docs.binary.ninja/dev/annotation.html?h=install+path#signature-library) for more info)
+
+If you encounter malloc errors or instability try and adjust the number of parallel threads using `RAYON_NUM_THREADS` environment variable (ex. `RAYON_NUM_THREADS=1 ./warp_headless -o ./sigs.warp mylib.a`)
\ No newline at end of file
diff --git a/plugins/warp/examples/headless/src/main.rs b/plugins/warp/examples/headless/src/main.rs
new file mode 100644
index 00000000..c08103a5
--- /dev/null
+++ b/plugins/warp/examples/headless/src/main.rs
@@ -0,0 +1,182 @@
+use binaryninja::headless::Session;
+use clap::Parser;
+use indicatif::{ProgressBar, ProgressStyle};
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::time::Duration;
+use warp_ninja::processor::{
+    CompressionTypeField, ProcessingFileState, ProcessingState, WarpFileProcessor,
+};
+use warp_ninja::warp::WarpFile;
+
+/// Generate WARP files using Binary Ninja
+///
+/// Examples:
+///
+/// - Merge multiple inputs into a single output:
+///   ./binaryninja-warp-headless -o result.warp ./bin1 ./bin2
+///
+/// - Write one .warp per input into a directory:
+///   ./binaryninja-warp-headless -o ./results ./bin1 ./bin2 ./dir_with_bins
+///
+/// - Skip already-existing .warp files inside inputs:
+///   ./binaryninja-warp-headless -o ./results --skip-warp-files ./dir_with_mixed_inputs
+///
+/// - Use a cache directory to speed up repeated runs:
+///   ./binaryninja-warp-headless -o ./results --cache-dir ./bn-cache ./binaries
+#[derive(Parser, Debug)]
+#[command(version, about, long_about)]
+struct Args {
+    /// Input files and directories (positional).
+    #[arg(value_name = "PATH", required = true)]
+    input: Vec<PathBuf>,
+
+    /// Output destination: file path or directory for per-input output
+    #[arg(short = 'o', long = "output")]
+    output: PathBuf,
+
+    /// Skip processing of any existing WARP files within any input path
+    #[arg(long, default_value_t = false)]
+    skip_warp_files: bool,
+
+    /// Optional directory to cache analysis information to
+    #[arg(long)]
+    cache_dir: Option<PathBuf>,
+
+    /// Whether to compress the output file data
+    #[arg(short, long, default_value_t = true)]
+    compressed: bool,
+}
+
+fn main() {
+    let _session = Session::new().expect("Failed to create session");
+    let args = Args::parse();
+    env_logger::init();
+
+    let compression_ty = match args.compressed {
+        true => CompressionTypeField::Zstd,
+        false => CompressionTypeField::None,
+    };
+    let mut processor = WarpFileProcessor::new()
+        .with_skip_warp_files(args.skip_warp_files)
+        .with_compression_type(compression_ty);
+    if let Some(cache_dir) = args.cache_dir {
+        processor = processor.with_cache_path(cache_dir);
+    }
+
+    // Cancel the processor on ctrl+c, otherwise it will block forever.
+    let ctrlc_state = processor.state().clone();
+    ctrlc::set_handler(move || {
+        ctrlc_state.cancel();
+    })
+    .expect("Error setting Ctrl-C handler");
+
+    let progress_guard = run_progress_bar(processor.state());
+
+    let outputs: HashMap<PathBuf, WarpFile<'static>> = args
+        .input
+        .into_iter()
+        .filter_map(|i| match processor.process(i.clone()) {
+            Ok(o) => Some((i, o)),
+            Err(err) => {
+                log::error!("{}", err);
+                None
+            }
+        })
+        .collect();
+
+    // Stop progress UI
+    progress_guard.finish();
+
+    match args.output.is_dir() {
+        true => {
+            // Given a directory, place each output individually using the input name.
+            for (input, output) in &outputs {
+                let output_name = match input.file_name() {
+                    Some(name) => name,
+                    None => input.components().last().unwrap().as_os_str(),
+                };
+                let output_path = args.output.join(output_name).with_extension("warp");
+                log::info!("Writing to {:?}", output_path);
+                std::fs::write(output_path, output.to_bytes()).unwrap();
+            }
+        }
+        false => {
+            // Given a non-existing directory, merge all outputs and place at the output path.
+            match processor.merge_files(outputs.values().cloned().collect()) {
+                Ok(output) => {
+                    log::info!("Writing to {:?}", args.output);
+                    std::fs::write(args.output, output.to_bytes()).unwrap();
+                }
+                Err(err) => {
+                    log::error!("{}", err);
+                }
+            }
+        }
+    }
+}
+
+// TODO: Also poll for background tasks and display them as independent progress bars.
+fn run_progress_bar(state: Arc<ProcessingState>) -> ProgressGuard {
+    let pb = ProgressBar::new(0);
+    pb.set_style(
+        ProgressStyle::with_template(
+            "{spinner} {bar:40.cyan/blue} {pos}/{len} ({percent}%) [{elapsed_precise}] {msg}",
+        )
+        .unwrap()
+        .progress_chars("=>-"),
+    );
+
+    let pb_clone = pb.clone();
+    let state_clone = state.clone();
+    let handle = std::thread::spawn(move || loop {
+        let total = state_clone.total_files() as u64;
+        let done = state_clone.files_with_state(ProcessingFileState::Processed) as u64;
+        let unprocessed = state_clone.files_with_state(ProcessingFileState::Unprocessed);
+        let analyzing = state_clone.files_with_state(ProcessingFileState::Analyzing);
+        let processing = state_clone.files_with_state(ProcessingFileState::Processing);
+
+        if pb_clone.length().unwrap_or(0) != total {
+            pb_clone.set_length(total);
+        }
+        pb_clone.set_position(done.min(total));
+        pb_clone.set_message(format!(
+            "{{u:{}|a:{}|p:{}|d:{}}}",
+            unprocessed, analyzing, processing, done
+        ));
+
+        if pb_clone.is_finished() {
+            break;
+        }
+        std::thread::sleep(Duration::from_millis(100));
+    });
+
+    ProgressGuard {
+        pb,
+        handle: Some(handle),
+    }
+}
+
+struct ProgressGuard {
+    pb: ProgressBar,
+    handle: Option<std::thread::JoinHandle<()>>,
+}
+
+impl ProgressGuard {
+    fn finish(mut self) {
+        self.pb.finish_and_clear();
+        if let Some(h) = self.handle.take() {
+            let _ = h.join();
+        }
+    }
+}
+
+impl Drop for ProgressGuard {
+    fn drop(&mut self) {
+        self.pb.finish_and_clear();
+        if let Some(h) = self.handle.take() {
+            let _ = h.join();
+        }
+    }
+}