WARP: User friendly file picker

Instead of using the interaction api or text input just use a rust crate to manage the file picker dialog

4 files changed, 63 insertions, 57 deletions

diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index 1b4f70a0..dee5d26b 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -1,8 +1,9 @@
+use std::path::PathBuf;
 use binaryninja::architecture::{
     Architecture, ImplicitRegisterExtend, Register as BNRegister, RegisterInfo,
 };
 use binaryninja::basicblock::BasicBlock as BNBasicBlock;
-use binaryninja::binaryview::BinaryViewExt;
+use binaryninja::binaryview::{BinaryView, BinaryViewExt};
 use binaryninja::function::{Function as BNFunction, NativeBlock};
 use binaryninja::llil;
 use binaryninja::llil::{
@@ -25,6 +26,24 @@ mod matcher;
 /// Only used when compiled for cdylib target.
 mod plugin;
 
+pub fn core_signature_dir() -> PathBuf {
+    // Get core signatures for the given platform
+    let install_dir = binaryninja::install_directory().unwrap();
+    // macOS core dir is separate from the install dir.
+    #[cfg(target_os = "macos")]
+    let core_dir = install_dir
+        .parent()
+        .unwrap()
+        .join("Resources");
+    #[cfg(not(target_os = "macos"))]
+    let core_dir = install_dir;
+    core_dir.join("signatures")
+}
+
+pub fn user_signature_dir() -> PathBuf {
+    binaryninja::user_directory().unwrap().join("signatures/")
+}
+
 pub fn build_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
     func: &BNFunction,
     llil: &llil::Function<A, M, NonSSA<V>>,
diff --git a/plugins/warp/src/matcher.rs b/plugins/warp/src/matcher.rs
index 52180777..86985cc2 100644
--- a/plugins/warp/src/matcher.rs
+++ b/plugins/warp/src/matcher.rs
@@ -23,6 +23,7 @@ use crate::cache::{
     try_cached_function_guid,
 };
 use crate::convert::to_bn_type;
+use crate::{core_signature_dir, user_signature_dir};
 use crate::plugin::on_matched_function;
 
 pub static PLAT_MATCHER_CACHE: OnceLock<DashMap<PlatformID, Matcher>> = OnceLock::new();
@@ -61,24 +62,12 @@ impl Matcher {
     /// Create a matcher from the platforms signature subdirectory.
     pub fn from_platform(platform: BNRef<Platform>) -> Self {
         let platform_name = platform.name().to_string();
-        // Get core signatures for the given platform
-        let install_dir = binaryninja::install_directory().unwrap();
-        #[cfg(target_os = "macos")]
-        let root_core_sig_dir = install_dir
-            .parent()
-            .unwrap()
-            .join("Resources")
-            .join("signatures");
-        #[cfg(not(target_os = "macos"))]
-        let root_core_sig_dir = install_dir.join("signatures");
-        let plat_core_sig_dir = root_core_sig_dir.join(&platform_name);
+        
+        // Get core and user signatures.
+        // TODO: Separate each file into own bucket for filtering?
+        let plat_core_sig_dir = core_signature_dir().join(&platform_name);
         let mut data = get_data_from_dir(&plat_core_sig_dir);
-
-        // Get user signatures for the given platform
-        let user_dir = binaryninja::user_directory().unwrap();
-        let root_user_sig_dir = user_dir.join("signatures");
-        let plat_user_sig_dir = root_user_sig_dir.join(&platform_name);
-        // If the dir has not been created, create it.
+        let plat_user_sig_dir = user_signature_dir().join(&platform_name);
         let user_data = get_data_from_dir(&plat_user_sig_dir);
 
         data.extend(user_data);
diff --git a/plugins/warp/src/plugin/add.rs b/plugins/warp/src/plugin/add.rs
index ef4f9f26..8b50425b 100644
--- a/plugins/warp/src/plugin/add.rs
+++ b/plugins/warp/src/plugin/add.rs
@@ -5,11 +5,14 @@ use binaryninja::command::FunctionCommand;
 use binaryninja::function::Function;
 use std::io::Write;
 use std::thread;
+use crate::user_signature_dir;
 
 pub struct AddFunctionSignature;
 
 impl FunctionCommand for AddFunctionSignature {
     fn action(&self, view: &BinaryView, func: &Function) {
+        let func_plat_name = func.platform().name().to_string();
+        let signature_dir = user_signature_dir().join(func_plat_name);
         let view = view.to_owned();
         let func = func.to_owned();
         thread::spawn(move || {
@@ -17,12 +20,14 @@ impl FunctionCommand for AddFunctionSignature {
                 log::error!("Could not get low level IL for function.");
                 return;
             };
-
-            let Some(save_file) = binaryninja::interaction::get_save_filename_input(
-                "Use Signature File",
-                "*.sbin",
-                "user.sbin",
-            ) else {
+            
+            // 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("user.sbin")
+                .set_directory(signature_dir)
+                .save_file() else {
                 return;
             };
 
@@ -50,17 +55,13 @@ impl FunctionCommand for AddFunctionSignature {
                 data.types.extend(referenced_types);
             }
 
-            if let Ok(mut file) = std::fs::File::create(&save_file) {
-                match file.write_all(&data.to_bytes()) {
-                    Ok(_) => {
-                        log::info!("Signature file saved successfully.");
-                        // Force rebuild platform matcher.
-                        invalidate_function_matcher_cache();
-                    }
-                    Err(e) => log::error!("Failed to write data to signature file: {:?}", e),
+            match std::fs::write(&save_file, data.to_bytes()) {
+                Ok(_) => {
+                    log::info!("Signature file saved successfully.");
+                    // Force rebuild platform matcher.
+                    invalidate_function_matcher_cache();
                 }
-            } else {
-                log::error!("Could not create signature file: {:?}", save_file);
+                Err(e) => log::error!("Failed to write data to signature file: {:?}", e),
             }
         });
     }
diff --git a/plugins/warp/src/plugin/create.rs b/plugins/warp/src/plugin/create.rs
index 70abcba9..beb866d8 100644
--- a/plugins/warp/src/plugin/create.rs
+++ b/plugins/warp/src/plugin/create.rs
@@ -5,11 +5,11 @@ use binaryninja::command::Command;
 use binaryninja::function::Function;
 use binaryninja::rc::Guard;
 use rayon::prelude::*;
-use std::io::Write;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::Relaxed;
 use std::thread;
 use std::time::Instant;
+use crate::user_signature_dir;
 
 pub struct CreateSignatureFile;
 
@@ -20,10 +20,10 @@ impl Command for CreateSignatureFile {
         let is_function_named = |f: &Guard<Function>| {
             !f.symbol().short_name().as_str().contains("sub_") || f.has_user_annotations()
         };
-
-        let mut signature_dir = binaryninja::user_directory().unwrap().join("signatures/");
+        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());
         }
         let view = view.to_owned();
@@ -69,29 +69,26 @@ impl Command for CreateSignatureFile {
             }
 
             log::info!("Signature generation took {:?}", start.elapsed());
+            background_task.finish();
 
-            if let Some(sig_file_name) = binaryninja::interaction::get_text_line_input(
-                "Signature File",
-                "Create Signature File",
-            ) {
-                let save_file = signature_dir.join(sig_file_name + ".sbin");
-                log::info!("Saving to signatures to {:?}...", &save_file);
-                // TODO: Should we overwrite? Prompt user.
-                if let Ok(mut file) = std::fs::File::create(&save_file) {
-                    match file.write_all(&data.to_bytes()) {
-                        Ok(_) => {
-                            log::info!("Signature file saved successfully.");
-                            // Force rebuild platform matcher.
-                            invalidate_function_matcher_cache();
-                        }
-                        Err(e) => log::error!("Failed to write data to signature file: {:?}", e),
-                    }
-                } else {
-                    log::error!("Could not create signature file: {:?}", save_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().to_string()))
+                .set_directory(signature_dir)
+                .save_file() else {
+                return;
+            };
+            
+            match std::fs::write(&save_file, data.to_bytes()) {
+                Ok(_) => {
+                    log::info!("Signature file saved successfully.");
+                    // Force rebuild platform matcher.
+                    invalidate_function_matcher_cache();
                 }
+                Err(e) => log::error!("Failed to write data to signature file: {:?}", e),
             }
-
-            background_task.finish();
         });
     }