[BNTL] Misc improvements

6 files changed, 150 insertions, 118 deletions

diff --git a/plugins/bntl_utils/src/command/diff.rs b/plugins/bntl_utils/src/command/diff.rs
index ae93dafe..48501433 100644
--- a/plugins/bntl_utils/src/command/diff.rs
+++ b/plugins/bntl_utils/src/command/diff.rs
@@ -1,10 +1,10 @@
 use crate::command::OutputDirectoryField;
 use crate::diff::TILDiff;
+use crate::helper::path_to_type_libraries;
 use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::BinaryView;
 use binaryninja::command::Command;
 use binaryninja::interaction::{Form, FormInputField};
-use binaryninja::types::TypeLibrary;
 use std::path::PathBuf;
 use std::thread;
 
@@ -12,17 +12,15 @@ pub struct InputFileAField;
 
 impl InputFileAField {
     pub fn field() -> FormInputField {
-        FormInputField::OpenFileName {
-            prompt: "Library A".to_string(),
-            // TODO: This is called extension but is really a filter.
-            extension: Some("*.bntl".to_string()),
+        FormInputField::DirectoryName {
+            prompt: "Directory A".to_string(),
             default: None,
             value: None,
         }
     }
 
     pub fn from_form(form: &Form) -> Option<PathBuf> {
-        let field = form.get_field_with_name("Library A")?;
+        let field = form.get_field_with_name("Directory A")?;
         let field_value = field.try_value_string()?;
         Some(PathBuf::from(field_value))
     }
@@ -32,17 +30,15 @@ pub struct InputFileBField;
 
 impl InputFileBField {
     pub fn field() -> FormInputField {
-        FormInputField::OpenFileName {
-            prompt: "Library B".to_string(),
-            // TODO: This is called extension but is really a filter.
-            extension: Some("*.bntl".to_string()),
+        FormInputField::DirectoryName {
+            prompt: "Directory B".to_string(),
             default: None,
             value: None,
         }
     }
 
     pub fn from_form(form: &Form) -> Option<PathBuf> {
-        let field = form.get_field_with_name("Library B")?;
+        let field = form.get_field_with_name("Directory B")?;
         let field_value = field.try_value_string()?;
         Some(PathBuf::from(field_value))
     }
@@ -63,30 +59,39 @@ impl Diff {
         let b_path = InputFileBField::from_form(&form).unwrap();
         let output_dir = OutputDirectoryField::from_form(&form).unwrap();
 
-        let _bg_task = BackgroundTask::new("Diffing type libraries...", false).enter();
-        let Some(type_lib_a) = TypeLibrary::load_from_file(&a_path) else {
-            tracing::error!("Failed to load type library: {}", a_path.display());
-            return;
-        };
-        let Some(type_lib_b) = TypeLibrary::load_from_file(&b_path) else {
-            tracing::error!("Failed to load type library: {}", b_path.display());
-            return;
-        };
+        let bg_task = BackgroundTask::new("Diffing type libraries...", true).enter();
+
+        let b_libraries = path_to_type_libraries(&a_path);
+        let a_libraries = path_to_type_libraries(&b_path);
+        // TODO: Make this parallel
+        for a_lib in &a_libraries {
+            for b_lib in &b_libraries {
+                if bg_task.is_cancelled() {
+                    return;
+                }
 
-        let diff_result = match TILDiff::new().diff((&a_path, &type_lib_a), (&b_path, &type_lib_b))
-        {
-            Ok(diff_result) => diff_result,
-            Err(err) => {
-                tracing::error!("Failed to diff type libraries: {}", err);
-                return;
+                if a_lib.name() != b_lib.name() {
+                    continue;
+                }
+
+                bg_task.set_progress_text(&format!("Diffing '{}'...", a_lib.name()));
+                let diff_result = match TILDiff::new().diff_with_dependencies(
+                    (&a_lib, a_libraries.clone()),
+                    (&b_lib, b_libraries.clone()),
+                ) {
+                    Ok(diff_result) => diff_result,
+                    Err(err) => {
+                        tracing::error!("Failed to diff type libraries: {}", err);
+                        continue;
+                    }
+                };
+                tracing::info!("Similarity Ratio: {}", diff_result.ratio);
+
+                let output_path = output_dir.join(a_lib.name()).with_extension("diff");
+                std::fs::write(&output_path, diff_result.diff).unwrap();
+                tracing::info!("Diff written to: {}", output_path.display());
             }
-        };
-        tracing::info!("Similarity Ratio: {}", diff_result.ratio);
-        let output_path = output_dir
-            .join(type_lib_a.dependency_name())
-            .with_extension("diff");
-        std::fs::write(&output_path, diff_result.diff).unwrap();
-        tracing::info!("Diff written to: {}", output_path.display());
+        }
     }
 }
 
diff --git a/plugins/bntl_utils/src/command/dump.rs b/plugins/bntl_utils/src/command/dump.rs
index 4b68cd79..96c3b264 100644
--- a/plugins/bntl_utils/src/command/dump.rs
+++ b/plugins/bntl_utils/src/command/dump.rs
@@ -1,49 +1,59 @@
-use crate::command::{InputFileField, OutputDirectoryField};
+use crate::command::{InputDirectoryField, OutputDirectoryField};
 use crate::dump::TILDump;
 use crate::helper::path_to_type_libraries;
+use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::BinaryView;
 use binaryninja::command::Command;
 use binaryninja::interaction::Form;
-use binaryninja::types::TypeLibrary;
 
 pub struct Dump;
 
-impl Command for Dump {
-    // TODO: We need a command type that does not require a binary view.
-    fn action(&self, _view: &BinaryView) {
+impl Dump {
+    pub fn execute() {
         let mut form = Form::new("Dump to C Header");
         // TODO: The choice to select what to include?
-        form.add_field(InputFileField::field());
+        form.add_field(InputDirectoryField::field());
         form.add_field(OutputDirectoryField::field());
         if !form.prompt() {
             return;
         }
         let output_dir = OutputDirectoryField::from_form(&form).unwrap();
-        let input_path = InputFileField::from_form(&form).unwrap();
+        let input_path = InputDirectoryField::from_form(&form).unwrap();
 
-        let type_lib = match TypeLibrary::load_from_file(&input_path) {
-            Some(type_lib) => type_lib,
-            None => {
-                tracing::error!("Failed to load type library from {}", input_path.display());
-                return;
-            }
-        };
+        let bg_task = BackgroundTask::new("Dumping type libraries...", true).enter();
 
-        // TODO: Currently we collect input path dependencies from the platform and the parent directory.
-        let dependencies = path_to_type_libraries(input_path.parent().unwrap());
-        let dump = match TILDump::new().with_type_libs(dependencies).dump(&type_lib) {
-            Ok(dump) => dump,
-            Err(err) => {
-                tracing::error!("Failed to dump type library: {}", err);
+        let type_libraries = path_to_type_libraries(&input_path);
+        for type_lib in &type_libraries {
+            if bg_task.is_cancelled() {
                 return;
             }
-        };
+            bg_task.set_progress_text(&format!("Dumping '{}'...", type_lib.name()));
+            let dump = match TILDump::new()
+                .with_type_libs(type_libraries.clone())
+                .dump(&type_lib)
+            {
+                Ok(dump) => dump,
+                Err(err) => {
+                    tracing::error!("Failed to dump type library: {}", err);
+                    return;
+                }
+            };
 
-        let output_path = output_dir.join(format!("{}.h", type_lib.name()));
-        if let Err(e) = std::fs::write(&output_path, dump) {
-            tracing::error!("Failed to write dump to {}: {}", output_path.display(), e);
+            let output_path = output_dir.join(format!("{}.h", type_lib.name()));
+            if let Err(e) = std::fs::write(&output_path, dump) {
+                tracing::error!("Failed to write dump to {}: {}", output_path.display(), e);
+            }
+            tracing::info!("Dump written to {}", output_path.display());
         }
-        tracing::info!("Dump written to {}", output_path.display());
+    }
+}
+
+impl Command for Dump {
+    // TODO: We need a command type that does not require a binary view.
+    fn action(&self, _view: &BinaryView) {
+        std::thread::spawn(move || {
+            Dump::execute();
+        });
     }
 
     fn valid(&self, _view: &BinaryView) -> bool {
diff --git a/plugins/bntl_utils/src/command/validate.rs b/plugins/bntl_utils/src/command/validate.rs
index 8f0095ae..a0507974 100644
--- a/plugins/bntl_utils/src/command/validate.rs
+++ b/plugins/bntl_utils/src/command/validate.rs
@@ -1,76 +1,79 @@
+use crate::command::{InputDirectoryField, OutputDirectoryField};
 use crate::helper::path_to_type_libraries;
 use crate::validate::TypeLibValidater;
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::command::Command;
-use binaryninja::interaction::get_open_filename_input;
+use binaryninja::interaction::Form;
 use binaryninja::platform::Platform;
-use binaryninja::types::TypeLibrary;
 
 pub struct Validate;
 
-impl Command for Validate {
-    fn action(&self, _view: &BinaryView) {
-        let Some(input_path) =
-            get_open_filename_input("Select a type library to validate", "*.bntl")
-        else {
-            return;
-        };
-
-        let type_lib = match TypeLibrary::load_from_file(&input_path) {
-            Some(type_lib) => type_lib,
-            None => {
-                tracing::error!("Failed to load type library from {}", input_path.display());
-                return;
-            }
-        };
-
-        // Type libraries should always have at least one platform associated with them.
-        if type_lib.platform_names().is_empty() {
-            tracing::error!("Type library {} has no platforms!", input_path.display());
+impl Validate {
+    pub fn execute() {
+        let mut form = Form::new("Validate Type Libraries");
+        form.add_field(InputDirectoryField::field());
+        form.add_field(OutputDirectoryField::field());
+        if !form.prompt() {
             return;
         }
+        let output_dir = OutputDirectoryField::from_form(&form).unwrap();
+        let input_path = InputDirectoryField::from_form(&form).unwrap();
 
-        // TODO: Currently we collect input path dependencies from the platform and the parent directory.
-        let dependencies = path_to_type_libraries(input_path.parent().unwrap());
-
-        let validator = TypeLibValidater::new().with_type_libraries(dependencies);
-        // Validate for every platform so that we can find issues in lesser used platforms.
-        for platform_name in &type_lib.platform_names() {
-            let Some(platform) = Platform::by_name(platform_name) else {
-                tracing::error!("Failed to find platform with name {}", platform_name);
-                continue;
-            };
-            let results = validator
-                .clone()
-                .with_platform(&platform)
-                .validate(&type_lib);
-            if results.issues.is_empty() {
-                tracing::info!(
-                    "No issues found for type library {} on platform {}",
-                    type_lib.name(),
-                    platform_name
-                );
+        let type_libraries = path_to_type_libraries(&input_path);
+        // TODO: Run this in parallel.
+        for type_lib in &type_libraries {
+            // Type libraries should always have at least one platform associated with them.
+            if type_lib.platform_names().is_empty() {
+                tracing::error!("Type library {} has no platforms!", input_path.display());
                 continue;
             }
-            let rendered = match results.render_report() {
-                Ok(rendered) => rendered,
-                Err(err) => {
-                    tracing::error!("Failed to render validation report: {}", err);
+
+            // TODO: Currently we collect input path dependencies from the platform and the parent directory.
+            let validator = TypeLibValidater::new().with_type_libraries(type_libraries.clone());
+            // Validate for every platform so that we can find issues in lesser used platforms.
+            for platform_name in &type_lib.platform_names() {
+                let Some(platform) = Platform::by_name(platform_name) else {
+                    tracing::error!("Failed to find platform with name {}", platform_name);
+                    continue;
+                };
+                let results = validator
+                    .clone()
+                    .with_platform(&platform)
+                    .validate(&type_lib);
+                if results.issues.is_empty() {
+                    tracing::info!(
+                        "No issues found for type library {} on platform {}",
+                        type_lib.name(),
+                        platform_name
+                    );
                     continue;
                 }
-            };
-            let out_path = input_path.with_extension(format!("{}.html", platform_name));
-            let out_name = format!("{} ({})", type_lib.name(), platform_name);
-            _view.show_html_report(&out_name, &rendered, "");
-            if let Err(e) = std::fs::write(out_path, rendered) {
-                tracing::error!(
-                    "Failed to write validation report to {}: {}",
-                    input_path.display(),
-                    e
-                );
+                let rendered = match results.render_report() {
+                    Ok(rendered) => rendered,
+                    Err(err) => {
+                        tracing::error!("Failed to render validation report: {}", err);
+                        continue;
+                    }
+                };
+                let out_path = output_dir.with_extension(format!("{}.html", platform_name));
+                if let Err(e) = std::fs::write(out_path, rendered) {
+                    tracing::error!(
+                        "Failed to write validation report to {}: {}",
+                        output_dir.display(),
+                        e
+                    );
+                }
             }
         }
     }
+}
+
+impl Command for Validate {
+    fn action(&self, _view: &BinaryView) {
+        std::thread::spawn(move || {
+            Validate::execute();
+        });
+    }
 
     fn valid(&self, _view: &BinaryView) -> bool {
         true
diff --git a/plugins/bntl_utils/src/diff.rs b/plugins/bntl_utils/src/diff.rs
index 18d4dce8..f55e8446 100644
--- a/plugins/bntl_utils/src/diff.rs
+++ b/plugins/bntl_utils/src/diff.rs
@@ -1,5 +1,6 @@
 use crate::dump::TILDump;
 use crate::helper::path_to_type_libraries;
+use binaryninja::rc::Ref;
 use binaryninja::types::TypeLibrary;
 use similar::{Algorithm, TextDiff};
 use std::path::{Path, PathBuf};
@@ -48,9 +49,17 @@ impl TILDiff {
             .parent()
             .ok_or_else(|| TILDiffError::InvalidPath(b_path.to_path_buf()))?;
 
-        let a_dependencies = path_to_type_libraries(a_parent);
-        let b_dependencies = path_to_type_libraries(b_parent);
+        self.diff_with_dependencies(
+            (&a_type_lib, path_to_type_libraries(a_parent)),
+            (&b_type_lib, path_to_type_libraries(b_parent)),
+        )
+    }
 
+    pub fn diff_with_dependencies(
+        &self,
+        (a_type_lib, a_dependencies): (&TypeLibrary, Vec<Ref<TypeLibrary>>),
+        (b_type_lib, b_dependencies): (&TypeLibrary, Vec<Ref<TypeLibrary>>),
+    ) -> Result<DiffResult, TILDiffError> {
         let dumped_a = TILDump::new()
             .with_type_libs(a_dependencies)
             .dump(a_type_lib)
@@ -70,8 +79,8 @@ impl TILDiff {
             .unified_diff()
             .context_radius(3)
             .header(
-                a_path.to_string_lossy().as_ref(),
-                b_path.to_string_lossy().as_ref(),
+                &format!("A/{}", a_type_lib.name()),
+                &format!("B/{}", b_type_lib.name()),
             )
             .to_string();
 
diff --git a/plugins/bntl_utils/src/process.rs b/plugins/bntl_utils/src/process.rs
index b287a1d4..b81986af 100644
--- a/plugins/bntl_utils/src/process.rs
+++ b/plugins/bntl_utils/src/process.rs
@@ -143,6 +143,11 @@ impl ProcessedData {
         }
     }
 
+    /// Finalizes the processed data, deduplicating types and pruning empty type libraries.
+    ///
+    /// The `default_name` should be the library name for which you want deduplicated types to be
+    /// relocated to. This does not need to be a logical-shared library name like `mylib.dll` as it will
+    /// be only referenced by other loaded type libraries (it cannot contain named objects).
     pub fn finalized(mut self, default_name: &str) -> Self {
         self.deduplicate_types(default_name);
         // TODO: Run remap.
diff --git a/plugins/bntl_utils/src/validate.rs b/plugins/bntl_utils/src/validate.rs
index ecd0978d..b7bd4d38 100644
--- a/plugins/bntl_utils/src/validate.rs
+++ b/plugins/bntl_utils/src/validate.rs
@@ -198,7 +198,7 @@ impl TypeLibValidater {
             .extend(self.validate_external_references(type_lib));
 
         // TODO: This is currently disabled because it's too slow.
-        // result.issues.extend(self.validate_source_files(type_lib));
+        result.issues.extend(self.validate_source_files(type_lib));
 
         result
     }