1 files changed, 211 insertions, 73 deletions

diff --git a/plugins/warp/src/plugin/create.rs b/plugins/warp/src/plugin/create.rs
index 1dd83e6e..9ebcba2d 100644
--- a/plugins/warp/src/plugin/create.rs
+++ b/plugins/warp/src/plugin/create.rs
@@ -1,94 +1,232 @@
-use crate::cache::{cached_function, cached_type_references};
-use crate::matcher::invalidate_function_matcher_cache;
-use crate::user_signature_dir;
+use crate::processor::{
+    new_processing_state_background_thread, CompressionTypeField, FileDataKindField,
+    IncludedFunctionsField, SaveReportToDiskField, WarpFileProcessor,
+};
+use crate::report::{ReportGenerator, ReportKindField};
+use crate::{user_signature_dir, INCLUDE_TAG_NAME};
+use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::command::Command;
-use binaryninja::function::Function;
-use binaryninja::rc::Guard;
-use rayon::prelude::*;
-use std::sync::atomic::AtomicUsize;
-use std::sync::atomic::Ordering::Relaxed;
+use binaryninja::interaction::form::{Form, FormInputField};
+use binaryninja::interaction::{MessageBoxButtonResult, MessageBoxButtonSet, MessageBoxIcon};
+use binaryninja::rc::Ref;
+use std::path::PathBuf;
 use std::thread;
-use std::time::Instant;
+use warp::chunk::Chunk;
+use warp::WarpFile;
 
-pub struct CreateSignatureFile;
+pub struct SaveFileField;
 
