IDB Import refactor

1 files changed, 350 insertions, 0 deletions

diff --git a/plugins/idb_import/src/mapper.rs b/plugins/idb_import/src/mapper.rs
new file mode 100644
index 00000000..4301ca9c
--- /dev/null
+++ b/plugins/idb_import/src/mapper.rs
@@ -0,0 +1,350 @@
+//! Map the IDB data we parsed into the [`BinaryView`].
+
+use crate::parse::{CommentInfo, FunctionInfo, IDBInfo, LabelInfo, NameInfo, SegmentInfo};
+use crate::translate::TILTranslator;
+use binaryninja::architecture::Architecture;
+use binaryninja::binary_view::{BinaryView, BinaryViewBase, BinaryViewExt};
+use binaryninja::qualified_name::QualifiedName;
+use binaryninja::rc::Ref;
+use binaryninja::section::{SectionBuilder, Semantics};
+use binaryninja::symbol::{Symbol, SymbolType};
+use binaryninja::types::Type;
+use idb_rs::id0::SegmentType;
+use std::collections::HashSet;
+
+/// Maps IDB data into a [`BinaryView`].
+///
+/// The mapper can be re-used if mapping into multiple views.
+pub struct IDBMapper {
+    info: IDBInfo,
+}
+
+impl IDBMapper {
+    pub fn new(info: IDBInfo) -> Self {
+        Self { info }
+    }
+
+    pub fn map_to_view(&self, view: &BinaryView) {
+        let Some(id0) = &self.info.id0 else {
+            tracing::warn!("No ID0 data found, skipping mapping.");
+            return;
+        };
+
+        // TODO: Actually the below comment belongs in an IDBVerifier that tries to determine if the idb
+        // TODO: Will process correctly for the given view.
+        // TODO: Have a shasum check of the file to make sure we are not mapping to bad data?
+
+        // Rebase the address from ida -> binja without this rebased views will fail to map.
+        let bn_base_address = view.start();
+        let ida_base_address = id0.base_address.unwrap_or(bn_base_address);
+        let base_address_delta = bn_base_address.wrapping_sub(ida_base_address);
+        let rebase = |addr: u64| -> u64 { addr.wrapping_add(base_address_delta) };
+
+        for segment in &id0.segments {
+            let mut rebased_segment = segment.clone();
+            rebased_segment.region.start = rebase(segment.region.start);
+            rebased_segment.region.end = rebase(segment.region.end);
+            self.map_segment_to_view(view, &rebased_segment);
+        }
+
+        // Create the type translator, adding all referencable types from both the TIL and the binary view.
+        let platform = view.default_platform().unwrap();
+        // TODO: Need to remove this, but for now keeping this since it gets referenced a lot.
+        view.define_auto_type(
+            "size_t",
+            "IDA",
+            &Type::int(platform.arch().default_integer_size(), false),
+        );
+        let til_translator =
+            TILTranslator::new_from_platform(&platform).with_type_container(&view.type_container());
+        let til_translator = match &self.info.til {
+            Some(til) => til_translator.with_til_info(&til),
+            None => til_translator,
+        };
+
+        self.map_types_to_view(view, &til_translator, &self.info.merged_types());
+
+        for func in &self.info.merged_functions() {
+            let mut rebased_func = func.clone();
+            rebased_func.address = rebase(func.address);
+            self.map_func_to_view(view, &til_translator, &rebased_func);
+        }
+
+        // TODO: The below undo and ignore is not thread safe, this means that the mapper itself
+        // TODO: should be the only thing running at the time of the mapping process.
+        let undo = view.file().begin_undo_actions(true);
+        for comment in &id0.comments {
+            let mut rebased_comment = comment.clone();
+            rebased_comment.address = rebase(comment.address);
+            self.map_comment_to_view(view, &rebased_comment);
+        }
+        view.file().forget_undo_actions(&undo);
+
+        for name in &self.info.merged_names() {
+            let mut rebased_name = name.clone();
+            rebased_name.address = rebase(name.address);
+            self.map_name_to_view(view, &til_translator, &rebased_name);
+        }
+
+        // self.map_used_types_to_view(view, &til_translator);
+    }
+
+    pub fn map_types_to_view(
+        &self,
+        view: &BinaryView,
+        til_translator: &TILTranslator,
+        types: &[idb_rs::til::TILTypeInfo],
+    ) {
+        for ty in types {
+            let ty_name = ty.name.to_string();
+            if view.type_by_name(&ty_name).is_some() {
+                tracing::debug!("Type already exists in view: {}", ty_name);
+                continue;
+            }
+            match til_translator.translate_type_info(&ty.tinfo) {
+                Ok(bn_ty) => {
+                    tracing::debug!("Mapping type: {:?}", ty);
+                    view.define_auto_type(&ty_name, "IDA", &bn_ty);
+                }
+                Err(err) => {
+                    tracing::warn!("Failed to map type {:?}: {}", ty, err)
+                }
+            }
+        }
+    }
+
+    pub fn map_used_types_to_view(&self, view: &BinaryView, til_translator: &TILTranslator) {
+        let type_archives: Vec<_> = view
+            .attached_type_archives()
+            .iter()
+            .filter_map(|id| view.type_archive_by_id(&id))
+            .collect();
+
+        let mut til_type_map = std::collections::HashMap::new();
+        if let Some(til) = &self.info.til {
+            til_type_map = til
+                .types
+                .iter()
+                .map(|ty| (ty.name.to_string(), ty))
+                .collect();
+        }
+        if let Some(dir_tree) = &self.info.dir_tree {
+            til_type_map = dir_tree
+                .types
+                .iter()
+                .map(|ty| (ty.name.to_string(), ty))
+                .collect();
+        }
+
+        let mut used_types = HashSet::new();
+        if let Ok(_used_types) = til_translator.used_types.lock() {
+            used_types = _used_types.clone();
+        }
+        // TODO: Adding types to view after the types have been applied to the functions is not a
+        // TODO: great idea, I imagine the NTR's will have stale references until the analysis runs again.
+        'found: for used_ty in &used_types {
+            // 0. Make sure the type doesn't already exist in the view
+            if view.type_by_name(&used_ty.name).is_some() {
+                tracing::debug!("Type already exists in view: {:?}", used_ty.name);
+                continue 'found;
+            }
+
+            // 1. Check in BN type libraries.
+            if let Some(found_ty) = view.import_type_library_type(&used_ty.name, None) {
+                tracing::debug!("Found type in type library: {:?}", found_ty);
+                continue 'found;
+            }
+
+            // 2. Check in type archives
+            for type_archive in &type_archives {
+                if let Some(found_ty) =
+                    type_archive.get_type_by_name(QualifiedName::from(&used_ty.name))
+                {
+                    tracing::debug!("Found type in type archive: {:?}", found_ty);
+                    view.define_auto_type(&used_ty.name, "IDA", &found_ty);
+                    continue 'found;
+                }
+            }
+
+            // // 3. Check in the TIL of the IDB info.
+            if let Some(ty) = til_type_map.get(&used_ty.name) {
+                if let Ok(bn_ty) = til_translator.translate_type_info(&ty.tinfo) {
+                    tracing::debug!("Found type in TIL: {:?}", ty);
+                    view.define_auto_type(&used_ty.name, "IDA", &bn_ty);
+                    continue 'found;
+                }
+            }
+
+            tracing::warn!("Failed to find type: {:?}", used_ty);
+            // 4. TODO: Look through the idb attached tils?
+        }
+    }
+
+    pub fn map_segment_to_view(&self, view: &BinaryView, segment: &SegmentInfo) {
+        let semantics = match segment.ty {
+            SegmentType::Norm => Semantics::DefaultSection,
+            SegmentType::Xtrn => {
+                // IDA definition of extern is an actual section like '.idata' whereas extern in BN
+                // is a synthetic section, do NOT use [`Semantics::External`].
+                Semantics::ReadWriteData
+            }
+            SegmentType::Code => Semantics::ReadOnlyCode,
+            SegmentType::Data => Semantics::ReadWriteData,
+            SegmentType::Imp => Semantics::DefaultSection,
+            SegmentType::Grp => Semantics::DefaultSection,
+            SegmentType::Null => Semantics::DefaultSection,
+            SegmentType::Undf => {
+                // Don't map undefined segment i guess?
+                return;
+            }
+            SegmentType::Bss => Semantics::ReadWriteData,
+            SegmentType::Abssym => Semantics::DefaultSection,
+            SegmentType::Comm => Semantics::DefaultSection,
+            SegmentType::Imem => Semantics::DefaultSection,
+        };
+
+        // TODO: Is this section already mapped using address range not name.
+        if view.section_by_name(&segment.name).is_some() {
+            tracing::debug!(
+                "Section with name '{}' already exists, skipping...",
+                segment.name
+            );
+            return;
+        }
+
+        tracing::info!(
+            "Mapping segment '{}': {:0x} - {:0x} ({:?})",
+            segment.name,
+            segment.region.start,
+            segment.region.end,
+            segment.ty
+        );
+
+        let section = SectionBuilder::new(segment.name.clone(), segment.region.clone())
+            .semantics(semantics)
+            .is_auto(true);
+        view.add_section(section);
+    }
+
+    pub fn map_func_to_view(
+        &self,
+        view: &BinaryView,
+        til_translator: &TILTranslator,
+        func: &FunctionInfo,
+    ) {
+        let Some(bn_func) = view.add_auto_function(func.address) else {
+            tracing::warn!("Failed to add function for {:0x}", func.address);
+            return;
+        };
+
+        if let Some(func_ty) = &func.ty {
+            match til_translator.translate_type_info(&func_ty) {
+                Ok(bn_func_ty) => {
+                    tracing::debug!("Mapping function with type: {:0x}", func.address);
+                    bn_func.apply_auto_discovered_type(&bn_func_ty);
+                }
+                Err(err) => {
+                    tracing::warn!(
+                        "Failed to translate type {:?} for function {:0x}: {}",
+                        func_ty,
+                        func.address,
+                        err
+                    );
+                }
+            }
+        }
+
+        // TODO: Attach a platform tuple to the FunctionInfo?
+        if let Some(func_sym) = symbol_from_func(func) {
+            tracing::debug!(
+                "Mapping function symbol: {:0x} => {}",
+                func.address,
+                func_sym
+            );
+            view.define_auto_symbol(&func_sym);
+        }
+    }
+
+    pub fn map_label_to_view(&self, view: &BinaryView, label: &LabelInfo) {
+        let symbol = Symbol::builder(SymbolType::LocalLabel, &label.label, label.address).create();
+        tracing::debug!("Mapping label: {:0x} => {}", label.address, symbol);
+        view.define_auto_symbol(&symbol);
+    }
+
+    pub fn map_name_to_view(
+        &self,
+        view: &BinaryView,
+        til_translator: &TILTranslator,
+        name: &NameInfo,
+    ) {
+        // Currently, we only want to use name info to map data variables, so skip anything in code.
+        let within_code_section = view
+            .sections_at(name.address)
+            .iter()
+            .find(|s| s.semantics() == Semantics::ReadOnlyCode)
+            .is_some();
+        if within_code_section || !view.functions_containing(name.address).is_empty() {
+            tracing::debug!("Skipping name contained in code: {:0x}", name.address);
+            return;
+        }
+
+        if let Some(label) = &name.label {
+            let symbol = Symbol::builder(SymbolType::Data, &label, name.address).create();
+            tracing::debug!("Mapping name label: {:0x} => {}", name.address, symbol);
+            view.define_auto_symbol(&symbol);
+        }
+
+        if let Some(data_ty) = &name.ty {
+            match til_translator.translate_type_info(&data_ty) {
+                Ok(data_ty) => {
+                    tracing::debug!("Mapping name with type: {:0x}", name.address);
+                    view.define_auto_data_var(name.address, &data_ty);
+                }
+                Err(err) => {
+                    tracing::warn!(
+                        "Failed to translate type {:?} for name {:0x}: {}",
+                        data_ty,
+                        name.address,
+                        err
+                    );
+                }
+            }
+        }
+    }
+
+    pub fn map_comment_to_view(&self, view: &BinaryView, comment: &CommentInfo) {
+        // NOTE: This (`set_comment`) will generate an undo action.
+        // First try and attach the comment to the containing functions, if that fails, then
+        // attach the comment to the view. Attaching to the containing function can help with
+        // the comments' placement.
+        let functions = view.functions_containing(comment.address);
+        for func in &functions {
+            if func.start() == comment.address {
+                func.set_comment(&comment.comment);
+            } else {
+                func.set_comment_at(comment.address, &comment.comment);
+            }
+        }
+
+        // We did not find any functions containing the comment, so attach it to the view.
+        if functions.is_empty() {
+            view.set_comment_at(comment.address, &comment.comment);
+        }
+    }
+}
+
+fn symbol_from_func(func: &FunctionInfo) -> Option<Ref<Symbol>> {
+    let Some(func_name) = &func.name else {
+        return None;
+    };
+    let short_func_name = binaryninja::demangle::demangle_llvm(&func_name, true)
+        .map(|qn| qn.to_string())
+        .unwrap_or(func_name.clone());
+    let sym_type = match func.is_library {
+        true => SymbolType::LibraryFunction,
+        false => SymbolType::Function,
+    };
+    let symbol_builder =
+        Symbol::builder(sym_type, func_name, func.address).short_name(short_func_name);
+    Some(symbol_builder.create())
+}