IDB Import refactor

1 files changed, 471 insertions, 0 deletions

diff --git a/plugins/idb_import/src/parse.rs b/plugins/idb_import/src/parse.rs
new file mode 100644
index 00000000..72ac1baa
--- /dev/null
+++ b/plugins/idb_import/src/parse.rs
@@ -0,0 +1,471 @@
+//! Parse the provided IDB / TIL file and extract information into a struct for further processing.
+
+use idb_rs::addr_info::AddressInfo;
+use idb_rs::id0::function::{FuncIdx, FuncordsIdx, IDBFunctionType};
+use idb_rs::id0::{ID0Section, Netdelta, SegmentType};
+use idb_rs::id1::ID1Section;
+use idb_rs::id2::ID2Section;
+use idb_rs::til::section::TILSection;
+use idb_rs::til::TILTypeInfo;
+use idb_rs::{identify_idb_file, Address, IDAKind, IDAUsize, IDAVariants, IDBFormat, IDBFormats};
+use serde::Serialize;
+use std::ffi::CString;
+use std::io::{BufRead, Seek};
+use std::ops::Range;
+
+#[derive(Debug, Clone, Serialize)]
+pub struct SegmentInfo {
+    pub name: String,
+    pub region: Range<u64>,
+    pub ty: SegmentType,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct FunctionInfo {
+    pub name: Option<String>,
+    pub ty: Option<idb_rs::til::Type>,
+    pub address: u64,
+    pub is_library: bool,
+    pub is_no_return: bool,
+    pub is_entry: bool,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct NameInfo {
+    pub address: u64,
+    pub ty: Option<idb_rs::til::Type>,
+    pub label: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct CommentInfo {
+    pub address: u64,
+    pub comment: String,
+    pub is_repeatable: bool,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct LabelInfo {
+    pub address: u64,
+    pub label: String,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct FunctionCordInfo {
+    comments: Vec<CommentInfo>,
+    labels: Vec<LabelInfo>,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct ID0Info {
+    pub base_address: Option<u64>,
+    pub segments: Vec<SegmentInfo>,
+    pub functions: Vec<FunctionInfo>,
+    pub comments: Vec<CommentInfo>,
+    pub labels: Vec<LabelInfo>,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct DirTreeInfo {
+    pub functions: Vec<FunctionInfo>,
+    pub types: Vec<TILTypeInfo>,
+    /// Contains both function and data names (along with their types).
+    pub names: Vec<NameInfo>,
+}
+
+#[derive(Debug, Clone, Serialize, Default)]
+pub struct IDBInfo {
+    pub sha256: Option<String>,
+    pub id0: Option<ID0Info>,
+    // NOTE: TILSection is self-contained, so we do no pre-processing.
+    pub til: Option<TILSection>,
+    pub dir_tree: Option<DirTreeInfo>,
+}
+
+impl IDBInfo {
+    /// Retrieve the functions from both the `id0` and `dir_tree` sections, with function information
+    /// merged; this is the "sanitized" version of the functions contained in an IDB.
+    pub fn merged_functions(&self) -> Vec<FunctionInfo> {
+        let mut id0_functions = self
+            .id0
+            .as_ref()
+            .map(|id0| id0.functions.clone())
+            .unwrap_or_default();
+        let dir_tree_functions = self
+            .dir_tree
+            .as_ref()
+            .map(|dir_tree| dir_tree.functions.clone())
+            .unwrap_or_default();
+        id0_functions.extend(dir_tree_functions);
+        id0_functions.sort_by_key(|f| f.address);
+        id0_functions.dedup_by(|a, b| {
+            if a.address != b.address {
+                return false;
+            }
+            // We need to move data from one function to another, if a name is missing
+            // in one of the functions, add it to the function we are keeping, if we are missing
+            // a type, do the same.
+            if a.name.is_some() {
+                b.name = a.name.clone();
+            }
+            if a.ty.is_some() {
+                b.ty = a.ty.clone();
+            }
+            true
+        });
+        id0_functions
+    }
+
+    pub fn merged_names(&self) -> Vec<NameInfo> {
+        let mut names = Vec::new();
+        if let Some(id0) = &self.id0 {
+            for label in &id0.labels {
+                names.push(NameInfo {
+                    address: label.address,
+                    ty: None,
+                    label: Some(label.label.clone()),
+                });
+            }
+
+            for func in &id0.functions {
+                names.push(NameInfo {
+                    address: func.address,
+                    ty: func.ty.clone(),
+                    label: func.name.clone(),
+                });
+            }
+        }
+        if let Some(dir_tree) = &self.dir_tree {
+            names.extend(dir_tree.names.clone());
+        }
+        names.sort_by_key(|n| n.address);
+        names.dedup_by(|a, b| {
+            if a.address != b.address {
+                return false;
+            }
+            if a.label.is_some() {
+                b.label = a.label.clone();
+            }
+            if a.ty.is_some() {
+                b.ty = a.ty.clone();
+            }
+            true
+        });
+        names
+    }
+
+    pub fn merged_types(&self) -> Vec<TILTypeInfo> {
+        let mut types = Vec::new();
+        if let Some(dir_tree) = &self.dir_tree {
+            types.extend(dir_tree.types.clone());
+        }
+        if let Some(til) = &self.til {
+            types.extend(til.types.clone());
+        }
+        types.sort_by_key(|t| t.name.to_string());
+        types.dedup_by(|a, b| {
+            if a.name.to_string() != b.name.to_string() {
+                return false;
+            }
+            // TODO: Merge types instead of just picking b and not transferring.
+            true
+        });
+        types
+    }
+}
+
+/// Parsed the IDB data into [`IDBInfo`].
+pub struct IDBFileParser;
+
+impl IDBFileParser {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    pub fn parse<I: BufRead + Seek>(&self, data: &mut I) -> anyhow::Result<IDBInfo> {
+        match identify_idb_file(data)? {
+            IDBFormats::Separated(f) => match f {
+                IDAVariants::IDA32(f_32) => self.parse_format(&mut *data, f_32),
+                IDAVariants::IDA64(f_64) => self.parse_format(&mut *data, f_64),
+            },
+            IDBFormats::InlineUncompressed(f) => self.parse_format(&mut *data, f),
+            IDBFormats::InlineCompressed(f) => {
+                let mut decompressed = Vec::new();
+                let uncompressed_format = f.decompress_into_memory(data, &mut decompressed)?;
+                let mut decompressed_data = std::io::Cursor::new(decompressed);
+                self.parse_format(&mut decompressed_data, uncompressed_format)
+            }
+        }
+    }
+
+    pub fn parse_format<I: BufRead + Seek, K: IDAKind>(
+        &self,
+        data: &mut I,
+        format: impl IDBFormat<K>,
+    ) -> anyhow::Result<IDBInfo> {
+        let mut id0 = None;
+        if let Some(id0_loc) = format.id0_location() {
+            id0 = Some(format.read_id0(&mut *data, id0_loc)?);
+        }
+
+        let mut id1 = None;
+        if let Some(id1_loc) = format.id1_location() {
+            id1 = Some(format.read_id1(&mut *data, id1_loc)?);
+        }
+
+        let mut id2 = None;
+        if let Some(id2_loc) = format.id2_location() {
+            id2 = Some(format.read_id2(&mut *data, id2_loc)?);
+        }
+
+        // TODO: Decompress til
+        let mut til = None;
+        if let Some(til_loc) = format.til_location() {
+            til = Some(format.read_til(&mut *data, til_loc)?);
+        };
+
+        let dir_tree_info = match (id0.as_ref(), id1.as_ref(), id2.as_ref(), til.as_ref()) {
+            (Some(id0), Some(id1), id2, til) => Some(self.parse_dir_tree(id0, id1, id2, til)?),
+            _ => None,
+        };
+
+        let id0_info = id0.as_ref().map(|id0| self.parse_id0(id0)).transpose()?;
+
+        Ok(IDBInfo {
+            sha256: None,
+            id0: id0_info,
+            til,
+            dir_tree: dir_tree_info,
+        })
+    }
+
+    pub fn parse_id0<K: IDAKind>(&self, id0: &ID0Section<K>) -> anyhow::Result<ID0Info> {
+        let root_info_idx = id0.root_node()?;
+        let root_info = id0.ida_info(root_info_idx)?;
+        let netdelta = root_info.netdelta();
+
+        let mut segments = Vec::new();
+        if let Some(seg_idx) = id0.segments_idx()? {
+            for entry in id0.segments(seg_idx) {
+                let Ok(segment) = entry else {
+                    tracing::warn!("Failed to read segment entry");
+                    continue;
+                };
+                let name = id0.segment_name(segment.name)?.map(|s| s.to_string());
+                let seg_start = segment.address.start.into_raw().into_u64();
+                let seg_end = segment.address.end.into_raw().into_u64();
+                segments.push(SegmentInfo {
+                    name: name.unwrap_or_else(|| format!("seg_{:0x}", seg_start)),
+                    region: seg_start..seg_end,
+                    ty: segment.seg_type,
+                });
+            }
+        }
+
+        let mut functions = Vec::new();
+        let mut comments = Vec::new();
+        let mut labels = Vec::new();
+        if let Some(funcs_idx) = id0.funcs_idx()? {
+            if let Some(funcords_idx) = id0.funcords_idx()? {
+                let info = self.parse_func_cord(&id0, netdelta, funcords_idx, funcs_idx)?;
+                comments.extend(info.comments);
+                labels.extend(info.labels);
+            }
+
+            for entry in id0.fchunks(funcs_idx) {
+                let Ok(func) = entry else {
+                    tracing::warn!("Failed to read function entry");
+                    continue;
+                };
+
+                let func_start = func.address.start.into_raw().into_u64();
+                match &func.extra {
+                    IDBFunctionType::Tail(_) => {
+                        tracing::debug!("Skipping tail function... {:0x}", func_start);
+                    }
+                    IDBFunctionType::NonTail(_func_ext) => {
+                        if func.flags.is_outline() {
+                            tracing::debug!("Skipping outlined function... {:0x}", func_start);
+                            continue;
+                        }
+
+                        // TODO: Parse function registers and params
+                        // for def_reg in id0.function_defined_registers(netdelta, &func, &func_ext) {
+                        //     tracing::info!("{:0x} : Function register: {:?}", func_start, def_reg);
+                        //     let Ok(_def_reg) = def_reg else {
+                        //         tracing::warn!("Failed to read function register entry");
+                        //         continue;
+                        //     };
+                        // }
+
+                        // if let Ok(stack_names) =
+                        //     id0.function_defined_variables(&root_info, &func, &func_ext)
+                        // {
+                        //     tracing::info!(
+                        //         "{:0x} : Function stack variables: {:#?}",
+                        //         func_start,
+                        //         stack_names
+                        //     );
+                        // }
+
+                        functions.push(FunctionInfo {
+                            name: None,
+                            ty: None,
+                            address: func_start,
+                            is_library: func.flags.is_lib(),
+                            is_no_return: func.flags.is_no_return(),
+                            is_entry: false,
+                        });
+                    }
+                }
+            }
+        }
+
+        if let Ok(entry_points) = id0.entry_points(&root_info) {
+            for entry in entry_points {
+                // TODO: What to do with entry.forwarded?
+                functions.push(FunctionInfo {
+                    name: Some(entry.name),
+                    ty: entry.entry_type,
+                    address: entry.address.into_u64(),
+                    is_library: false,
+                    is_no_return: false,
+                    is_entry: true,
+                });
+            }
+        }
+
+        let base_address = match root_info.addresses.loading_base.into_u64() {
+            // An IDB with zero loading base is possibly not loaded there.
+            // For example, see the FlawedGrace.idb in the idb-rs resources directory.
+            0 => None,
+            loading_base => Some(loading_base.into_u64()),
+        };
+
+        Ok(ID0Info {
+            base_address,
+            segments,
+            functions,
+            comments,
+            labels,
+        })
+    }
+
+    pub fn parse_func_cord<K: IDAKind>(
+        &self,
+        id0: &ID0Section<K>,
+        netdelta: Netdelta<K>,
+        funcords_idx: FuncordsIdx<K>,
+        funcs_idx: FuncIdx<K>,
+    ) -> anyhow::Result<FunctionCordInfo> {
+        let mut comments = Vec::new();
+        let mut labels = Vec::new();
+
+        for entry in id0.funcords(funcords_idx)? {
+            let Ok(address) = entry else {
+                tracing::warn!("Failed to read function address entry");
+                continue;
+            };
+
+            for (label_addr, label_data) in id0.local_labels(netdelta, address)? {
+                if let Ok(label_data_cstr) = CString::new(label_data) {
+                    let label_data_str = label_data_cstr.to_string_lossy();
+                    labels.push(LabelInfo {
+                        address: label_addr.into_raw().into_u64(),
+                        label: label_data_str.to_string(),
+                    });
+                }
+            }
+
+            if let Some(comment) = id0.func_cmt(funcs_idx, netdelta, address)? {
+                comments.push(CommentInfo {
+                    address: address.into_raw().into_u64(),
+                    comment: comment.to_string(),
+                    is_repeatable: false,
+                });
+            }
+
+            if let Some(comment) = id0.func_repeatable_cmt(funcs_idx, netdelta, address)? {
+                comments.push(CommentInfo {
+                    address: address.into_raw().into_u64(),
+                    comment: comment.to_string(),
+                    is_repeatable: true,
+                });
+            }
+        }
+
+        Ok(FunctionCordInfo { comments, labels })
+    }
+
+    pub fn parse_dir_tree<K: IDAKind>(
+        &self,
+        id0: &ID0Section<K>,
+        id1: &ID1Section<K>,
+        id2: Option<&ID2Section<K>>,
+        til: Option<&TILSection>,
+    ) -> anyhow::Result<DirTreeInfo> {
+        let root_info_idx = id0.root_node()?;
+        let root_info = id0.ida_info(root_info_idx)?;
+        let netdelta = root_info.netdelta();
+
+        // sha256
+
+        let func_info_from_addr =
+            |addr_info: &AddressInfo<K>| -> anyhow::Result<Option<FunctionInfo>> {
+                let func_name = addr_info.label()?.map(|s| s.to_string());
+                let func_ty = addr_info.tinfo(&root_info)?;
+                let func_addr = addr_info.address().into_raw().into_u64();
+                Ok(Some(FunctionInfo {
+                    name: func_name,
+                    ty: func_ty,
+                    address: func_addr,
+                    is_library: false,
+                    is_no_return: false,
+                    is_entry: false,
+                }))
+            };
+
+        let mut functions = Vec::new();
+        if let Some(func_dir_tree) = id0.dirtree_function_address()? {
+            func_dir_tree.visit_leafs(|addr_raw| {
+                let addr = Address::from_raw(*addr_raw);
+                if let Some(info) = AddressInfo::new(id0, id1, id2, netdelta, addr) {
+                    if let Ok(Some(func_info)) = func_info_from_addr(&info) {
+                        functions.push(func_info);
+                    }
+                }
+            });
+        }
+
+        let mut names = Vec::new();
+        if let Some(names_dir_tree) = id0.dirtree_names()? {
+            names_dir_tree.visit_leafs(|name_raw| {
+                let addr = Address::from_raw(*name_raw);
+                if let Some(info) = AddressInfo::new(id0, id1, id2, netdelta, addr) {
+                    names.push(NameInfo {
+                        address: info.address().into_raw().into_u64(),
+                        ty: info.tinfo(&root_info).ok().flatten().map(|t| t.clone()),
+                        label: info.label().ok().flatten().map(|s| s.to_string()),
+                    });
+                }
+            });
+        }
+
+        let mut types = Vec::new();
+        if let Some(til) = til {
+            if let Some(type_dir_tree) = id0.dirtree_tinfos()? {
+                type_dir_tree.visit_leafs(|type_ord_raw| {
+                    if let Some(type_info) = til.get_ord(type_ord_raw.into_u64()) {
+                        types.push(type_info.clone());
+                    }
+                })
+            }
+        }
+
+        Ok(DirTreeInfo {
+            functions,
+            types,
+            names,
+        })
+    }
+}