-// TODO: Prompt the user to add the newly created signature file to the signature blacklist (so that it doesn't keep getting applied)
+impl SaveFileField {
+    pub fn field(view: &BinaryView) -> FormInputField {
+        let default_name = view
+            .file()
+            .filename()
+            .split('/')
+            .last()
+            .unwrap_or("file")
+            .to_string();
+        let signature_dir = user_signature_dir();
+        let default_file_path = signature_dir.join(&default_name).with_extension("warp");
+        FormInputField::SaveFileName {
+            prompt: "File Path".to_string(),
+            // TODO: This is called extension but is really a filter.
+            extension: Some("*.warp".to_string()),
+            default_name: Some(default_name),
+            default: Some(default_file_path.to_string_lossy().to_string()),
+            value: None,
+        }
+    }
 
-impl Command for CreateSignatureFile {
-    fn action(&self, view: &BinaryView) {
-        let is_function_named = |f: &Guard<Function>| {
-            !f.symbol().short_name().to_string_lossy().contains("sub_") || f.has_user_annotations()
-        };
-        let mut signature_dir = user_signature_dir();
-        if let Some(default_plat) = view.default_platform() {
-            // If there is a default platform, put the signature in there.
-            // TODO: We should instead use the platform of the function.
-            signature_dir.push(default_plat.name().to_string());
+    pub fn from_form(form: &Form) -> Option<PathBuf> {
+        let field = form.get_field_with_name("File Path")?;
+        let field_value = field.try_value_string()?;
+        Some(PathBuf::from(field_value))
+    }
+}
+
+pub struct OpenFileField;
+
+impl OpenFileField {
+    pub fn field() -> FormInputField {
+        FormInputField::OpenFileName {
+            prompt: "Input File Path".to_string(),
+            extension: None,
+            default: None,
+            value: None,
         }
-        let view = view.to_owned();
-        thread::spawn(move || {
-            let total_functions = view.functions().len();
-            let done_functions = AtomicUsize::default();
-            let background_task = binaryninja::background_task::BackgroundTask::new(
-                &format!("Generating signatures... ({}/{})", 0, total_functions),
-                true,
-            );
+    }
+
+    pub fn from_form(form: &Form) -> Option<PathBuf> {
+        let field = form.get_field_with_name("Input File Path")?;
+        let field_value = field.try_value_string()?;
+        Some(PathBuf::from(field_value))
+    }
+}
+
+pub struct CreateFromCurrentView;
 
-            let start = Instant::now();
+impl CreateFromCurrentView {
+    pub fn execute(view: Ref<BinaryView>, external_file: bool) -> Option<()> {
+        // Prompt the user first so that they can go do other things and not worry about a popup.
+        let mut form = Form::new("Create From View");
 
-            let mut data = warp::signature::Data::default();
-            data.functions.par_extend(
-                view.functions()
-                    .par_iter()
-                    .inspect(|_| {
-                        done_functions.fetch_add(1, Relaxed);
-                        background_task.set_progress_text(&format!(
-                            "Generating signatures... ({}/{})",
-                            done_functions.load(Relaxed),
-                            total_functions
-                        ))
-                    })
-                    .filter(is_function_named)
-                    .filter(|f| !f.analysis_skipped())
-                    .filter_map(|func| {
-                        let llil = func.low_level_il().ok()?;
-                        Some(cached_function(&func, &llil))
-                    }),
+        if external_file {
+            form.add_field(OpenFileField::field());
+        }
+
+        form.add_field(SaveFileField::field(&view));
+
+        let fd_field = FileDataKindField::default();
+        form.add_field(fd_field.to_field());
+
+        let compression_field = CompressionTypeField::default();
+        form.add_field(compression_field.to_field());
+
+        let mut included_field = IncludedFunctionsField::default();
+        // If the view has the include tag, we better set the default to the selected functions.
+        if view.tag_type_by_name(INCLUDE_TAG_NAME).is_some() {
+            included_field = IncludedFunctionsField::Selected;
+        }
+        form.add_field(included_field.to_field());
+
+        let report_field = ReportKindField::default();
+        form.add_field(report_field.to_field());
+        let report_to_disk_field = SaveReportToDiskField::default();
+        form.add_field(report_to_disk_field.to_field());
+
+        if !form.prompt() {
+            return None;
+        }
+        let compression_type = CompressionTypeField::from_form(&form).unwrap_or_default();
+        let file_path = SaveFileField::from_form(&form)?;
+        let file_data_kind = FileDataKindField::from_form(&form).unwrap_or_default();
+        let file_included_functions = IncludedFunctionsField::from_form(&form).unwrap_or_default();
+        let report_kind = ReportKindField::from_form(&form).unwrap_or_default();
+        let save_report_to_disk = SaveReportToDiskField::from_form(&form).unwrap_or_default();
+        let open_file_path = OpenFileField::from_form(&form);
+
+        // If we already have a file, prompt the user if they want to add the data.
+        let mut existing_chunks = Vec::new();
+        if file_path.exists() {
+            let prompt_result = binaryninja::interaction::show_message_box(
+                "Keep existing file data?",
+                "The file already exists. Do you want to keep the existing data?",
+                MessageBoxButtonSet::YesNoCancelButtonSet,
+                MessageBoxIcon::QuestionIcon,
             );
 
-            if let Some(ref_ty_cache) = cached_type_references(&view) {
-                let referenced_types = ref_ty_cache
-                    .cache
-                    .iter()
-                    .filter_map(|t| t.to_owned())
-                    .collect::<Vec<_>>();
+            match prompt_result {
+                MessageBoxButtonResult::NoButton => {
+                    // User wants to overwrite the file.
+                }
+                MessageBoxButtonResult::YesButton | MessageBoxButtonResult::OKButton => {
+                    // User wants to keep the existing data.
+                    let data = std::fs::read(&file_path).ok()?;
+                    let existing_file = WarpFile::from_owned_bytes(data)?;
+                    existing_chunks.extend(existing_file.chunks);
+                }
+                MessageBoxButtonResult::CancelButton => {
+                    log::info!(
+                        "User cancelled signature file creation, no operations were performed."
+                    );
+                    return None;
+                }
+            }
+        }
+
+        let processor = WarpFileProcessor::new()
+            .with_compression_type(compression_type)
+            .with_file_data(file_data_kind)
+            .with_included_functions(file_included_functions);
 
-                data.types.extend(referenced_types);
+        let file = match open_file_path {
+            None => {
+                // We are processing the current view. NOT an external file.
+                // Reference path is just used for the state tracking. Does not need to be readable.
+                let reference_path = file_path.clone();
+                processor.process_view(reference_path, &view)
+            }
+            Some(open_file_path) => {
+                // 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 file = processor.process(open_file_path);
+                background_task.finish();
+                file
             }
+        };
+
+        if let Err(err) = file {
+            binaryninja::interaction::show_message_box(
+                "Error",
+                &format!("Failed to create signature file: {}", err),
+                MessageBoxButtonSet::OKButtonSet,
+                MessageBoxIcon::ErrorIcon,
+            );
+            log::error!("Failed to create signature file: {}", err);
+            return None;
+        }
 
-            log::info!("Signature generation took {:?}", start.elapsed());
-            background_task.finish();
+        let mut file = file.unwrap();
+        // Add back the existing chunks if the user selected to keep them.
+        file.chunks.extend(existing_chunks);
+        // TODO: Make merging optional?
+        file.chunks = Chunk::merge(&file.chunks, compression_type.into());
+
+        if std::fs::write(&file_path, file.to_bytes()).is_err() {
+            log::error!("Failed to write data to signature file!");
+        }
 
-            // NOTE: Because we only can consume signatures from a specific directory, we don't need to use the interaction API.
-            // If we did need to save signature files to a project than this would need to change.
-            let Some(save_file) = rfd::FileDialog::new()
-                .add_filter("Signature Files", &["sbin"])
-                .set_file_name(format!("{}.sbin", view.file().filename()))
-                .set_directory(signature_dir)
-                .save_file()
-            else {
-                return;
-            };
+        // Show a report of the generate signatures, if desired.
+        let report_generator = ReportGenerator::new();
+        if let Some(report_string) = report_generator.report(&report_kind, &file) {
+            if save_report_to_disk == SaveReportToDiskField::Yes {
+                let report_ext = report_generator
+                    .report_extension(&report_kind)
+                    .unwrap_or_default();
+                let report_path = file_path.with_extension(report_ext);
+                let _ = std::fs::write(report_path, &report_string);
+            }
 
-            match std::fs::write(&save_file, data.to_bytes()) {
-                Ok(_) => {
-                    log::info!("Signature file saved successfully.");
-                    // Force rebuild platform matcher.
-                    invalidate_function_matcher_cache();
+            match report_kind {
+                ReportKindField::None => {}
+                ReportKindField::Html => {
+                    view.show_html_report("Generated WARP File", report_string.as_str(), "");
+                }
+                ReportKindField::Markdown => {
+                    view.show_markdown_report("Generated WARP File", report_string.as_str(), "");
+                }
+                ReportKindField::Json => {
+                    view.show_plaintext_report("Generated WARP File", report_string.as_str());
                 }
-                Err(e) => log::error!("Failed to write data to signature file: {:?}", e),
             }
+        }
+
+        Some(())
+    }
+}
+
+impl Command for CreateFromCurrentView {
+    fn action(&self, view: &BinaryView) {
+        let view = view.to_owned();
+        thread::spawn(move || {
+            CreateFromCurrentView::execute(view, false);
+        });
+    }
+
+    fn valid(&self, _view: &BinaryView) -> bool {
+        true
+    }
+}
+
+pub struct CreateFromFiles;
+
+impl Command for CreateFromFiles {
+    fn action(&self, view: &BinaryView) {
+        let view = view.to_owned();
+        thread::spawn(move || {
+            CreateFromCurrentView::execute(view, true);
         });
     }