[WARP] Add more options to creating signatures in projects

2 files changed, 215 insertions, 74 deletions

diff --git a/plugins/warp/src/plugin/project.rs b/plugins/warp/src/plugin/project.rs
index ea5a5600..fc95342c 100644
--- a/plugins/warp/src/plugin/project.rs
+++ b/plugins/warp/src/plugin/project.rs
@@ -1,16 +1,19 @@
+use crate::processor::{
+    new_processing_state_background_thread, CompressionTypeField, FileDataKindField,
+    FileFilterField, WarpFileProcessor,
+};
+use crate::report::{ReportGenerator, ReportKindField};
 use binaryninja::background_task::BackgroundTask;
 use binaryninja::command::ProjectCommand;
 use binaryninja::interaction::{Form, FormInputField};
+use binaryninja::project::folder::ProjectFolder;
 use binaryninja::project::Project;
 use binaryninja::rc::Ref;
 use regex::Regex;
+use std::path::{Path, PathBuf};
 use std::thread;
 use std::time::Instant;
-
-use crate::processor::{
-    new_processing_state_background_thread, FileDataKindField, FileFilterField, WarpFileProcessor,
-};
-use crate::report::{ReportGenerator, ReportKindField};
+use warp::WarpFile;
 
 pub struct CreateSignaturesForm {
     form: Form,
@@ -22,6 +25,9 @@ impl CreateSignaturesForm {
         form.add_field(Self::file_data_field());
         form.add_field(Self::file_filter_field());
         form.add_field(Self::generated_report_field());
+        form.add_field(Self::compression_type_field());
+        form.add_field(Self::save_individual_files_field());
+        form.add_field(Self::skip_existing_warp_files_field());
         // TODO: Threads (we run the analysis in the background)
         Self { form }
     }
@@ -50,6 +56,48 @@ impl CreateSignaturesForm {
         ReportKindField::from_form(&self.form).unwrap_or_default()
     }
 
+    pub fn compression_type_field() -> FormInputField {
+        CompressionTypeField::default().to_field()
+    }
+
+    pub fn compression_type(&self) -> CompressionTypeField {
+        CompressionTypeField::from_form(&self.form).unwrap_or_default()
+    }
+
+    pub fn save_individual_files_field() -> FormInputField {
+        FormInputField::Checkbox {
+            prompt: "Save individual files".to_string(),
+            default: None,
+            value: false,
+        }
+    }
+
+    pub fn save_individual_files(&self) -> bool {
+        let field = self.form.get_field_with_name("Save individual files");
+        let field_value = field.and_then(|f| f.try_value_int()).unwrap_or(0);
+        match field_value {
+            1 => true,
+            _ => false,
+        }
+    }
+
+    pub fn skip_existing_warp_files_field() -> FormInputField {
+        FormInputField::Checkbox {
+            prompt: "Skip existing WARP files".to_string(),
+            default: Some(true),
+            value: false,
+        }
+    }
+
+    pub fn skip_existing_warp_files(&self) -> bool {
+        let field = self.form.get_field_with_name("Skip existing WARP files");
+        let field_value = field.and_then(|f| f.try_value_int()).unwrap_or(0);
+        match field_value {
+            1 => true,
+            _ => false,
+        }
+    }
+
     pub fn prompt(&mut self) -> bool {
         self.form.prompt()
     }
@@ -65,47 +113,91 @@ impl CreateSignatures {
         }
         let file_data_kind = form.file_data_kind();
         let report_kind = form.generated_report_kind();
+        let compression_type = form.compression_type();
+        let save_individual_files = form.save_individual_files();
 
-        let mut processor = WarpFileProcessor::new().with_file_data(file_data_kind);
+        // Save the warp file to the project.
+        let save_warp_file = move |project: &Project,
+                                   folder: Option<&ProjectFolder>,
+                                   name: &str,
+                                   warp_file: &WarpFile| {
+            if project
+                .create_file(&warp_file.to_bytes(), folder, name, "")
+                .is_err()
+            {
+                log::error!("Failed to create project file!");
+            }
 
-        // This thread will show the state in a background task.
-        let background_task = BackgroundTask::new("Processing started...", true);
-        new_processing_state_background_thread(background_task.clone(), processor.state());
+            let report = ReportGenerator::new();
+            if let Some(generated) = report.report(&report_kind, &warp_file) {
+                let ext = report.report_extension(&report_kind).unwrap_or_default();
+                let file_name = format!("{}_report.{}", name, ext);
+                if project
+                    .create_file(&generated.into_bytes(), folder, &file_name, "Warp file")
+                    .is_err()
+                {
+                    log::error!("Failed to create project file!");
+                }
+            }
+        };
 
+        // Optional callback for saving off the individual project files.
+        let callback_project = project.clone();
+        let save_individual_files_cb = move |path: &Path, file: &WarpFile| {
+            if file.chunks.is_empty() {
+                log::debug!("Skipping empty file: {}", path.display());
+                return;
+            }
+            // The path returned will be the one on disk, so we will go and grab the project for it.
+            let Some(project_file) = callback_project.file_by_path(path) else {
+                log::error!("Failed to find project file for path: {}", path.display());
+                return;
+            };
+            let project_file = project_file.to_owned();
+            let file_name = format!("{}.warp", project_file.name());
+            let project_folder = project_file.folder();
+            save_warp_file(
+                &callback_project,
+                project_folder.as_deref(),
+                &file_name,
+                file,
+            );
+        };
+
+        let mut processor = WarpFileProcessor::new()
+            .with_file_data(file_data_kind)
+            .with_compression_type(compression_type);
+
+        if save_individual_files {
+            processor = processor.with_processed_file_callback(save_individual_files_cb);
+        }
+
+        // Construct the user-supplied file filter, we also have an appended filter for warp files.
+        let mut filter = form.file_filter();
+        // This checkbox is here as this is very common to filter for.
+        // And we want to do it by default.
+        if form.skip_existing_warp_files() {
+            let warp_filter = Regex::new(".*\\.warp").unwrap();
+            filter = match filter {
+                Some(existing) => {
+                    let combined = format!("{}|.*\\.warp", existing.as_str());
+                    Some(Regex::new(&combined).unwrap())
+                }
+                None => Some(warp_filter),
+            };
+        }
         if let Some(filter) = form.file_filter() {
             processor = processor.with_file_filter(filter);
         }
 
+        // This thread will show the state in a background task.
+        let background_task = BackgroundTask::new("Processing started...", true);
+        new_processing_state_background_thread(background_task.clone(), processor.state());
+
         let start = Instant::now();
         match processor.process_project(&project) {
             Ok(warp_file) => {
-                // Print the processor string into the description of the file, so we know how it was generated.
-                let processor_str = format!("{:#?}", &processor);
-
-                // TODO: File name needs to be configurable.
-                if project
-                    .create_file(
-                        &warp_file.to_bytes(),
-                        None,
-                        "generated.warp",
-                        &processor_str,
-                    )
-                    .is_err()
-                {
-                    log::error!("Failed to create project file!");
-                }
-
-                let report = ReportGenerator::new();
-                if let Some(generated) = report.report(&report_kind, &warp_file) {
-                    let ext = report.report_extension(&report_kind).unwrap_or_default();
-                    let file_name = format!("report.{}", ext);
-                    if project
-                        .create_file(&generated.into_bytes(), None, &file_name, "Warp file")
-                        .is_err()
-                    {
-                        log::error!("Failed to create project file!");
-                    }
-                }
+                save_warp_file(&project, None, "generated.warp", &warp_file);
             }
             Err(e) => {
                 log::error!("Failed to process project: {}", e);
diff --git a/plugins/warp/src/processor.rs b/plugins/warp/src/processor.rs
index 4b973753..056b8a89 100644
--- a/plugins/warp/src/processor.rs
+++ b/plugins/warp/src/processor.rs
@@ -1,4 +1,5 @@
 use std::collections::{HashMap, HashSet};
+use std::fmt::{Debug, Formatter};
 use std::fs::File;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::Ordering::Relaxed;
@@ -285,6 +286,10 @@ pub enum ProcessingFileState {
     Processed,
 }
 
+/// A callback for when a file has been processed, use this if you intend to save off individual
+/// files inside a directory, project or archive.
+pub type ProcessedFileCallback = Arc<dyn Fn(&Path, &WarpFile) + Send + Sync>;
+
 #[derive(Debug, Default)]
 pub struct ProcessingState {
     pub cancelled: AtomicBool,
@@ -319,7 +324,7 @@ impl ProcessingState {
 }
 
 /// Create a new [`WarpFile`] from files, projects, and directories.
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct WarpFileProcessor {
     /// The Binary Ninja settings to use when analyzing the binaries.
     analysis_settings: Value,
@@ -333,6 +338,7 @@ pub struct WarpFileProcessor {
     file_data: FileDataKindField,
     included_functions: IncludedFunctionsField,
     compression_type: CompressionTypeField,
+    processed_file_callback: Option<ProcessedFileCallback>,
     /// Regex pattern used to filter out files.
     file_filter: Option<Regex>,
     /// Processor state, this is shareable between threads, so the processor and the consumer can
@@ -357,6 +363,7 @@ impl WarpFileProcessor {
             file_data: Default::default(),
             included_functions: Default::default(),
             compression_type: Default::default(),
+            processed_file_callback: None,
             file_filter: None,
             state: Arc::new(ProcessingState::default()),
         }
@@ -397,6 +404,14 @@ impl WarpFileProcessor {
         self
     }
 
+    pub fn with_processed_file_callback(
+        mut self,
+        processed_file_callback: impl Fn(&Path, &WarpFile) + Send + Sync + 'static,
+    ) -> Self {
+        self.processed_file_callback = Some(Arc::new(processed_file_callback));
+        self
+    }
+
     pub fn with_file_filter(mut self, file_filter: Regex) -> Self {
         self.file_filter = Some(file_filter);
         self
@@ -417,14 +432,34 @@ impl WarpFileProcessor {
         }
     }
 
+    pub fn merge_files(
+        &self,
+        files: Vec<WarpFile<'static>>,
+    ) -> Result<WarpFile<'static>, ProcessingError> {
+        let chunks: Vec<_> = files.into_iter().flat_map(|f| f.chunks.clone()).collect();
+        let merged_chunks = Chunk::merge(&chunks, self.compression_type.into());
+        Ok(WarpFile::new(WarpFileHeader::new(), merged_chunks))
+    }
+
     pub fn process(&self, path: PathBuf) -> Result<WarpFile<'static>, ProcessingError> {
-        match path.extension() {
-            Some(ext) if ext == "a" || ext == "lib" || ext == "rlib" => self.process_archive(path),
-            Some(ext) if ext == "warp" => self.process_warp_file(path),
+        let file = match path.extension() {
+            Some(ext) if ext == "a" || ext == "lib" || ext == "rlib" => {
+                self.process_archive(path.clone())
+            }
+            Some(ext) if ext == "warp" => self.process_warp_file(path.clone()),
             _ if path.is_dir() => self.process_directory(&path),
             // TODO: process_database?
-            _ => self.process_file(path),
+            _ => self.process_file(path.clone()),
+        }?;
+
+        // We do this right after we process the file so that all possible paths to this will be caught.
+        // This callback is typically used to write the file out to some other place for caching or
+        // for distributing smaller unmerged files.
+        if let Some(callback) = &self.processed_file_callback {
+            callback(&path, &file);
         }
+
+        Ok(file)
     }
 
     pub fn process_project(&self, project: &Project) -> Result<WarpFile<'static>, ProcessingError> {
@@ -465,12 +500,7 @@ impl WarpFileProcessor {
             })
             .collect();
 
-        let unmerged_chunks: Vec<_> = unmerged_files?
-            .iter()
-            .flat_map(|f| f.chunks.clone())
-            .collect();
-        let merged_chunks = Chunk::merge(&unmerged_chunks, self.compression_type.into());
-        Ok(WarpFile::new(WarpFileHeader::new(), merged_chunks))
+        self.merge_files(unmerged_files?)
     }
 
     pub fn process_project_file(
@@ -482,11 +512,22 @@ impl WarpFileProcessor {
         let path = project_file
             .path_on_disk()
             .ok_or_else(|| ProcessingError::NoPathToProjectFile(project_file.to_owned()))?;
-        match extension {
-            Some(ext) if ext == "a" || ext == "lib" || ext == "rlib" => self.process_archive(path),
-            Some("warp") => self.process_warp_file(path),
-            _ => self.process_file(path),
+        let file = match extension {
+            Some(ext) if ext == "a" || ext == "lib" || ext == "rlib" => {
+                self.process_archive(path.clone())
+            }
+            Some("warp") => self.process_warp_file(path.clone()),
+            _ => self.process_file(path.clone()),
+        }?;
+
+        // We do this right after we process the file so that all possible paths to this will be caught.
+        // This callback is typically used to write the file out to some other place for caching or
+        // for distributing smaller unmerged files.
+        if let Some(callback) = &self.processed_file_callback {
+            callback(&path, &file);
         }
+
+        Ok(file)
     }
 
     pub fn process_warp_file(&self, path: PathBuf) -> Result<WarpFile<'static>, ProcessingError> {
@@ -601,12 +642,7 @@ impl WarpFileProcessor {
             })
             .collect();
 
-        let unmerged_chunks: Vec<_> = unmerged_files?
-            .iter()
-            .flat_map(|f| f.chunks.clone())
-            .collect();
-        let merged_chunks = Chunk::merge(&unmerged_chunks, self.compression_type.into());
-        Ok(WarpFile::new(WarpFileHeader::new(), merged_chunks))
+        self.merge_files(unmerged_files?)
     }
 
     pub fn process_archive(&self, path: PathBuf) -> Result<WarpFile<'static>, ProcessingError> {
@@ -650,7 +686,6 @@ impl WarpFileProcessor {
                 .set_file_state(entry_file.clone(), ProcessingFileState::Unprocessed);
         }
 
-        // TODO: Par iter?
         // Process all the entries.
         let unmerged_files: Result<Vec<_>, _> = entry_files
             .into_par_iter()
@@ -669,12 +704,7 @@ impl WarpFileProcessor {
             })
             .collect();
 
-        let unmerged_chunks: Vec<_> = unmerged_files?
-            .iter()
-            .flat_map(|f| f.chunks.clone())
-            .collect();
-        let merged_chunks = Chunk::merge(&unmerged_chunks, self.compression_type.into());
-        Ok(WarpFile::new(WarpFileHeader::new(), merged_chunks))
+        self.merge_files(unmerged_files?)
     }
 
     pub fn process_view(
@@ -689,20 +719,25 @@ impl WarpFileProcessor {
         if self.file_data != FileDataKindField::Types {
             let mut signature_chunks = self.create_signature_chunks(view)?;
             for (target, signature_chunk) in signature_chunks.drain() {
-                let chunk = Chunk::new_with_target(
-                    ChunkKind::Signature(signature_chunk),
-                    self.compression_type.into(),
-                    target,
-                );
-                chunks.push(chunk)
+                if signature_chunk.raw_functions().count() != 0 {
+                    let chunk = Chunk::new_with_target(
+                        ChunkKind::Signature(signature_chunk),
+                        self.compression_type.into(),
+                        target,
+                    );
+                    chunks.push(chunk)
+                }
             }
         }
 
         if self.file_data != FileDataKindField::Signatures {
-            chunks.push(Chunk::new(
-                ChunkKind::Type(self.create_type_chunk(view)?),
-                self.compression_type.into(),
-            ));
+            let type_chunk = self.create_type_chunk(view)?;
+            if type_chunk.raw_types().count() != 0 {
+                chunks.push(Chunk::new(
+                    ChunkKind::Type(type_chunk),
+                    self.compression_type.into(),
+                ));
+            }
         }
 
         self.state
@@ -814,6 +849,20 @@ impl WarpFileProcessor {
     }
 }
 
+impl Debug for WarpFileProcessor {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("WarpFileProcessor")
+            .field("file_data", &self.file_data)
+            .field("compression_type", &self.compression_type)
+            .field("included_functions", &self.included_functions)
+            .field("file_filter", &self.file_filter)
+            .field("state", &self.state)
+            .field("cache_path", &self.cache_path)
+            .field("analysis_settings", &self.analysis_settings)
+            .finish()
+    }
+}
+
 fn project_file_path(file: &ProjectFile) -> PathBuf {
     // Recurse up the folders to build a string like /foldera/folderb/myfile
     let mut path = PathBuf::new();