diff --git a/Cargo.toml b/Cargo.toml index 8c43039..7f0a364 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ edition = "2021" [dependencies] r2pipe = "0.6.0" serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0.60" +serde_json = { version = "1.0.60", features = ["raw_value"] } clap = { version = "4.0.29", features = ["derive"] } goblin = { version = "0.6.0", optional = true } walkdir = "2" @@ -19,7 +19,8 @@ tokenizers = "0.13.2" rust_search = "2.1.0" tch = {version = "0.10.3", optional = true} mimalloc = { version = "*", default-features = false } -indicatif = { version = "0.17.3", features = ["rayon"]} +indicatif = { version = "0.18.0", features = ["rayon"]} +indicatif-log-bridge = "0.2.3" itertools = "0.10.5" prettytable = "0.10.0" serde-aux = "4.2.0" @@ -32,6 +33,8 @@ hex = { version = "0.4", features = ["alloc"] } md5 = "0.7.0" sha1 = "0.10.6" sha2 = "0.10.8" +csv = "1.3" +once_cell = "1.19" [dependencies.petgraph] version = "0.6.2" diff --git a/src/afij.rs b/src/afij.rs index eb43764..0776329 100644 --- a/src/afij.rs +++ b/src/afij.rs @@ -6,9 +6,10 @@ use serde_json::Value; #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct AFIJFunctionInfo { + #[serde(alias = "addr")] pub offset: u64, pub name: String, - pub size: i128, + pub size: u64, #[serde(rename = "is-pure")] pub is_pure: String, pub realsz: u64, @@ -23,21 +24,23 @@ pub struct AFIJFunctionInfo { pub nbbs: u64, #[serde(rename = "is-lineal")] pub is_lineal: bool, - pub ninstrs: i64, - pub edges: i64, + pub ninstrs: u64, + pub edges: u64, pub ebbs: u64, pub signature: String, + #[serde(alias = "minaddr")] pub minbound: u64, - pub maxbound: i128, + #[serde(alias = "maxaddr")] + pub maxbound: u64, pub callrefs: Option>, // TODO: Need to fix this and change to string instead of i64 to get round large random numbers pub datarefs: Option>, pub codexrefs: Option>, - pub dataxrefs: Option>, - pub indegree: Option, - pub outdegree: Option, - pub nlocals: Option, - pub nargs: Option, + pub dataxrefs: Option>, // TODO: Check this data type + pub indegree: Option, + pub outdegree: Option, + pub nlocals: Option, + pub nargs: Option, pub bpvars: Option>, // Cannot find a good example of an spvars yet pub spvars: Option>, @@ -103,12 +106,12 @@ pub struct Regvar { #[serde(rename_all = "camelCase")] pub struct AFIJFeatureSubset { pub name: String, - pub ninstrs: i64, - pub edges: i64, - pub indegree: i64, - pub outdegree: i64, - pub nlocals: i64, - pub nargs: i64, + pub ninstrs: u64, + pub edges: u64, + pub indegree: u64, + pub outdegree: u64, + pub nlocals: u64, + pub nargs: u64, pub signature: String, } @@ -130,12 +133,12 @@ impl From<&AFIJFunctionInfo> for AFIJFeatureSubset { #[derive(Default, Debug, Clone, PartialEq, Hash, Serialize, Deserialize)] pub struct AFIJFeatureSubsetExtended { pub name: String, - pub ninstrs: i64, - pub edges: i64, - pub indegree: i64, - pub outdegree: i64, - pub nlocals: i64, - pub nargs: i64, + pub ninstrs: u64, + pub edges: u64, + pub indegree: u64, + pub outdegree: u64, + pub nlocals: u64, + pub nargs: u64, pub nbbs: u64, pub avg_ins_bb: OrderedFloat, } diff --git a/src/agcj.rs b/src/agcj.rs index 87097bb..8ef994b 100644 --- a/src/agcj.rs +++ b/src/agcj.rs @@ -3,7 +3,7 @@ use crate::networkx::{ CallGraphFuncNameNode, CallGraphFuncWithMetadata, CallGraphTikNibFeatures, CallGraphTikNibFinfoFeatures, NetworkxDiGraph, }; -use crate::utils::{check_or_create_dir, get_save_file_path}; +use crate::utils::{check_or_create_dir, get_save_file_path, sanitize_filename}; use itertools::Itertools; use petgraph::prelude::Graph; use serde::{Deserialize, Serialize}; @@ -14,14 +14,14 @@ use std::path::{Path, PathBuf}; #[serde(rename_all = "camelCase")] pub struct AGCJFunctionCallGraph { pub name: String, - pub size: i64, + pub size: u64, pub imports: Option>, } #[derive(Serialize, Deserialize, Debug)] pub struct AGCJParsedObjects { pub edge_property: String, - pub edges: Vec>, + pub edges: Vec>, pub node_holes: Vec, pub nodes: Vec, } @@ -45,10 +45,7 @@ impl AGCJFunctionCallGraph { let mut function_name = self.name.clone(); - // This is a pretty dirty fix and may break things - if function_name.chars().count() > 100 { - function_name = self.name[..75].to_string(); - } + function_name = sanitize_filename(&function_name); let filename = format!("{}-{}.json", function_name, type_suffix); @@ -82,11 +79,7 @@ impl AGCJFunctionCallGraph { check_or_create_dir(&full_output_path); let mut function_name = self.name.clone(); - - // This is a pretty dirty fix and may break things - if function_name.chars().count() > 100 { - function_name = self.name[..75].to_string(); - } + function_name = sanitize_filename(&function_name); let filename = format!( "{}/{}-{}.json", @@ -121,11 +114,7 @@ impl AGCJFunctionCallGraph { check_or_create_dir(&full_output_path); let mut function_name = self.name.clone(); - - // This is a pretty dirty fix and may break things - if function_name.chars().count() > 100 { - function_name = self.name[..75].to_string(); - } + function_name = sanitize_filename(&function_name); let filename = format!( "{}/{}-{}.json", @@ -159,11 +148,7 @@ impl AGCJFunctionCallGraph { check_or_create_dir(&full_output_path); debug!("Built Path: {:?}", full_output_path); let mut function_name = self.name.clone(); - - // This is a pretty dirty fix and may break things - if function_name.chars().count() > 100 { - function_name = self.name[..75].to_string(); - } + function_name = sanitize_filename(&function_name); let filename = format!("{}-{}.json", function_name, type_suffix); // Normalise string for windows diff --git a/src/agfj.rs b/src/agfj.rs index 90adcdf..289b46c 100644 --- a/src/agfj.rs +++ b/src/agfj.rs @@ -35,18 +35,19 @@ struct EdgePair { wt: u16, } -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug)] pub struct AGFJFunc { pub name: String, + #[serde(alias = "addr")] + pub offset: u64, + pub size: u64, + pub ninstr: u64, nargs: u64, - ninstr: u64, nlocals: u64, - offset: u64, - size: Option, stack: u64, r#type: String, pub blocks: Vec, - addr_idx: Option>, + addr_idx: Option>, pub edge_list: Option>, graph: Option>, } @@ -143,7 +144,7 @@ impl AGFJFunc { pub fn create_bb_edge_list(&mut self, min_blocks: &u16) { if self.blocks.len() > >::into(*min_blocks) && self.blocks[0].offset != 1 { - let bb_start_addrs: Vec = self.blocks.iter().map(|x| x.offset).collect::>(); + let bb_start_addrs: Vec = self.blocks.iter().map(|x| x.offset).collect::>(); let mut edge_list = Vec::<(u32, u32, u32)>::new(); for bb in &self.blocks { @@ -423,7 +424,7 @@ impl AGFJFunc { } }; - let bb_start_addrs: Vec = + let bb_start_addrs: Vec = self.blocks.iter().map(|x| x.offset).collect::>(); match feature_type { @@ -606,7 +607,7 @@ impl AGFJFunc { } // Convert string memory address to hex / string - fn str_to_hex_node_idxs(graph: &mut Graph, addr_idxs: &[i64]) { + fn str_to_hex_node_idxs(graph: &mut Graph, addr_idxs: &[u64]) { for idx in graph.node_indices() { let i_idx = idx.index(); let hex = addr_idxs[i_idx]; @@ -732,6 +733,7 @@ impl From<(&String, Vec)> for TikNibFunc { #[cfg(test)] mod tests { use crate::bb::FeatureType; + use crate::utils::get_default_addr; use std::path::PathBuf; use crate::AGFJFile; @@ -792,7 +794,7 @@ mod tests { assert_eq!(file.functions.as_ref().unwrap()[0][0].ninstr, 78); assert_eq!(file.functions.as_ref().unwrap()[0][0].nargs, 2); assert_eq!(file.functions.as_ref().unwrap()[0][0].nlocals, 0); - assert_eq!(file.functions.as_ref().unwrap()[0][0].size, Some(334)); + assert_eq!(file.functions.as_ref().unwrap()[0][0].size, 334); assert_eq!(file.functions.as_ref().unwrap()[0][0].stack, 56); assert!(!file.functions.as_ref().unwrap()[0][0].blocks.is_empty()); @@ -811,7 +813,10 @@ mod tests { assert!(!file.functions.as_ref().unwrap()[0][0].blocks[0] .ops .is_empty()); - assert_eq!(file.functions.as_ref().unwrap()[0][0].blocks[0].fail, -1); + assert_eq!( + file.functions.as_ref().unwrap()[0][0].blocks[0].fail, + get_default_addr() + ); assert!(file.functions.as_ref().unwrap()[0][0].blocks[0] .switchop diff --git a/src/bb.rs b/src/bb.rs index 2694a03..a397d40 100644 --- a/src/bb.rs +++ b/src/bb.rs @@ -2,6 +2,7 @@ use crate::consts::*; #[cfg(feature = "inference")] use crate::inference::InferenceJob; use crate::normalisation::{normalise_disasm_simple, normalise_esil_simple}; +use crate::utils::get_default_addr; use serde::{Deserialize, Serialize}; use serde_aux::prelude::*; use serde_json::Value; @@ -58,8 +59,9 @@ pub enum InstructionMode { #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] pub struct SwitchOpCase { - pub jump: i64, - pub offset: i64, + pub jump: u64, + #[serde(alias = "addr")] + pub offset: u64, #[serde(deserialize_with = "deserialize_string_from_number")] // This will make it challenging to use downstream but has been // added because sometimes it is a VERY large int (bigger than i64). @@ -74,6 +76,7 @@ pub struct SwitchOp { pub defval: u16, pub maxval: u16, pub minval: u16, + #[serde(alias = "addr")] pub offset: u64, } #[derive(Serialize, Deserialize, Debug, Clone)] @@ -86,6 +89,7 @@ pub struct Op { pub fcn_addr: Option, pub fcn_last: Option, pub flags: Option>, + #[serde(alias = "addr")] pub offset: u64, pub opcode: Option, pub ptr: Option, @@ -100,26 +104,21 @@ pub struct Op { pub val: Option, } -// Function to set offset, jump and fail to default values -fn return_minus_one() -> i64 { - -1 -} - #[serde_as] -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug)] pub struct ACFJBlock { - #[serde(default = "return_minus_one")] - pub offset: i64, - #[serde(default = "return_minus_one")] + #[serde(default = "get_default_addr", alias = "addr")] + pub offset: u64, + #[serde(default = "get_default_addr")] #[serde_as(deserialize_as = "DefaultOnError")] // This has been added to eliminate an error where // the jump address from x86-64 binaries is larger than // an i64. - pub jump: i64, - #[serde(default = "return_minus_one")] - pub fail: i64, + pub jump: u64, + #[serde(default = "get_default_addr")] + pub fail: u64, pub ops: Vec, - pub size: Option, + pub size: Option, pub switchop: Option, } @@ -457,18 +456,18 @@ impl ACFJBlock { } num_offspring } - pub fn get_block_edges(&self, bb_start_addrs: &[i64], edge_list: &mut Vec<(u32, u32, u32)>) { + pub fn get_block_edges(&self, bb_start_addrs: &[u64], edge_list: &mut Vec<(u32, u32, u32)>) { let offset_idx = bb_start_addrs.iter().position(|&p| p == self.offset); if let Some(offset_idx) = offset_idx { - if self.jump != -1 { + if self.jump != get_default_addr() { let jump_idx = bb_start_addrs.iter().position(|&p| p == self.jump); if let Some(jump_idx) = jump_idx { edge_list.push((offset_idx as u32, jump_idx as u32, 1)); } } - if self.fail != -1 { + if self.fail != get_default_addr() { let fail_idx = bb_start_addrs.iter().position(|&p| p == self.fail); if let Some(fail_idx) = fail_idx { edge_list.push((offset_idx as u32, fail_idx as u32, 1)); diff --git a/src/combos.rs b/src/combos.rs index f2d675f..a11371d 100644 --- a/src/combos.rs +++ b/src/combos.rs @@ -133,11 +133,11 @@ impl ComboJob { #[derive(Default, Hash, PartialEq, Clone, Debug, Deserialize, Serialize)] pub struct FinfoTiknib { pub name: String, - pub edges: i64, - pub indegree: i64, - pub outdegree: i64, - pub nlocals: i64, - pub nargs: i64, + pub edges: u64, + pub indegree: u64, + pub outdegree: u64, + pub nlocals: u64, + pub nargs: u64, pub avg_arithshift: OrderedFloat, pub avg_compare: OrderedFloat, pub avg_ctransfer: OrderedFloat, diff --git a/src/extract.rs b/src/extract.rs index 8e9e668..d46bf53 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1,7 +1,7 @@ use crate::afij::AFIJFunctionInfo; use crate::agcj::AGCJFunctionCallGraph; - -use std::io; +use crate::agfj::AGFJFunc; +use crate::utils::sanitize_filename; use anyhow::anyhow; use anyhow::bail; @@ -17,8 +17,9 @@ use serde_json; use glob::glob; use md5; -use regex::Regex; -use serde_json::{json, Deserializer, Value}; +use once_cell::sync::OnceCell; +use serde::ser::{SerializeSeq, Serializer}; +use serde_json::{json, value::RawValue, Deserializer, Value}; use sha1; use sha1::Digest as Sha1Digest; use sha2::Digest as Sha2Digest; @@ -27,9 +28,10 @@ use std::collections::HashMap; use std::env; use std::fs; use std::fs::File; -use std::io::{BufReader, Read}; -use std::path::{Path, PathBuf}; +use std::io::{BufReader, BufWriter, Read}; +use std::path::PathBuf; use std::string::String; +use std::sync::LazyLock; use walkdir::WalkDir; #[derive(PartialEq, Debug)] @@ -39,12 +41,14 @@ pub enum PathType { Dir, Unk, } -#[derive(Debug, PartialEq, Clone, Copy)] + +#[derive(Debug, PartialEq, Clone, Copy, Hash, Eq)] pub enum ExtractionJobType { BinInfo, // Extract high level information from the binary (r2 ij) RegisterBehaviour, FunctionXrefs, CFG, + FunctionCFG, // Like CFG, but in a separate file for each function CallGraphs, FuncInfo, FunctionVariables, @@ -54,9 +58,40 @@ pub enum ExtractionJobType { LocalVariableXrefs, GlobalStrings, FunctionBytes, + FunctionBytesMasked, FunctionZignatures, } +static JOB_TYPE_TO_SUFFIX: LazyLock> = + LazyLock::new(|| { + HashMap::from([ + (ExtractionJobType::BinInfo, "bininfo"), + (ExtractionJobType::RegisterBehaviour, "reg"), + (ExtractionJobType::FunctionXrefs, "func-xrefs"), + (ExtractionJobType::CFG, "cfg"), + (ExtractionJobType::FunctionCFG, "func-cfg"), + (ExtractionJobType::CallGraphs, "cg"), + (ExtractionJobType::FuncInfo, "finfo"), + (ExtractionJobType::FunctionVariables, "fvars"), + (ExtractionJobType::Decompilation, "decomp"), + (ExtractionJobType::PCodeFunc, "pcode-func"), + (ExtractionJobType::PCodeBB, "pcode-bb"), + (ExtractionJobType::LocalVariableXrefs, "localvar-xrefs"), + (ExtractionJobType::GlobalStrings, "strings"), + (ExtractionJobType::FunctionBytes, "bytes"), + (ExtractionJobType::FunctionBytesMasked, "bytes-masked"), + (ExtractionJobType::FunctionZignatures, "zigs"), + ]) + }); + +static SUFFIX_TO_JOB_TYPE: LazyLock> = + LazyLock::new(|| { + JOB_TYPE_TO_SUFFIX + .iter() + .map(|(job_type, suffix)| (*suffix, *job_type)) + .collect() + }); + #[derive(Debug)] pub struct FileToBeProcessed { pub file_path: PathBuf, @@ -64,6 +99,21 @@ pub struct FileToBeProcessed { pub job_types: Vec, pub r2p_config: R2PipeConfig, pub with_annotations: bool, + pub keep_raw_bytes: bool, + pub retry_aborted: bool, + pub func_filename_template: String, + pub function_list: OnceCell>, + pub min_basic_blocks: Option, + needs_func_list: bool, +} + +#[derive(Debug, Clone)] +pub struct FunctionToBeProcessed { + pub name: String, + pub addr: u64, + pub size: u64, + pub ninstrs: u64, + pub nblocks: u64, } #[derive(Debug)] @@ -75,11 +125,13 @@ pub struct ExtractionJob { pub output_path: PathBuf, } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone)] pub struct R2PipeConfig { pub debug: bool, - pub extended_analysis: bool, + pub analysis_mode: String, pub use_curl_pdb: bool, + pub apply_relocations: bool, + pub disable_pseudo_asm: bool, pub timeout: Option, } @@ -97,6 +149,7 @@ impl std::fmt::Display for ExtractionJob { #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct AFLJFuncDetails { + #[serde(alias = "addr")] pub offset: u64, pub name: String, pub size: u64, @@ -104,6 +157,8 @@ pub struct AFLJFuncDetails { pub is_pure: String, pub realsz: u64, pub noreturn: bool, + #[serde(default)] + pub recursive: bool, pub stackframe: u64, pub calltype: String, pub cost: u64, @@ -112,15 +167,26 @@ pub struct AFLJFuncDetails { #[serde(rename = "type")] pub type_field: String, pub nbbs: u64, + #[serde(default)] + pub tracecov: u64, #[serde(rename = "is-lineal")] pub is_lineal: bool, pub ninstrs: u64, pub edges: u64, pub ebbs: u64, + #[serde(default)] pub signature: String, - pub minbound: i64, + #[serde(alias = "minaddr")] + pub minbound: u64, + #[serde(alias = "maxaddr")] pub maxbound: u64, #[serde(default)] + pub maxbbins: Option, + #[serde(default)] + pub midbbins: Option, + #[serde(default)] + pub ratbbins: Option, + #[serde(default)] pub callrefs: Vec, #[serde(default)] pub datarefs: Vec, @@ -233,9 +299,33 @@ impl std::fmt::Display for AFLJFuncDetails { } } -impl From<(String, String, Vec, R2PipeConfig, bool)> for FileToBeProcessed { +impl + From<( + String, + String, + Vec, + R2PipeConfig, + bool, + bool, + bool, + String, + Option, + bool, + )> for FileToBeProcessed +{ fn from( - orig: (String, String, Vec, R2PipeConfig, bool), + orig: ( + String, + String, + Vec, + R2PipeConfig, + bool, + bool, + bool, + String, + Option, + bool, + ), ) -> FileToBeProcessed { FileToBeProcessed { file_path: PathBuf::from(orig.0), @@ -243,6 +333,60 @@ impl From<(String, String, Vec, R2PipeConfig, bool)> for File job_types: orig.2, r2p_config: orig.3, with_annotations: orig.4, + keep_raw_bytes: orig.5, + retry_aborted: orig.6, + func_filename_template: orig.7, + function_list: OnceCell::new(), + min_basic_blocks: orig.8, + needs_func_list: orig.9, + } + } +} + +impl From<(String, u64, u64, u64, u64)> for FunctionToBeProcessed { + fn from(orig: (String, u64, u64, u64, u64)) -> Self { + FunctionToBeProcessed { + name: orig.0, + addr: orig.1, + size: orig.2, + ninstrs: orig.3, + nblocks: orig.4, + } + } +} + +impl From for FunctionToBeProcessed { + fn from(func_info: AFIJFunctionInfo) -> Self { + FunctionToBeProcessed { + name: func_info.name, + addr: func_info.offset, + size: func_info.size, + ninstrs: func_info.ninstrs, + nblocks: func_info.nbbs, + } + } +} + +impl From for FunctionToBeProcessed { + fn from(func_details: AFLJFuncDetails) -> Self { + FunctionToBeProcessed { + name: func_details.name, + addr: func_details.offset, + size: func_details.size, + ninstrs: func_details.ninstrs, + nblocks: func_details.nbbs, + } + } +} + +impl From for FunctionToBeProcessed { + fn from(func: AGFJFunc) -> Self { + FunctionToBeProcessed { + name: func.name, + addr: func.offset, + size: func.size, + ninstrs: func.ninstr, + nblocks: func.blocks.len() as u64, } } } @@ -296,6 +440,15 @@ pub struct PCodeJsonWithBBAndFuncName { // Structs for afbj - Basic Block JSON output pub type BasicBlockInfo = Vec; +// Custom deserializer for converting integer (0 or 1) to boolean +fn deserialize_bool_from_int<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + let num = i64::deserialize(deserializer)?; + Ok(num != 0) // 0 is false, anything else is true +} + #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct BasicBlockMetadataEntry { pub addr: u64, @@ -307,6 +460,7 @@ pub struct BasicBlockMetadataEntry { pub outputs: u64, pub ninstr: u64, pub instrs: Vec, + #[serde(deserialize_with = "deserialize_bool_from_int")] pub traced: bool, } @@ -353,6 +507,7 @@ pub struct StringEntry { #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct FuncBytes { pub bytes: Vec, + pub mask: Option>, } // Structs for zj - Function signatures (called "zignatures" in r2) @@ -476,144 +631,195 @@ impl ExtractionJob { output_path: &PathBuf, modes: &Vec, debug: &bool, - extended_analysis: &bool, + analysis_mode: &str, use_curl_pdb: &bool, + apply_relocations: &bool, + disable_pseudo_asm: &bool, + func_filename_template: &str, timeout: &Option, with_annotations: &bool, + keep_raw_bytes: &bool, + retry_aborted: &bool, + min_basic_blocks: &Option, ) -> Result { - fn get_path_type(bin_path: &PathBuf) -> PathType { - // Handle pattern first since it would raise NotFound error - let path_str = bin_path.to_string_lossy(); - if path_str.contains('*') || path_str.contains('?') || path_str.contains('[') { - return PathType::Pattern; - } + let mut job_types = vec![]; + let mut extraction_job_types = vec![]; - let fpath_md = fs::metadata(bin_path).unwrap(); - if fpath_md.is_file() { - PathType::File - } else if fpath_md.is_dir() { - PathType::Dir + for mode in modes { + let job_type = Self::extraction_job_matcher(mode)?; + + // FuncInfo should always be processed first + if job_type == ExtractionJobType::FuncInfo && !job_types.is_empty() { + job_types.insert(0, (job_type, mode.clone())); + extraction_job_types.insert(0, job_type); } else { - PathType::Unk + job_types.push((job_type, mode.clone())); + extraction_job_types.push(job_type); } } - // This function is used to validate modes and convert them to job types - fn extraction_job_matcher(mode: &str) -> Result { - match mode { - "bininfo" => Ok(ExtractionJobType::BinInfo), - "finfo" => Ok(ExtractionJobType::FuncInfo), - "fvars" => Ok(ExtractionJobType::FunctionVariables), - "reg" => Ok(ExtractionJobType::RegisterBehaviour), - "cfg" => Ok(ExtractionJobType::CFG), - "func-xrefs" => Ok(ExtractionJobType::FunctionXrefs), - "cg" => Ok(ExtractionJobType::CallGraphs), - "decomp" => Ok(ExtractionJobType::Decompilation), - "pcode-func" => Ok(ExtractionJobType::PCodeFunc), - "pcode-bb" => Ok(ExtractionJobType::PCodeBB), - "localvar-xrefs" => Ok(ExtractionJobType::LocalVariableXrefs), - "strings" => Ok(ExtractionJobType::GlobalStrings), - "bytes" => Ok(ExtractionJobType::FunctionBytes), - "zigs" => Ok(ExtractionJobType::FunctionZignatures), - _ => bail!("Incorrect command type - got {}", mode), - } + // Warn the user if mode-specific flags are turned on for no reason + if !extraction_job_types.contains(&ExtractionJobType::Decompilation) && *with_annotations { + let mode = ExtractionJob::get_job_type_suffix(&ExtractionJobType::Decompilation); + warn!( + "Annotations are only supported for decompilation extraction (mode: {})", + mode + ); } - - let mut job_types = vec![]; - let mut extraction_job_types = vec![]; - - for mode in modes { - let job_type = extraction_job_matcher(mode)?; - job_types.push((job_type, mode.clone())); - extraction_job_types.push(job_type); // Store just the job type - - if job_type != ExtractionJobType::Decompilation && *with_annotations { - warn!( - "Annotations are only supported for decompilation extraction (mode: {})", - mode - ); - } + if !extraction_job_types.contains(&ExtractionJobType::FunctionBytesMasked) + && *keep_raw_bytes + { + let mode = ExtractionJob::get_job_type_suffix(&ExtractionJobType::FunctionBytesMasked); + warn!( + "Keep raw bytes is only supported for masked bytes extraction (mode: {})", + mode + ); } let r2_handle_config = R2PipeConfig { debug: *debug, - extended_analysis: *extended_analysis, + analysis_mode: analysis_mode.to_string(), + apply_relocations: *apply_relocations, + disable_pseudo_asm: *disable_pseudo_asm, use_curl_pdb: *use_curl_pdb, timeout: *timeout, }; - let p_type = get_path_type(input_path); - - if p_type == PathType::File { - // For a single file, create one FileToBeProcessed object - // but track all the job types - let file = FileToBeProcessed { - file_path: input_path.to_owned(), - output_path: output_path.to_owned(), - job_types: extraction_job_types, // Use the vector of just ExtractionJobType - r2p_config: r2_handle_config, - with_annotations: *with_annotations, - }; - - Ok(ExtractionJob { - input_path: input_path.to_owned(), - input_path_type: p_type, - job_types, - files_to_be_processed: vec![file], - output_path: output_path.to_owned(), - }) - } else if p_type == PathType::Dir { - // For a directory, get all file paths - let files = ExtractionJob::get_file_paths_dir(input_path); - - // Create FileToBeProcessed objects for each file with all job types - let files_to_be_processed = files - .into_iter() - .map(|f| FileToBeProcessed { - file_path: PathBuf::from(f), + // Check if any of the job types requires the function list + let needs_func_list = extraction_job_types + .iter() + .any(|jt| Self::job_type_needs_function_list(jt)); + + let p_type = Self::get_path_type(input_path); + match p_type { + PathType::File => { + // For a single file, create one FileToBeProcessed object + // but track all the job types + let file = FileToBeProcessed { + file_path: input_path.to_owned(), output_path: output_path.to_owned(), - job_types: extraction_job_types.clone(), + job_types: extraction_job_types, // Use the vector of just ExtractionJobType r2p_config: r2_handle_config, with_annotations: *with_annotations, + keep_raw_bytes: *keep_raw_bytes, + retry_aborted: *retry_aborted, + func_filename_template: func_filename_template.to_string(), + function_list: OnceCell::new(), + min_basic_blocks: *min_basic_blocks, + needs_func_list, + }; + + Ok(ExtractionJob { + input_path: input_path.to_owned(), + input_path_type: p_type, + job_types, + files_to_be_processed: vec![file], + output_path: output_path.to_owned(), }) - .collect(); - - Ok(ExtractionJob { - input_path: input_path.to_owned(), - input_path_type: p_type, - job_types, - files_to_be_processed, - output_path: output_path.to_owned(), - }) - } else if p_type == PathType::Pattern { - // For a match pattern get the list of matching file paths - let pattern = input_path.to_string_lossy(); - let files = ExtractionJob::get_file_paths_pattern(&pattern); - - // Create FileToBeProcessed objects for each file with all job types - let files_to_be_processed = files - .into_iter() - .map(|f| FileToBeProcessed { - file_path: PathBuf::from(f), + } + PathType::Dir | PathType::Pattern => { + // Get file paths based on path type + let files = match p_type { + PathType::Dir => ExtractionJob::get_file_paths_dir(input_path), + PathType::Pattern => { + let pattern = input_path.to_string_lossy(); + ExtractionJob::get_file_paths_pattern(&pattern) + } + _ => unreachable!(), + }; + + // Create FileToBeProcessed objects for each file with all job types + let files_to_be_processed = files + .into_iter() + .map(|f| FileToBeProcessed { + file_path: PathBuf::from(f), + output_path: output_path.to_owned(), + job_types: extraction_job_types.clone(), + r2p_config: r2_handle_config.clone(), + with_annotations: *with_annotations, + keep_raw_bytes: *keep_raw_bytes, + retry_aborted: *retry_aborted, + func_filename_template: func_filename_template.to_string(), + function_list: OnceCell::new(), + min_basic_blocks: *min_basic_blocks, + needs_func_list, + }) + .collect(); + + Ok(ExtractionJob { + input_path: input_path.to_owned(), + input_path_type: PathType::Dir, // For using parallel processing + job_types, + files_to_be_processed, output_path: output_path.to_owned(), - job_types: extraction_job_types.clone(), - r2p_config: r2_handle_config, - with_annotations: *with_annotations, }) - .collect(); - - Ok(ExtractionJob { - input_path: input_path.to_owned(), - input_path_type: PathType::Dir, // For using parallel processing - job_types, - files_to_be_processed, - output_path: output_path.to_owned(), - }) + } + PathType::Unk => bail!("Failed to create ExtractionJob"), + } + } + + /// Get the type of the input path (file, directory, or pattern) + fn get_path_type(bin_path: &PathBuf) -> PathType { + // Handle pattern first since it would raise NotFound error + let path_str = bin_path.to_string_lossy(); + if path_str.contains('*') || path_str.contains('?') || path_str.contains('[') { + return PathType::Pattern; + } + + let fpath_md = fs::metadata(bin_path).unwrap(); + if fpath_md.is_file() { + PathType::File + } else if fpath_md.is_dir() { + PathType::Dir } else { - bail!("Failed to create ExtractionJob") + PathType::Unk + } + } + + /// Validate extraction modes and convert them to job types + fn extraction_job_matcher(mode: &str) -> Result { + SUFFIX_TO_JOB_TYPE + .get(mode) + .copied() + .ok_or_else(|| anyhow!("Incorrect command type - got {}", mode)) + } + + fn get_job_type_suffix(job_type: &ExtractionJobType) -> String { + JOB_TYPE_TO_SUFFIX + .get(job_type) + .copied() + .expect("Incorrect command type") + .to_string() + } + + fn get_output_extension(job_type: &ExtractionJobType) -> Option<&str> { + match job_type { + // Add here if output is not a JSON file (None if a directory) + ExtractionJobType::FunctionBytes => None, + ExtractionJobType::FunctionBytesMasked => None, + ExtractionJobType::FunctionCFG => None, + _ => Some("json"), } } + /// Check if a job type requires the function list to be populated + fn job_type_needs_function_list(job_type: &ExtractionJobType) -> bool { + matches!( + job_type, + ExtractionJobType::RegisterBehaviour + | ExtractionJobType::FunctionXrefs + | ExtractionJobType::FunctionCFG + | ExtractionJobType::FunctionVariables + | ExtractionJobType::Decompilation + | ExtractionJobType::PCodeFunc + | ExtractionJobType::PCodeBB + | ExtractionJobType::LocalVariableXrefs + | ExtractionJobType::FunctionBytes + | ExtractionJobType::FunctionBytesMasked + ) + } + + /// Get all file paths in the input_path directory fn get_file_paths_dir(input_path: &PathBuf) -> Vec { let mut str_vec: Vec = Vec::new(); for file in WalkDir::new(input_path) @@ -631,9 +837,9 @@ impl ExtractionJob { str_vec } + /// Get all file paths that match the pattern fn get_file_paths_pattern(pattern: &str) -> Vec { let mut paths = Vec::new(); - // glob returns an iterator over Result for entry in glob(pattern) .expect("Failed to read glob pattern") .flatten() @@ -642,83 +848,558 @@ impl ExtractionJob { paths.push(entry.to_string_lossy().to_string()); } } - paths } } +impl FunctionToBeProcessed { + /// Write function bytes to a binary file with suffix .bin + /// If apply_mask is true, write function masked bytes to a binary file with suffix .masked.bin + fn write_to_bin( + &self, + r2p: &mut R2Pipe, + output_dirpath: &PathBuf, + filename_template: &str, + apply_mask: bool, + keep_raw_bytes: bool, + ) -> Result<()> { + let func_bytes = self + .get_bytes(r2p, apply_mask) + .map_err(|e| anyhow::anyhow!("Failed to get bytes: {}", e))?; + + // Store the bytes if the user specified to keep the raw bytes + // or if the byte mask is not applied + if keep_raw_bytes || !apply_mask { + // Setup output filepaths for function bytes + let bytes_ext = + FunctionToBeProcessed::get_function_file_ext(&ExtractionJobType::FunctionBytes); + let bytes_filepath = + self.get_output_filepath(output_dirpath, filename_template, bytes_ext); + + debug!("Writing function bytes to file: {:?}", bytes_filepath); + std::fs::write(&bytes_filepath, func_bytes.bytes).with_context(|| { + format!( + "Failed to write function bytes to file: {:?}", + bytes_filepath + ) + })?; + } + + if apply_mask { + let bytes_mask = func_bytes + .mask + .context("Masked bytes missing from get_bytes output")?; + + // Setup output filepaths for masked bytes + let masked_bytes_ext = FunctionToBeProcessed::get_function_file_ext( + &ExtractionJobType::FunctionBytesMasked, + ); + let masked_bytes_filepath = + self.get_output_filepath(output_dirpath, filename_template, masked_bytes_ext); + + debug!( + "Writing function masked bytes to file: {:?}", + masked_bytes_filepath + ); + std::fs::write(&masked_bytes_filepath, bytes_mask).with_context(|| { + format!( + "Failed to write function masked bytes to file: {:?}", + masked_bytes_filepath + ) + })?; + } + Ok(()) + } + + fn write_to_json( + &self, + json_obj: &Value, + output_dirpath: &PathBuf, + filename_template: &str, + ) -> Result<()> { + let output_filepath = self.get_output_filepath(output_dirpath, filename_template, "json"); + debug!("Writing JSON to {:?}", output_filepath); + let file = File::create(&output_filepath) + .with_context(|| format!("Unable to create file {:?}", output_filepath))?; + serde_json::to_writer(file, json_obj) + .with_context(|| format!("Unable to write JSON to {:?}", output_filepath))?; + Ok(()) + } + + fn write_cfg_to_json( + &self, + r2p: &mut R2Pipe, + output_dirpath: &PathBuf, + filename_template: &str, + ) -> Result<()> { + let cfg = self + .get_cfg(r2p) + .context(format!("Failed to get CFG @ {}", self.addr))?; + // Make sure self.addr and cfg_json.addr are the same + if self.addr != cfg.offset { + // If they aren't, write the CFG to a JSON file with the other function's details + warn!("Function address mismatch: {} != {}", self.addr, cfg.offset); + let other_function = FunctionToBeProcessed::from(cfg.clone()); + other_function.write_to_json(&json!(cfg), output_dirpath, filename_template)?; + } else { + self.write_to_json(&json!(cfg), output_dirpath, filename_template)?; + } + Ok(()) + } + + // Getters + pub fn get_function_file_ext(job_type: &ExtractionJobType) -> &str { + // Based on the job type, return the file extension for the output function file + match job_type { + ExtractionJobType::FunctionBytes => "raw.bin", + ExtractionJobType::FunctionBytesMasked => "masked.bin", + ExtractionJobType::FunctionCFG => "json", + _ => "", + } + } + + fn get_hex_address(&self) -> String { + format!("{:x}", self.addr) + .trim_start_matches("0x") + .to_string() + } + + fn get_output_filename(&self, template: &str, ext: &str) -> String { + let mut func_filename = match template { + "symbol" => self.name.clone(), + "address" => self.get_hex_address(), + _ => template + .replace("{symbol}", &self.name) + .replace("{address}", &self.get_hex_address()) + .replace("{ext}", ext), + }; + + func_filename = sanitize_filename(&func_filename); + if ["symbol", "address"].contains(&template) && ext != "" { + // Add an extension only if the user did not specify a custom template + // and the ext string is not empty + func_filename = func_filename + "." + ext; + } + func_filename + } + + pub fn get_output_filepath( + &self, + output_dirpath: &PathBuf, + filename_template: &str, + ext: &str, + ) -> PathBuf { + let output_filename = self.get_output_filename(filename_template, ext); + + let mut output_filepath = PathBuf::new(); + output_filepath.push(output_dirpath); + output_filepath.push(output_filename); + output_filepath + } + + /// Build a length-aligned mask for the function containing `func_addr`. + /// + /// Uses func_bytes_len (from p8fm) as ground truth since that's what's saved to disk. + /// Strategy: Try abmj per-BB, fall back to batch aoj for speed. + fn rebuild_mask_with_bb_and_instr( + &self, + r2p: &mut R2Pipe, + func_addr: u64, + func_bytes_len: usize, + ) -> Result, Error> { + // Get function start for BB enumeration + let afij = r2p.cmdj(&format!("afij @ {}", func_addr))?; + let fobj = afij + .as_array() + .and_then(|a| a.get(0)) + .ok_or_else(|| anyhow::anyhow!("`afij` returned no function at 0x{:x}", func_addr))?; + + let func_base = fobj + .get("offset") + .or_else(|| fobj.get("addr")) + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("Missing function offset/addr in `afij`"))?; + + // Pre-size mask to match p8fm output (ground truth) + let mut mask = vec![0u8; func_bytes_len]; + + // Enumerate basic blocks + let afbj = r2p.cmdj(&format!("afbj @ {}", func_base))?; + let blocks = afbj + .as_array() + .ok_or_else(|| anyhow::anyhow!("`afbj` did not return a JSON array"))?; + + for bb in blocks { + let bb_addr = bb + .get("addr") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("BB missing `addr` in `afbj`"))?; + let bb_size = bb + .get("size") + .and_then(|v| v.as_u64()) + .ok_or_else(|| anyhow::anyhow!("BB missing `size` in `afbj`"))?; + + // Skip empty BBs or BBs outside our range + if bb_size == 0 || bb_addr < func_base { + continue; + } + + let off = (bb_addr - func_base) as usize; + let blen = bb_size as usize; + + if off >= mask.len() { + continue; + } + + let cap = off.saturating_add(blen).min(mask.len()); + let slice_len = cap - off; + + // Tier 1: abmj (BB bytes+mask) — fastest when available + if let Ok(v) = r2p.cmdj(&format!("abmj @ {}", bb_addr)) { + if let (Some(bh), Some(mh)) = ( + v.get("bytes").and_then(|x| x.as_str()), + v.get("mask").and_then(|x| x.as_str()), + ) { + // Check hex string lengths before decoding (2 chars per byte) + if bh.len() == mh.len() && bh.len() / 2 == blen { + if let Ok(bb_mask) = hex::decode(mh) { + if bb_mask.len() >= slice_len { + mask[off..cap].copy_from_slice(&bb_mask[..slice_len]); + continue; + } + } + } + } + } + + // Tier 2: batch aoj (fast and works reliably) + let max_instrs = blen; + if let Ok(aoj_result) = r2p.cmdj(&format!("aoj {} @ {}", max_instrs, bb_addr)) { + if let Some(instrs) = aoj_result.as_array() { + let mut p = 0usize; + for instr in instrs { + if p >= blen || (off + p) >= mask.len() { + break; + } + let sz = instr.get("size").and_then(|v| v.as_u64()).unwrap_or(0) as usize; + if sz == 0 { + mask[off + p] = 0x00; + p += 1; + } else { + let endp = (p + sz).min(blen); + for i in p..endp { + if (off + i) < mask.len() { + mask[off + i] = 0xFF; + } + } + p = endp; + } + } + while p < blen && (off + p) < mask.len() { + mask[off + p] = 0x00; + p += 1; + } + continue; + } + } + + // Tier 3: last resort - wildcard entire BB + for i in 0..slice_len { + mask[off + i] = 0x00; + } + } + + Ok(mask) + } + + fn get_bytes(&self, r2p: &mut R2Pipe, apply_mask: bool) -> Result { + let cmd_str; + let function_bytes; + let mut function_mask; + let function_bytes_and_mask; + let mut masked_bytes: Option> = None; + + if apply_mask { + cmd_str = format!("p8fm @ {}", self.addr); + debug!( + "Getting function bytes and mask for function: `{}`", + cmd_str + ); + function_bytes_and_mask = r2p + .cmd(&cmd_str) + .context("Failed to execute `{}`")? + .trim() + .to_string(); + let parts: Vec<&str> = function_bytes_and_mask.split(":").collect(); + + if parts.len() < 2 { + return Err(anyhow::anyhow!( + "Invalid output format: expected 'bytes:mask'.\n\ + Output from `{}`: '{}'", + cmd_str, + function_bytes_and_mask + )); + } + function_bytes = hex::decode(parts[0]) + .with_context(|| format!("Failed to decode hex function bytes: '{}'", parts[0]))?; + function_mask = hex::decode(parts[1]) + .with_context(|| format!("Failed to decode hex function mask: '{}'", parts[1]))?; + + // Ensure function_bytes and bytes_mask have the same length + if function_bytes.len() != function_mask.len() { + warn!("Function bytes length ({}) and mask length ({}) do not match. Rebuilding mask.", function_bytes.len(), function_mask.len()); + // Fix-up: rebuild a length-aligned mask from BBs and (if needed) per-instruction + function_mask = + self.rebuild_mask_with_bb_and_instr(r2p, self.addr, function_bytes.len())?; + } + + // Try again with the fixed mask + if function_mask.len() != function_bytes.len() { + error!( + "After fix-up, lengths still differ: bytes={} mask={}. Original p8fm was '{}'", + function_bytes.len(), + function_mask.len(), + function_bytes_and_mask + ); + masked_bytes = None; + } else { + let mut masked_bytes_tmp = vec![0; function_bytes.len()]; + for i in 0..function_bytes.len() { + masked_bytes_tmp[i] = function_bytes[i] & function_mask[i]; + } + masked_bytes = Some(masked_bytes_tmp); + } + } else { + cmd_str = format!("p8f @ {}", self.addr); + debug!("Getting function bytes for function: `{}`", cmd_str); + let function_bytes_str = r2p + .cmd(&cmd_str) + .context("Failed to execute `{}`")? + .trim() + .to_string(); + function_bytes = hex::decode(function_bytes_str.clone()).with_context(|| { + format!( + "Failed to decode hex function bytes: '{}'", + function_bytes_str + ) + })?; + } + + Ok(FuncBytes { + bytes: function_bytes, + mask: masked_bytes, + }) + } + + fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { + let cmd_str = format!("afbj @ {}", self.addr); + debug!("Getting Basic Block Info for function: `{}`", cmd_str); + let value = r2p.cmdj(&cmd_str).context("Command afbj failed")?; + + let bb_info: BasicBlockInfo = + serde_json::from_value(value).context("Unable to convert to BasicBlockInfo struct!")?; + Ok(bb_info) + } + + fn get_local_variable_xref_details( + &self, + r2p: &mut R2Pipe, + ) -> Result { + info!( + "Getting local variable xref details for function @ {}", + self.addr + ); + let cmd_str = format!("axvj @ {}", self.addr); + debug!( + "Getting Local Variable Xref details for function: `{}`", + cmd_str + ); + let json = r2p.cmdj(&cmd_str).context("Command axvj failed")?; + + let local_variable_xrefs: LocalVariableXrefs = serde_json::from_value(json) + .context("Unable to convert to LocalVariableXrefs struct!")?; + Ok(local_variable_xrefs) + } + + fn get_xref_details(&self, r2p: &mut R2Pipe) -> Result, Error> { + info!("Getting xref details for function @ {}", self.addr); + let cmd_str = format!("axffj @ {}", self.addr); + debug!("Getting Xref details for function: `{}`", cmd_str); + let json = r2p.cmdj(&cmd_str).context("Command axffj failed")?; + let mut json_obj: Vec = serde_json::from_value(json) + .context("Unable to convert to FunctionXrefDetails struct!")?; + + // TODO: There is a minor bug in this where functions without any xrefs are included. + // Been left in as may be useful later down the line. + if !json_obj.is_empty() { + debug!("Replacing all CALL xrefs with actual function name"); + for element in json_obj.iter_mut() { + if element.type_field == "CALL" { + let cmd_str = format!("afi. @ {}", &element.ref_field); + let function_name = r2p.cmd(cmd_str.as_str()).context("Command afi. failed")?; + element.name = function_name.trim().to_string(); + } + } + }; + Ok(json_obj) + } + + fn get_ghidra_decomp( + &self, + r2p: &mut R2Pipe, + with_annotations: bool, + ) -> Result { + let cmd_str = format!("pdgj @ {}", self.addr); + debug!("Getting Ghidra Decomp for function: `{}`", cmd_str); + let json = r2p.cmdj(&cmd_str).context("Command pdgj failed")?; + + if with_annotations { + let json_obj: DecompJSON = + serde_json::from_value(json).context("Unable to convert to DecompJSON struct!")?; + Ok(json_obj) + } else { + let json_obj: Value = + serde_json::from_value(json).context("Unable to convert to JSON object!")?; + let parsed_code = json_obj["code"] + .as_str() + .context("Unable to get code from JSON object!")? + .to_string(); + let parsed_obj = DecompJSON { + code: parsed_code, + annotations: Vec::new(), + }; + Ok(parsed_obj) + } + } + + fn get_cfg(&self, r2p: &mut R2Pipe) -> Result { + let cmd_str = format!("agfj @ {}", self.addr); + debug!("Getting CFG for function: `{}`", cmd_str); + let json = r2p.cmdj(&cmd_str).context("Command agfj failed")?; + + // AGFJ returns an array of JSON objects, it should only ever be length 1 + let cfg: Vec = serde_json::from_value(json.clone()) + .with_context(|| format!("Unable to convert {:?} to AGFJFunc struct!", json))?; + if cfg.len() == 0 { + return Err(anyhow!("No CFG found for function @ {}", self.addr)); + } else if cfg.len() > 1 { + warn!( + "Multiple CFGs found for function @ {}; ignoring all but the first", + self.addr + ); + } + let func_cfg = cfg[0].clone(); + Ok(func_cfg) + } +} + impl FileToBeProcessed { - pub fn get_output_filename(&self, job_type_suffix: &str) -> String { - let mut fp_filename = self - .file_path + /// Returns the name of the input binary file + fn get_file_name(&self) -> Result { + self.file_path .file_name() - .expect("Unable to get filename") - .to_string_lossy() - .to_string(); + .ok_or_else(|| anyhow!("Unable to get file name from {:?}", self.file_path)) + .map(|os_str| os_str.to_string_lossy().to_string()) + } - fp_filename = if job_type_suffix == "bytes" || job_type_suffix == "bytes.__part" { - fp_filename + "_" + job_type_suffix - } else if self.with_annotations { - fp_filename + "_" + job_type_suffix + "_annotations" + ".json" + /// Returns the name of the output file for a given job type + fn get_output_filename(&self, job_type_suffix: &str) -> Result { + let job_type = ExtractionJob::extraction_job_matcher(job_type_suffix) + .context(format!("Incorrect job type suffix: {}", job_type_suffix))?; + let ext = ExtractionJob::get_output_extension(&job_type); + let ext_str = ext.map_or("".to_string(), |e| format!(".{}", e)); + let mut output_filename = self.get_file_name()?; + + if job_type == ExtractionJobType::Decompilation && self.with_annotations { + output_filename = output_filename + "_" + job_type_suffix + "_annotations" + &ext_str; } else { - fp_filename + "_" + job_type_suffix + ".json" + output_filename = output_filename + "_" + job_type_suffix + &ext_str; }; - fp_filename + + Ok(output_filename) } - pub fn get_output_filepath(&self, job_type_suffix: &str) -> PathBuf { - let fp_filename = self.get_output_filename(job_type_suffix); + /// Returns the path where an error message for a given job type should be stored + fn get_error_filepath(&self, job_type_suffix: &str) -> Result { + // Maintain the original name with suffix, but replace the extension with .error.log + // NOTE: The method .with_extension() won't work if the filename has no extension and contains a `.` + let mut output_filename = self.get_output_filename(job_type_suffix)?; + let job_type = ExtractionJob::extraction_job_matcher(job_type_suffix)?; + let suffix = if job_type != ExtractionJobType::Decompilation { + job_type_suffix.to_owned() + } else { + job_type_suffix.to_owned() + "_annotations" + }; + let suffix_with_ext = suffix.clone() + ".error.log"; + // Remove everything after the suffix in the output filename + output_filename = output_filename.split(&suffix).next().unwrap().to_string(); + output_filename = output_filename + &suffix_with_ext; + let mut error_filepath = PathBuf::from(self.output_path.clone()); + error_filepath.push(output_filename); + + Ok(error_filepath) + } - let mut output_filepath = PathBuf::new(); - output_filepath.push(self.output_path.clone()); - output_filepath.push(fp_filename); + /// Returns the path where the output data for a given job type should be stored + fn get_output_filepath(&self, job_type_suffix: &str) -> Result { + let output_filename = self.get_output_filename(job_type_suffix)?; + let mut output_filepath = PathBuf::from(self.output_path.clone()); + output_filepath.push(output_filename); + Ok(output_filepath) + } - output_filepath.clone() + fn get_tmp_output_filepath(&self, job_type_suffix: &str) -> Result { + let mut filepath_str = self + .get_output_filepath(job_type_suffix)? + .to_string_lossy() + .to_string(); + filepath_str = filepath_str + ".part"; + let output_filepath = PathBuf::from(filepath_str.clone()); + Ok(output_filepath) } pub fn process_mode(&self, r2p: &mut R2Pipe, job_type: &ExtractionJobType) -> Result<()> { - let job_type_suffix = self.get_job_type_suffix(job_type); + let job_type_suffix = ExtractionJob::get_job_type_suffix(job_type); + + // Prepare final output file path (could also be a directory) + let output_path = self.get_output_filepath(&job_type_suffix)?; // Use temporary name to keep track of incomplete extraction - let tmp_job_type_suffix = format!("{}.__part", job_type_suffix).to_string(); - let tmp_output_path = self.get_output_filepath(&tmp_job_type_suffix); + let tmp_output_path = self.get_tmp_output_filepath(&job_type_suffix)?; match job_type { - ExtractionJobType::BinInfo => self.extract_binary_info(r2p, tmp_job_type_suffix), + ExtractionJobType::BinInfo => self.extract_binary_info(r2p, &tmp_output_path), ExtractionJobType::RegisterBehaviour => { - self.extract_register_behaviour(r2p, tmp_job_type_suffix) + self.extract_register_behaviour(r2p, &tmp_output_path) } - ExtractionJobType::FunctionXrefs => { - self.extract_function_xrefs(r2p, tmp_job_type_suffix) - } - ExtractionJobType::CFG => self.extract_func_cfgs(r2p, tmp_job_type_suffix), + ExtractionJobType::FunctionXrefs => self.extract_function_xrefs(r2p, &tmp_output_path), + ExtractionJobType::CFG => self.extract_func_cfgs(r2p, &tmp_output_path, false), + ExtractionJobType::FunctionCFG => self.extract_func_cfgs(r2p, &tmp_output_path, true), ExtractionJobType::CallGraphs => { - self.extract_function_call_graphs(r2p, tmp_job_type_suffix) + self.extract_function_call_graphs(r2p, &tmp_output_path) } - ExtractionJobType::FuncInfo => self.extract_function_info(r2p, tmp_job_type_suffix), + ExtractionJobType::FuncInfo => self.extract_function_info(r2p, &tmp_output_path), ExtractionJobType::FunctionVariables => { - self.extract_function_variables(r2p, tmp_job_type_suffix) - } - ExtractionJobType::Decompilation => { - self.extract_decompilation(r2p, tmp_job_type_suffix) + self.extract_function_variables(r2p, &tmp_output_path) } - ExtractionJobType::PCodeFunc => self.extract_pcode_function(r2p, tmp_job_type_suffix), - ExtractionJobType::PCodeBB => self.extract_pcode_basic_block(r2p, tmp_job_type_suffix), + ExtractionJobType::Decompilation => self.extract_decompilation(r2p, &tmp_output_path), + ExtractionJobType::PCodeFunc => self.extract_pcode_function(r2p, &tmp_output_path), + ExtractionJobType::PCodeBB => self.extract_pcode_basic_block(r2p, &tmp_output_path), ExtractionJobType::LocalVariableXrefs => { - self.extract_local_variable_xrefs(r2p, tmp_job_type_suffix) - } - ExtractionJobType::GlobalStrings => { - self.extract_global_strings(r2p, tmp_job_type_suffix) + self.extract_local_variable_xrefs(r2p, &tmp_output_path) } + ExtractionJobType::GlobalStrings => self.extract_global_strings(r2p, &tmp_output_path), ExtractionJobType::FunctionZignatures => { - self.extract_function_zignatures(r2p, tmp_job_type_suffix) + self.extract_function_zignatures(r2p, &tmp_output_path) } ExtractionJobType::FunctionBytes => { - self.extract_function_bytes(r2p, tmp_job_type_suffix) + self.extract_function_bytes(r2p, &tmp_output_path, false) + } + ExtractionJobType::FunctionBytesMasked => { + self.extract_function_bytes(r2p, &tmp_output_path, true) } }?; // Apply final output file name when extraction is done - let output_path = self.get_output_filepath(&job_type_suffix); std::fs::rename(&tmp_output_path, &output_path).map_err(|e| { anyhow::anyhow!( "Failed to rename temporary path {:?}: {}", @@ -729,6 +1410,54 @@ impl FileToBeProcessed { Ok(()) } + pub fn write_function_index( + &self, + functions: &Vec, + output_dirpath: &PathBuf, + job_type: ExtractionJobType, + ) -> Result<()> { + let job_type_suffix = ExtractionJob::get_job_type_suffix(&job_type); + let ext = FunctionToBeProcessed::get_function_file_ext(&job_type); + let index_path = output_dirpath.join(format!("00-func-index_{}.csv", job_type_suffix)); + info!("Writing function index to {:?}", index_path); + let file = File::create(index_path)?; + let mut writer = csv::Writer::from_writer(file); + // Write header + writer.write_record(&[ + "name", + "address", + "size", + "ninstrs", + "nblocks", + "output_path", + ])?; + + // Write function records + for function in functions { + let output_path = + function.get_output_filepath(output_dirpath, &self.func_filename_template, ext); + writer.write_record(&[ + &function.name, + &function.addr.to_string(), + &function.size.to_string(), + &function.ninstrs.to_string(), + &function.nblocks.to_string(), + &output_path.to_string_lossy().to_string(), + ])?; + } + writer.flush()?; + Ok(()) + } + + /// Returns the function list for the current file + /// Uses a cached list if available, otherwise calls setup_function_list + /// NOTE: Add ExtractionJobType to job_type_needs_function_list so FuncInfo data can be added as cache + pub fn get_function_list(&self, r2p: &mut R2Pipe) -> Result<&[FunctionToBeProcessed]> { + self.function_list + .get_or_try_init(|| self.setup_function_list(r2p, None)) + .map(|v| v.as_slice()) + } + pub fn process_all_modes(&self) { info!( "Starting extraction for {} job types on {:?}", @@ -750,32 +1479,62 @@ impl FileToBeProcessed { info!("Processing job type: {:?}", job_type); // Check if the extracted data file already exists - let job_type_suffix = self.get_job_type_suffix(job_type); - let output_path = self.get_output_filepath(&job_type_suffix); - if Path::new(&output_path).exists() { - warn!( + let job_type_suffix = ExtractionJob::get_job_type_suffix(job_type); + let output_path = match self.get_output_filepath(&job_type_suffix) { + Ok(path) => path, + Err(e) => { + error!("Failed to get output filepath for {:?}: {}", job_type, e); + continue; + } + }; + let error_path = match self.get_error_filepath(&job_type_suffix) { + Ok(path) => path, + Err(e) => { + error!("Failed to get error filepath for {:?}: {}", job_type, e); + continue; + } + }; + + if output_path.exists() { + info!( "Skipping {:?} job for {:?}: already processed at {:?}.", job_type_suffix, self.file_path, output_path ); continue; + } else if error_path.exists() && !self.retry_aborted { + info!( + "Skipping {:?} job for {:?}: already processed and failed. Error log at {:?}.", + job_type_suffix, self.file_path, error_path + ); + continue; } // Lazily initialize r2p if not already done. - let r2p = maybe_r2p.get_or_insert_with(|| { - let mut pipe = self.setup_r2_pipe(); - self.analyse_r2_pipe(&mut pipe); - pipe - }); + maybe_r2p = match self.ensure_r2_pipe(maybe_r2p, 5) { + Ok(r2p) => Some(r2p), + Err(e) => { + error!("Failed to create R2Pipe for {:?}: {e}", self.file_path); + None + } + }; + if maybe_r2p.is_none() { + // We couldn't create a R2Pipe. We must give up on this file + break; + } + let r2p = maybe_r2p.as_mut().unwrap(); match self.process_mode(r2p, job_type) { Ok(_) => debug!( "Finished {:?} extraction job for {:?}: processed at {:?}.", job_type_suffix, self.file_path, output_path ), - Err(e) => error!( - "Aborted {:?} extraction job for {:?} due to error: {:?}.", - job_type_suffix, self.file_path, e - ), + Err(e) => { + error!( + "Aborted {:?} extraction job for {:?} due to error:\n{:?}\n\t> (stored in {:?}).", + job_type_suffix, self.file_path, e, error_path + ); + std::fs::write(error_path, format!("{}", e)).unwrap(); + } } } @@ -786,27 +1545,7 @@ impl FileToBeProcessed { } } - pub fn get_job_type_suffix(&self, job_type: &ExtractionJobType) -> String { - match job_type { - ExtractionJobType::BinInfo => "bininfo", - ExtractionJobType::RegisterBehaviour => "reg", - ExtractionJobType::FunctionXrefs => "func-xrefs", - ExtractionJobType::CFG => "cfg", - ExtractionJobType::CallGraphs => "cg", - ExtractionJobType::FuncInfo => "finfo", - ExtractionJobType::FunctionVariables => "fvars", - ExtractionJobType::Decompilation => "decomp", - ExtractionJobType::PCodeFunc => "pcode-func", - ExtractionJobType::PCodeBB => "pcode-bb", - ExtractionJobType::LocalVariableXrefs => "localvar-xrefs", - ExtractionJobType::GlobalStrings => "strings", - ExtractionJobType::FunctionZignatures => "zigs", - ExtractionJobType::FunctionBytes => "bytes", - } - .to_string() - } - - pub fn extract_binary_info(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { + pub fn extract_binary_info(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting binary information extraction"); let bininfo_json = r2p .cmd("ij") @@ -862,16 +1601,16 @@ impl FileToBeProcessed { info!("Binary information and checksums extracted."); info!("Writing extracted data to file"); - self.write_to_json(&json!(bininfo), job_type_suffix)?; + self.write_to_json(&json!(bininfo), output_path)?; Ok(()) } pub fn extract_register_behaviour( &self, r2p: &mut R2Pipe, - job_type_suffix: String, + output_path: &PathBuf, ) -> Result<()> { - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut register_behaviour_vec: HashMap = HashMap::new(); info!("Executing aeafj for each function"); for function in function_details.iter() { @@ -891,14 +1630,14 @@ impl FileToBeProcessed { } info!("All functions processed"); info!("Writing extracted data to file"); - self.write_to_json(&json!(register_behaviour_vec), job_type_suffix)?; + self.write_to_json(&json!(register_behaviour_vec), output_path)?; Ok(()) } pub fn extract_function_call_graphs( &self, r2p: &mut R2Pipe, - job_type_suffix: String, + output_path: &PathBuf, ) -> Result<()> { info!("Starting function call graph extraction"); let json = r2p @@ -908,28 +1647,45 @@ impl FileToBeProcessed { .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; info!("Function call graph extracted."); info!("Writing extracted data to file"); - self.write_to_json(&json!(function_call_graphs), job_type_suffix)?; + self.write_to_json(&json!(function_call_graphs), output_path)?; Ok(()) } - pub fn extract_function_info(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { - info!("Starting function metdata extraction"); - let function_details: Vec = self.get_function_name_list(r2p)?; + pub fn extract_function_info(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { + info!("Starting function info extraction"); + let json_raw = r2p + .cmd("aflj") + .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; + + info!("Writing function info to {:?}", output_path); + std::fs::write(output_path, &json_raw) + .with_context(|| format!("Failed to write JSON to {:?}", output_path))?; + info!("Function info written to {:?}", output_path); + + // Only populate function_list cache if other extraction modes need it + if self.needs_func_list { + // Parse the JSON and populate function_list cache for other extraction modes + // This avoids running aflj again later + let json: serde_json::Value = serde_json::from_str(&json_raw) + .with_context(|| format!("Failed to parse aflj output for {:?}", self.file_path))?; + + // Populate the function list cache using the existing logic in setup_function_list + if let Ok(func_list) = self.setup_function_list(r2p, Some(json)) { + // Ignore the result - if it's already set, that's fine + let _ = self.function_list.set(func_list); + } + } - info!("Writing extracted data to file"); - let json = json!(function_details); - self.write_to_json(&json, job_type_suffix) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; Ok(()) } pub fn extract_function_variables( &self, r2p: &mut R2Pipe, - job_type_suffix: String, + output_path: &PathBuf, ) -> Result<()> { info!("Starting function variables extraction"); - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut func_variables_vec: HashMap = HashMap::new(); info!("Executing aeafj for each function"); for function in function_details.iter() { @@ -947,87 +1703,294 @@ impl FileToBeProcessed { } info!("All functions processed"); info!("Writing extracted data to file"); - self.write_to_json(&json!(func_variables_vec), job_type_suffix)?; + self.write_to_json(&json!(func_variables_vec), output_path)?; Ok(()) } - pub fn extract_func_cfgs(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { - info!("Executing agfj @@f on {:?}", self.file_path); + /// Centralized function for extracting function-level data with resume capability. + /// Handles file scanning, skip logic, error logging, and aggregate statistics. + fn extract_functions_with_resume( + &self, + r2p: &mut R2Pipe, + output_dirpath: &PathBuf, + job_types: &[ExtractionJobType], + extraction_fn: F, + ) -> Result<()> + where + F: Fn(&FunctionToBeProcessed, &mut R2Pipe, usize) -> Result<()>, + { + let function_list = self.get_function_list(r2p)?.to_vec(); + let functions_count = function_list.len(); + + if functions_count == 0 { + warn!("No functions to process in {:?}", self.file_path); + return Ok(()); + } - let json_raw = r2p.cmd("agfj @@f").with_context(|| { - format!( - "Failed to extract control flow graph information from {:?}.", - self.file_path - ) - })?; + // Phase 1: Scan for existing files to enable resume functionality + struct FileStatus { + required_files_exist: Vec, // One per job_type + error_exists: bool, + } - info!("Starting JSON fixup for {:?}", self.file_path); - match self.fix_json_object(&json_raw) { - Ok(json) => { - info!("JSON fixup finished for {:?}", self.file_path); - // If the cleaned JSON is an empty array, log an error and skip. - if json == serde_json::Value::Array(vec![]) { - return Err(anyhow::anyhow!( - "File empty after JSON fixup - Only contains empty JSON array - {:?}", - self.file_path - )); - } else { - self.write_to_json(&json, job_type_suffix)?; + // Get extensions for all job types + let extensions: Vec<&str> = job_types + .iter() + .map(|jt| FunctionToBeProcessed::get_function_file_ext(jt)) + .collect(); + + let mut file_statuses: Vec = Vec::with_capacity(functions_count); + let mut last_file_index: Option = None; + let mut existing_complete_count = 0; + let mut skipped_error_count = 0; + + for (index, function) in function_list.iter().enumerate() { + let mut required_files_exist = Vec::with_capacity(job_types.len()); + + // Check if each required file exists + for ext in &extensions { + let file_path = + function.get_output_filepath(output_dirpath, &self.func_filename_template, ext); + required_files_exist.push(file_path.exists()); + } + + let error_path = function.get_output_filepath( + output_dirpath, + &self.func_filename_template, + "error.log", + ); + let error_exists = error_path.exists(); + + // Check if this function is "complete" (all required files exist) + let is_complete = required_files_exist.iter().all(|&exists| exists); + + if is_complete { + existing_complete_count += 1; + } + if error_exists && !self.retry_aborted { + skipped_error_count += 1; + } + + file_statuses.push(FileStatus { + required_files_exist, + error_exists, + }); + + if is_complete || error_exists { + last_file_index = Some(index); + } + } + + // Log resume information + if let Some(last_idx) = last_file_index { + info!( + "Found {} existing complete files. Resuming from function {} (regenerating last file). {} functions need processing.", + existing_complete_count, + last_idx, + functions_count - existing_complete_count + 1 - skipped_error_count + ); + if skipped_error_count > 0 { + info!( + "Skipping {} previously failed functions (retry_aborted=false)", + skipped_error_count + ); + } + } else { + info!( + "No existing files found, processing all {} functions", + functions_count + ); + } + + // Phase 2: Extract data for each function + let mut n_success: u32 = 0; + let mut n_skipped: u32 = 0; + let mut n_error: u32 = 0; + + for (index, function) in function_list.iter().enumerate() { + let status = &file_statuses[index]; + + // Skip logic: skip files before the last_file_index if they're already done or failed + if let Some(last_idx) = last_file_index { + if index < last_idx { + let is_complete = status.required_files_exist.iter().all(|&exists| exists); + + if is_complete { + debug!( + "Skipping function {:?} @ {:#x}: already extracted", + function.name, function.addr + ); + n_success += 1; + n_skipped += 1; + continue; + } + if status.error_exists && !self.retry_aborted { + debug!( + "Skipping function {:?} @ {:#x}: previously failed, retry_aborted=false", + function.name, function.addr + ); + n_skipped += 1; + continue; + } } } - Err(e) => { - return Err(anyhow::anyhow!( - "Unable to parse json for {:?}: {}: {}", - self.file_path, - json_raw, - e - )); + + // Call the extraction function + match extraction_fn(function, r2p, index) { + Ok(()) => { + n_success += 1; + } + Err(e) => { + let error_path = function.get_output_filepath( + output_dirpath, + &self.func_filename_template, + "error.log", + ); + error!( + "Failed to extract data for function {:?} @ {:#x}: {}", + function.name, function.addr, e + ); + if let Err(write_err) = std::fs::write(&error_path, e.to_string()) { + error!("Failed to write error to {:?}: {}", error_path, write_err); + } else { + debug!("Error stored at {:?}", error_path); + } + n_error += 1; + } + } + } + + // Final aggregate logging + info!( + "Extraction complete: {} successful, {} skipped, {} errors (out of {} functions)", + n_success, n_skipped, n_error, functions_count + ); + + // Return error if nothing succeeded + if n_success == 0 && n_skipped == 0 { + return Err(anyhow::anyhow!( + "Failed to extract data for all functions in {:?}", + self.file_path + )); + } + + Ok(()) + } + + pub fn extract_func_cfgs( + &self, + r2p: &mut R2Pipe, + output_path: &PathBuf, + split_by_function: bool, + ) -> Result<()> { + if !split_by_function { + // All CFGs stored in a single JSON file (potentially very large) + info!("Executing agfj @@f on {:?}", self.file_path); + + let json_raw = r2p + .cmd("agfj @@f") + .with_context(|| format!("Failed to extract CFGs from {:?}.", self.file_path))?; + + self.stream_write_to_json(&json_raw, output_path) + .with_context(|| format!("Failed to write CFGs to {:?}.", self.file_path))?; + } else { + // CFGs stored in a separate JSON file for each function + info!("Extracting CFGs for each function of {:?}", self.file_path); + + // The output path is a directory, create it if it doesn't exist + let output_dirpath = output_path.clone(); + if !output_dirpath.is_dir() { + std::fs::create_dir_all(&output_dirpath) + .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; } + + // Write index for the CFGs (need to get function_list first) + let function_list = self.get_function_list(r2p)?.to_vec(); + self.write_function_index( + &function_list, + &output_dirpath, + ExtractionJobType::FunctionCFG, + )?; + + // Use centralized extraction with resume capability + self.extract_functions_with_resume( + r2p, + &output_dirpath, + &[ExtractionJobType::FunctionCFG], + |function, r2p, _index| { + debug!( + "Extracting CFG for function {:?} @ {:#x}", + function.name, function.addr + ); + let result = function.write_cfg_to_json( + r2p, + &output_dirpath, + &self.func_filename_template, + ); + if result.is_ok() { + debug!( + "Successfully extracted CFG for function {:?} @ {:#x}", + function.name, function.addr + ); + } + result + }, + )?; + + info!( + "Finished extracting CFGs for each function of {:?}", + self.file_path + ); } Ok(()) } - pub fn extract_function_xrefs(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { - let function_details = self.get_function_name_list(r2p)?; + pub fn extract_function_xrefs(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { + let function_details = self.get_function_list(r2p)?; let mut function_xrefs: HashMap> = HashMap::new(); + info!("Extracting xrefs for each function"); - for function in function_details.iter() { - let ret = self - .get_function_xref_details(function.offset, r2p) - .with_context(|| { - format!("Unable to get function xrefs from {:?}", self.file_path) - })?; + for function in function_details { + let ret = function.get_xref_details(r2p).with_context(|| { + format!("Unable to get function xrefs from {:?}", self.file_path) + })?; function_xrefs.insert(function.name.clone(), ret); } - info!("All functions processed"); - info!("Writing extracted data to file"); - self.write_to_json(&json!(function_xrefs), job_type_suffix)?; + info!("All functions processed! Writing extracted data to file"); + self.write_to_json(&json!(function_xrefs), output_path)?; Ok(()) } - pub fn extract_decompilation(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { + pub fn extract_decompilation(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting decompilation extraction!"); - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut function_decomp: HashMap = HashMap::new(); - for function in function_details.iter() { - let ret = self.get_ghidra_decomp(function.offset, r2p); - function_decomp.insert(function.name.clone(), ret.unwrap()); + for function in function_details { + let ret = function + .get_ghidra_decomp(r2p, self.with_annotations) + .with_context(|| { + format!( + "Unable to get decompilation for {:?} @ {:?}", + self.file_path, function.addr + ) + })?; + function_decomp.insert(function.name.clone(), ret); } info!("Decompilation extracted successfully for all functions."); info!("Writing extracted data to file"); - self.write_to_json(&json!(function_decomp), job_type_suffix)?; + self.write_to_json(&json!(function_decomp), output_path)?; Ok(()) } - pub fn extract_pcode_function(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { + pub fn extract_pcode_function(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting pcode extraction at a function level"); - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut function_pcode = Vec::new(); - for function in function_details.iter() { - let ret = self.get_ghidra_pcode_function(function.offset, function.ninstrs, r2p); + for function in function_details { + let ret = self.get_ghidra_pcode(function.addr, function.ninstrs, r2p); let formatted_obj = PCodeJSONWithFuncName { function_name: function.name.clone(), @@ -1038,36 +2001,26 @@ impl FileToBeProcessed { } info!("Pcode extracted successfully for all functions."); info!("Writing extracted data to file"); - self.write_to_json(&json!(function_pcode), job_type_suffix)?; + self.write_to_json(&json!(function_pcode), output_path)?; Ok(()) } - pub fn extract_pcode_basic_block( - &self, - r2p: &mut R2Pipe, - job_type_suffix: String, - ) -> Result<()> { + pub fn extract_pcode_basic_block(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting pcode extraction for each basic block in each function within the binary"); - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut function_pcode = Vec::new(); - for function in function_details.iter() { - let bb_addresses = self - .get_basic_block_addresses(function.offset, r2p) - .with_context(|| { - format!( - "Unable to get basic block addresses in {:?} at offset {:?}", - self.file_path, function.offset - ) - })?; + for function in function_details { + let bb_info = function.get_basic_block_info(r2p).with_context(|| { + format!( + "Unable to get basic block addresses in {:?} @ {:?}", + self.file_path, function.addr + ) + })?; let mut bb_pcode: Vec = Vec::new(); - for bb in bb_addresses.iter() { + for bb in bb_info.iter() { let ret = self - .get_ghidra_pcode_function( - bb.addr, - bb.ninstr.try_into().unwrap(), // Convert u64 to i64 - r2p, - ) + .get_ghidra_pcode(bb.addr, bb.ninstr.try_into().unwrap(), r2p) .with_context(|| { format!( "Basic block decompilation failed in {:?} at offset {:?}", @@ -1090,48 +2043,47 @@ impl FileToBeProcessed { } info!("Pcode extracted successfully for all functions."); info!("Writing extracted data to file"); - self.write_to_json(&json!(function_pcode), job_type_suffix)?; + self.write_to_json(&json!(function_pcode), output_path)?; Ok(()) } pub fn extract_local_variable_xrefs( &self, r2p: &mut R2Pipe, - job_type_suffix: String, + output_path: &PathBuf, ) -> Result<()> { info!("Starting local variable xref extraction"); - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut function_local_variable_xrefs: HashMap = HashMap::new(); - for function in function_details.iter() { - let ret = self.get_local_variable_xref_details(function.offset, r2p)?; + for function in function_details { + let ret = function.get_local_variable_xref_details(r2p)?; function_local_variable_xrefs.insert(function.name.clone(), ret); } info!("Local variable xrefs extracted successfully for all functions."); info!("Writing extracted data to file"); - self.write_to_json(&json!(function_local_variable_xrefs), job_type_suffix)?; + self.write_to_json(&json!(function_local_variable_xrefs), output_path)?; Ok(()) } - pub fn extract_global_strings(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { + pub fn extract_global_strings(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting Global String Extraction"); let json = r2p .cmd("izj") .with_context(|| format!("Command izj failed in {:?}.", self.file_path))?; - debug!("{}", json); let json_obj: Vec = serde_json::from_str(&json) .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; - self.write_to_json(&json!(json_obj), job_type_suffix)?; + self.write_to_json(&json!(json_obj), output_path)?; Ok(()) } pub fn extract_function_zignatures( &self, r2p: &mut R2Pipe, - job_type_suffix: String, + output_path: &PathBuf, ) -> Result<()> { info!("Starting function zignatures extraction"); let _ = r2p @@ -1145,36 +2097,91 @@ impl FileToBeProcessed { .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; info!("Function zignatures extracted."); info!("Writing extracted data to file"); - self.write_to_json(&json!(function_zignatures), job_type_suffix)?; + self.write_to_json(&json!(function_zignatures), output_path)?; Ok(()) } - pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { + pub fn extract_function_bytes( + &self, + r2p: &mut R2Pipe, + output_dirpath: &PathBuf, + apply_mask: bool, + ) -> Result<()> { info!("Starting function bytes extraction"); - let function_details = self.get_function_name_list(r2p)?; - for function in function_details.iter() { - debug!( - "Function Name: {} Offset: {} Size: {}", - function.name, function.offset, function.size - ); - let function_bytes = self - .get_bytes_function(function.offset, function.size, r2p) - .with_context(|| { - format!( - "Bytes extraction failed in {:?} at function {:?}.", - self.file_path, function.name - ) - })?; - Self::write_to_bin( - self, - &function.name, - &function_bytes.bytes, - &job_type_suffix, - ) - .expect("Failed to write bytes to bin."); + let file_name = self.get_file_name()?; + let functions = self.get_function_list(r2p)?; + if !output_dirpath.is_dir() { + std::fs::create_dir_all(&output_dirpath) + .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; + } + + if !apply_mask || self.keep_raw_bytes { + // Write index for the raw bytes + self.write_function_index( + &functions.to_vec(), + &output_dirpath, + ExtractionJobType::FunctionBytes, + )?; + } + if apply_mask { + // Write index for the masked bytes + self.write_function_index( + &functions.to_vec(), + &output_dirpath, + ExtractionJobType::FunctionBytesMasked, + )?; } - info!("Function bytes successfully extracted"); + + // Determine which job types need to be extracted + let write_raw = !apply_mask || self.keep_raw_bytes; + let write_masked = apply_mask; + + let job_types: Vec = match (write_raw, write_masked) { + (true, true) => vec![ + ExtractionJobType::FunctionBytes, + ExtractionJobType::FunctionBytesMasked, + ], + (true, false) => vec![ExtractionJobType::FunctionBytes], + (false, true) => vec![ExtractionJobType::FunctionBytesMasked], + (false, false) => { + // This should never happen + warn!("Neither raw nor masked bytes requested for {:?}", file_name); + return Ok(()); + } + }; + + // Use centralized extraction with resume capability + self.extract_functions_with_resume( + r2p, + output_dirpath, + &job_types, + |function, r2p, _index| { + debug!( + "Extracting bytes for function {:?} @ {:?}, size: {}", + function.name, function.addr, function.size + ); + let result = function.write_to_bin( + r2p, + output_dirpath, + &self.func_filename_template, + apply_mask, + self.keep_raw_bytes, + ); + if result.is_ok() { + debug!( + "Successfully extracted bytes for function {:?} @ {:?}", + function.name, function.addr + ); + } + result + }, + )?; + + info!( + "Finished extracting bytes for each function of {:?}", + file_name + ); Ok(()) } @@ -1225,40 +2232,18 @@ impl FileToBeProcessed { } // r2 commands to structs - fn get_bytes_function( - &self, - function_addr: u64, - function_size: i128, - r2p: &mut R2Pipe, - ) -> Result { - Self::go_to_address(r2p, function_addr); - r2p.cmd(format!("s {}", function_addr).as_str())?; - let function_bytes = r2p.cmd(format!("p8 {}", function_size).as_str())?; - let function_bytes = function_bytes.trim(); - let function_bytes = hex::decode(function_bytes).map_err(|e| { - r2pipe::Error::Io(io::Error::new( - io::ErrorKind::InvalidData, - format!("Hex decode error: {}", e), - )) - })?; - - Ok(FuncBytes { - bytes: function_bytes, - }) - } - - fn get_ghidra_pcode_function( + fn get_ghidra_pcode( &self, - function_addr: u64, - num_instructons: i64, + address: u64, + num_instructons: u64, r2p: &mut R2Pipe, - ) -> Result { - Self::go_to_address(r2p, function_addr); - let pcode_ret = r2p.cmd(format!("pdgsd {}", num_instructons).as_str())?; + ) -> Result { + let cmd_str = format!("pdgsd {} @ {}", num_instructons, address); + debug!("Getting Ghidra PCode: `{}`", cmd_str); + let pcode_ret = r2p.cmd(&cmd_str).context("Failed to get Ghidra PCode")?; let lines = pcode_ret.lines(); let mut asm_ins = Vec::new(); let mut pcode_ins = Vec::new(); - for line in lines { if line.starts_with("0x") { asm_ins.push(line.trim().to_string()); @@ -1273,250 +2258,155 @@ impl FileToBeProcessed { }) } - fn get_ghidra_decomp( - &self, - function_addr: u64, - r2p: &mut R2Pipe, - ) -> Result { - Self::go_to_address(r2p, function_addr); - let json = r2p.cmd("pdgj")?; - - if self.with_annotations { - let json_obj: DecompJSON = serde_json::from_str(&json) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; - Ok(json_obj) - } else { - let json_obj: Value = serde_json::from_str(&json) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; - let parsed_code = json_obj["code"] - .as_str() - .with_context(|| format!("Unable to get code from {:?}!", json))? - .to_string(); - let parsed_obj = DecompJSON { - code: parsed_code, - annotations: Vec::new(), - }; - Ok(parsed_obj) - } - } - - fn get_function_name_list( + fn setup_function_list( &self, r2p: &mut R2Pipe, - ) -> Result, anyhow::Error> { - info!("Getting function information from binary"); - let json = r2p - .cmd("aflj") - .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; - - let json_obj: Vec = serde_json::from_str(json.as_ref()) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; - Ok(json_obj) - } - - fn get_basic_block_addresses( - &self, - function_addr: u64, - r2p: &mut R2Pipe, - ) -> Result { + aflj_json: Option, + ) -> Result> { info!( - "Getting the basic block information for function @ {}", - function_addr + "Setting up function list for {:?} ...", + self.get_file_name()? ); - Self::go_to_address(r2p, function_addr); - // Get basic block information - let json = r2p - .cmd("afbj") - .with_context(|| format!("Command afbj failed in {:?}", self.file_path))?; - // Parse the JSON into a mutable serde_json::Value. - let mut value: serde_json::Value = serde_json::from_str(&json) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; + // Determine the source of function list data + let json = match aflj_json { + // If JSON was provided, use it directly + Some(value) => value, - // Iterate over each object and convert "traced" from integer to boolean. - if let Some(array) = value.as_array_mut() { - for item in array.iter_mut() { - if let Some(traced_value) = item.get_mut("traced") { - // If traced is a number, convert it to a bool. - if let Some(num) = traced_value.as_i64() { - *traced_value = serde_json::Value::Bool(num != 0); - } + // Otherwise, try to load from existing finfo file or run aflj + None => { + let finfo_suffix = ExtractionJob::get_job_type_suffix(&ExtractionJobType::FuncInfo); + let finfo_path = self.get_output_filepath(&finfo_suffix).ok(); + + // Check if finfo file exists and use it + if let Some(path) = finfo_path.filter(|p| p.exists()) { + info!( + "Populating function list from existing finfo file at {:?}", + path + ); + let file = File::open(&path) + .with_context(|| format!("Failed to open finfo file: {:?}", path))?; + serde_json::from_reader(file).with_context(|| { + format!("Failed to parse finfo file as JSON: {:?}", path) + })? + } else { + // No existing finfo file, run aflj + info!("Populating function list from aflj"); + r2p.cmdj("aflj") + .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))? } } - } + }; - // Deserialize JSON into a BasicBlockInfo struct - let bb_addresses: BasicBlockInfo = - serde_json::from_value(value.clone()).with_context(|| { - format!( - "Unable to convert {:?} into a BasicBlockInfo struct!", - value - ) - })?; - Ok(bb_addresses) - } + let array = match json { + serde_json::Value::Array(arr) => arr, + _ => { + return Err(anyhow::anyhow!( + "aflj output is not an array for {:?}", + self.file_path + )) + } + }; - fn get_local_variable_xref_details( - &self, - function_addr: u64, - r2p: &mut R2Pipe, - ) -> Result { - info!("Getting local variable xref details"); - Self::go_to_address(r2p, function_addr); - let json = r2p.cmd("axvj"); - - // Convert returned JSON into a BasicBlockInfo struct - if let Ok(json_str) = json { - let local_variable_xrefs: LocalVariableXrefs = serde_json::from_str(json_str.as_ref()) - .expect("Unable to convert returned object into a BasicBlockInfo struct!"); - Ok(local_variable_xrefs) + if let Some(min_blocks) = self.min_basic_blocks { + info!( + "Filtering functions with less than {} basic blocks", + min_blocks + ); } else { - Err(json.unwrap_err()) + info!("No basic blocks filter applied"); } - } - - fn get_function_xref_details( - &self, - function_addr: u64, - r2p: &mut R2Pipe, - ) -> Result, anyhow::Error> { - info!("Getting function xref details"); - Self::go_to_address(r2p, function_addr); - let json = r2p - .cmd("axffj") - .with_context(|| format!("Command axffj failed in {:?}", self.file_path))?; - let mut json_obj: Vec = serde_json::from_str(&json) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; - debug!("Replacing all CALL xrefs with actual function name"); - // TODO: There is a minor bug in this where functions without any xrefs are included. - // Been left in as may be useful later down the line. - if !json_obj.is_empty() { - debug!("Replacing all CALL xrefs with actual function name"); - for element in json_obj.iter_mut() { - if element.type_field == "CALL" { - let cmd_str = format!("afi. @ {}", &element.ref_field); - let function_name = r2p.cmd(cmd_str.as_str()).with_context(|| { - format!("Command {:?} failed in {:?}", cmd_str, self.file_path) - })?; - element.name = function_name.trim().to_string(); - } - } - }; - Ok(json_obj) - } - // Helper Functions - fn fix_json_object(&self, json_raw: &str) -> Result { - // Collect all JSON objects into a vector. - let stream = Deserializer::from_str(json_raw).into_iter::(); - let json_objects: Result, _> = stream - .filter_map(|result| { + let functions_to_be_processed: Vec = array + .into_iter() + .map(|func_json| { + serde_json::from_value::(func_json) + .map(FunctionToBeProcessed::from) + }) + .filter(|result| { match result { - Ok(Value::Array(ref arr)) if arr.is_empty() => None, // skip empty arrays - other => Some(other), + Ok(func) => self + .min_basic_blocks + .map_or(true, |min_blocks| func.nblocks >= min_blocks as u64), + Err(_) => true, // Keep errors so they propagate through collect() } }) - .collect(); - // Map the collected vector into a JSON array. - json_objects.map(Value::Array) - } - - fn sanitize_function_name(&self, original: &str) -> String { - // Replace non-valid characters with '_' - // Valid characters: letters, digits, '_', '-', and '.' - let re = Regex::new(r"[^\w.-]").unwrap(); - re.replace_all(original, "_").into_owned() - } + .collect::, _>>() + .with_context(|| { + format!( + "Failed to parse aflj function details for {:?}", + self.file_path + ) + })?; + debug!( + "Done setting up function list for {:?}", + self.get_file_name()? + ); - fn get_file_name(&self) -> Result { - self.file_path - .file_name() - .ok_or_else(|| anyhow!("Unable to get file name from {:?}", self.file_path)) - .map(|os_str| os_str.to_string_lossy().to_string()) + Ok(functions_to_be_processed) } - fn write_to_json(&self, json_obj: &Value, job_type_suffix: String) -> Result<()> { - let mut fp_filename = self.get_file_name()?; - - fp_filename = if self.with_annotations { - fp_filename + "_" + &job_type_suffix + "_annotations" + ".json" - } else { - fp_filename + "_" + &job_type_suffix + ".json" - }; - - let mut output_filepath = PathBuf::new(); - output_filepath.push(self.output_path.clone()); - output_filepath.push(fp_filename); - debug!("Save filename: {:?}", output_filepath); + // Helper Functions + /// Write a JSON object to a file + fn write_to_json(&self, json_obj: &Value, output_filepath: &PathBuf) -> Result<()> { + info!("Writing JSON to {:?}", output_filepath); let file = File::create(&output_filepath) .with_context(|| format!("Unable to create file {:?}", output_filepath))?; - serde_json::to_writer(&file, &json_obj) + // Using a buffered writer to handle potentially huge JSON objects + let writer = BufWriter::new(file); + serde_json::to_writer(writer, &json_obj) .with_context(|| format!("Failed to write JSON to {:?}", output_filepath))?; - + info!("JSON written to {:?}", output_filepath); Ok(()) } - fn write_to_bin( - &self, - function_name: &String, - func_bytes: &[u8], - dirname_suffix: &String, - ) -> Result<()> { - // Extract the file stem from self.file_path and add context if missing. - let file_stem = self.get_file_name()?; - - // Construct the full output directory path. - let mut output_dir = self.output_path.clone(); - // Build the directory name by combining the file stem with the given suffix. - let dir_name = format!("{}_{}", file_stem, dirname_suffix); - output_dir.push(&dir_name); - - if !output_dir.is_dir() { - fs::create_dir_all(&output_dir) - .with_context(|| format!("Failed to create directory {:?}", output_dir))?; - } - - // Construct the full output file path. - let mut output_filepath = output_dir.clone(); - // Sanitize the function name to create a valid filename. - let sanitized_function_name = self.sanitize_function_name(function_name); - - // If sanitized_function_name larger than 100 characters, snip - let sanitized_function_name = if sanitized_function_name.len() > 100 { - String::from(&sanitized_function_name[0..75]) - } else { - sanitized_function_name - }; - - let file_name = format!("{}.bin", sanitized_function_name); - output_filepath.push(file_name); + /// Stream-parse a concatenated JSON string (e.g. from `agfj @@f`) + /// and write the combined results into a single JSON array on disk. + /// Unlike `write_to_json`, this never builds a full in-memory Vec. + fn stream_write_to_json(&self, json_raw: &str, output_filepath: &PathBuf) -> Result<()> { + info!("Stream writing JSON to {:?}", output_filepath); + // Stream the concatenated JSON values + let iter = Deserializer::from_str(json_raw).into_iter::>(); - // Check if a file with same name the sanitized name already exists. - if output_filepath.exists() { - debug!( - "Duplicate function binary file detected for '{}' at {:?}. Skipping.", - function_name, output_filepath - ); - return Ok(()); + // Stream the final JSON array to disk + let file = File::create(output_filepath) + .with_context(|| format!("Unable to create file {:?}", output_filepath))?; + let mut writer = BufWriter::new(file); + let mut ser = serde_json::Serializer::new(&mut writer); + let mut seq = ser.serialize_seq(None)?; // unknown length + + for item in iter { + match item { + Ok(raw) => { + // Skip exactly "[]" + if raw.get().trim() != "[]" { + // RawValue writes the raw JSON without re-encoding + debug!("Serializing JSON chunk"); + seq.serialize_element(&raw)?; + } else { + debug!("Skipping empty array"); + } + } + Err(e) => { + // Only tolerate a *trailing* truncation + if e.is_eof() { + warn!("Trailing truncation detected. Stopping stream"); + break; + } else { + error!("Malformed JSON chunk: {e}"); + return Err(e).context("Malformed JSON chunk"); + } + } + } } - debug!( - "Attempting to write function bytes to {:?}", - output_filepath - ); - // Write the file and attach context on error. - fs::write(&output_filepath, func_bytes) - .with_context(|| format!("Failed to write file {:?}", output_filepath))?; - + seq.end()?; + info!("JSON stream written to {:?}", output_filepath); Ok(()) } - fn go_to_address(r2p: &mut R2Pipe, function_addr: u64) { - r2p.cmd(format!("s {}", function_addr).as_str()) - .expect("failed to seek addr"); - } - + /// Downloads the PDB file if available fn handle_symbols_pdb(&self, r2p: &mut R2Pipe) -> Result<(), Error> { // Download symbols if available debug!("Downloading pdb file for {:?}", self.file_path); @@ -1534,43 +2424,70 @@ impl FileToBeProcessed { } } - fn setup_r2_pipe(&self) -> R2Pipe { + fn setup_r2_pipe(&self) -> Result { if self.r2p_config.use_curl_pdb { // Docs suggest this is unsafe env::set_var("R2_CURL", "1"); } - let opts = if self.r2p_config.debug { + let mut args = vec![]; + + // Set debug level and verbosity + if self.r2p_config.debug { debug!("Creating r2 handle with debugging"); - R2PipeSpawnOptions { - exepath: "radare2".to_owned(), - args: vec!["-e bin.cache=true", "-e log.level=0", "-e asm.pseudo=true"], - } + args.push("-e log.level=0"); } else { debug!("Creating r2 handle without debugging"); - R2PipeSpawnOptions { - exepath: "radare2".to_owned(), - args: vec![ - "-e bin.cache=true", - "-e log.level=1", - "-2", - "-e asm.pseudo=true", - ], - } + args.extend(["-e log.level=1", "-2"]); + } + + // Choose between cache or relocations (default is cache) + if self.r2p_config.apply_relocations { + args.push("-e bin.relocs.apply=true"); + } else { + args.push("-e bin.cache=true"); + } + + // Set pseudo-disassembly (default is true) + if self.r2p_config.disable_pseudo_asm { + args.push("-e asm.pseudo=false"); + } else { + args.push("-e asm.pseudo=true"); + } + + // Set timeout if any + if let Some(timeout) = self.r2p_config.timeout { + // Convert timeout from seconds to milliseconds + let timeout_ms = timeout * 1000; + let owned_arg_str = format!("-e anal.timeout={timeout_ms}"); + + // Leak the owned String to get a &'static str so the compiler doesn't complain + let arg_str: &'static str = Box::leak(owned_arg_str.into_boxed_str()); + args.push(arg_str); + } + + let opts = R2PipeSpawnOptions { + exepath: "radare2".to_owned(), + args, }; - debug!("Attempting to create r2pipe using {:?}", self.file_path); + debug!( + "Attempting to create r2pipe to analyze {:?} ...", + self.file_path + ); + debug!( + "-> {:?} {:?} {:?}", + opts.exepath, + opts.args.join(" "), + self.file_path + ); let mut r2p = match R2Pipe::in_session() { - Some(_) => R2Pipe::open().expect("Unable to open R2Pipe"), + Some(_) => R2Pipe::open() + .with_context(|| format!("Unable to open R2Pipe for {:?}", self.file_path))?, None => R2Pipe::spawn(self.file_path.to_str().unwrap(), Some(opts)) - .expect("Failed to spawn new R2Pipe"), + .with_context(|| format!("Failed to spawn new R2Pipe for {:?}", self.file_path))?, }; - if let Some(timeout) = self.r2p_config.timeout { - r2p.cmd(format!("e anal.timeout={}", timeout).as_str()) - .expect("Failed to set timeout"); - } - if self.r2p_config.use_curl_pdb { let info = r2p.cmdj("ij"); if info.is_ok() { @@ -1580,32 +2497,78 @@ impl FileToBeProcessed { let ret = self.handle_symbols_pdb(&mut r2p); if ret.is_err() { - error!("Unable to get PDB info") + warn!("Unable to get PDB info for {:?}", self.file_path); + info!("Continuing with analysis..."); + } else { + info!("PDB info obtained successfully for {:?}", self.file_path); } } } } - r2p + Ok(r2p) } - fn analyse_r2_pipe(&self, r2p: &mut R2Pipe) { - if self.r2p_config.extended_analysis { - debug!( - "Executing 'aaa' r2 command for {}", - self.file_path.display() - ); - r2p.cmd("aaa") - .expect("Unable to complete standard analysis!"); - debug!("'aaa' r2 command complete for {}", self.file_path.display()); - } else { - debug!("Executing 'aa' r2 command for {}", self.file_path.display()); - r2p.cmd("aa") - .expect("Unable to complete standard analysis!"); - debug!( - "'aa' r2 command complete for {:?}", - self.file_path.display() - ); - }; + fn analyse_r2_pipe(&self, r2p: &mut R2Pipe) -> Result<()> { + let analysis_mode = self.r2p_config.analysis_mode.as_str(); + debug!( + "Executing '{}' r2 command for {}", + analysis_mode, + self.file_path.display() + ); + r2p.cmd(analysis_mode) + .with_context(|| format!("Unable to complete analysis! ({analysis_mode})"))?; + debug!( + "'{}' r2 command complete for {}", + analysis_mode, + self.file_path.display() + ); + Ok(()) + } + + fn is_r2_pipe_alive(&self, r2p: &mut R2Pipe) -> bool { + // Check if we can get the version + match r2p.cmd("?V") { + Ok(_) => true, // The R2Pipe is alive + Err(_) => false, // The R2Pipe is dead + } + } + + fn ensure_r2_pipe(&self, maybe_r2p: Option, max_attempts: u16) -> Result { + match maybe_r2p { + Some(r2p) => { + let mut mutable_r2p = r2p; + if self.is_r2_pipe_alive(&mut mutable_r2p) { + return Ok(mutable_r2p); + } else { + return self.ensure_r2_pipe(None, max_attempts); + } + } + None => { + let mut n_attempts: u16 = 1; + while n_attempts <= max_attempts { + info!("Creating R2Pipe (attempt {n_attempts}/{max_attempts})..."); + match self.setup_r2_pipe() { + Ok(r2p) => { + let mut mutable_r2p = r2p; + self.analyse_r2_pipe(&mut mutable_r2p)?; + if self.is_r2_pipe_alive(&mut mutable_r2p) { + return Ok(mutable_r2p); + } + } + Err(e) => { + warn!("Failed to create R2Pipe: {e}"); + n_attempts += 1; + continue; + } + } + } + } + } + + bail!( + "Failed to create R2Pipe after {} attempts. Giving up.", + max_attempts + ); } } diff --git a/src/main.rs b/src/main.rs index c30b533..ef43464 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,7 +8,8 @@ use std::fmt; extern crate log; use clap::builder::TypedValueParser; use env_logger::Env; -use indicatif::{ParallelProgressIterator, ProgressIterator}; +use indicatif::MultiProgress; +use indicatif_log_bridge::LogWrapper; use glob::glob; use mimalloc::MiMalloc; @@ -16,6 +17,7 @@ use rayon::iter::ParallelIterator; use rayon::prelude::IntoParallelRefIterator; use std::path::{Path, PathBuf}; use std::process::exit; +use std::sync::Arc; use walkdir::WalkDir; pub mod afij; @@ -59,11 +61,30 @@ use inference::inference; #[cfg(feature = "inference")] use processors::agfj_graph_embedded_feats; use processors::agfj_graph_statistical_features; -use utils::get_json_paths_from_dir; +use utils::{get_json_paths_from_dir, validate_func_filename}; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; +// Global multi-progress instance for coordinating progress bars +static MULTI_PROGRESS: std::sync::OnceLock> = std::sync::OnceLock::new(); + +/// Get the global multi-progress instance +fn get_multi_progress() -> Arc { + MULTI_PROGRESS + .get() + .expect("Multi-progress not initialized") + .clone() +} + +/// Create a progress bar that works with the log bridge +fn create_progress_bar(len: u64) -> indicatif::ProgressBar { + let multi = get_multi_progress(); + let pb = indicatif::ProgressBar::new(len); + multi.add(pb.clone()); + pb +} + #[derive(PartialEq, Copy, Clone)] pub enum DataType { Cfg, @@ -93,6 +114,10 @@ impl fmt::Display for DataType { #[derive(Parser)] #[command(author, version, about, long_about = None)] struct Cli { + /// Set the logging level + #[arg(short, long, value_name = "LEVEL", default_value = "warn", value_parser = clap::builder::PossibleValuesParser::new(["off", "error", "warn", "info", "debug", "trace"]))] + log_level: String, + #[command(subcommand)] command: Commands, } @@ -283,8 +308,8 @@ enum Commands { /// The extraction modes (multiple can be specified) #[arg(short, long, value_name = "EXTRACT_MODE", value_parser = clap::builder::PossibleValuesParser::new([ - "bininfo", "finfo", "fvars", "reg", "cfg", "func-xrefs", "cg", "decomp", - "pcode-func", "pcode-bb", "localvar-xrefs", "strings", "bytes", "zigs" + "bininfo", "finfo", "fvars", "reg", "cfg", "func-cfg", "func-xrefs", "cg", "decomp", + "pcode-func", "pcode-bb", "localvar-xrefs", "strings", "bytes", "bytes-masked", "zigs" ]) .map(|s| s.parse::().unwrap()), num_args = 1.., @@ -301,14 +326,47 @@ enum Commands { #[arg(long, default_value = "false")] extended_analysis: bool, + #[arg(long, default_value = "false")] + experimental_analysis: bool, + #[arg(long, default_value = "false")] use_curl_pdb: bool, + /// Toggle to apply relocations instead of caching (slower but more accurate) + #[arg(long, default_value = "false")] + apply_relocations: bool, // Open r2 with bin.relocs.apply=true + // Otherwise defaults to bin.cache=true + /// Toggle to disable pseudo-disassembly (faster but more variable in cross-platform scenarios) + #[arg(long, default_value = "false")] + disable_pseudo_asm: bool, // Open r2 with asm.pseudo=false + // Otherwise defaults to asm.pseudo=true + /// Name function data files using symbol, address, or custom template + /// e.g. '{address}-{symbol}.{ext}' + #[arg( + long, + value_name = "TEMPLATE", + default_value = "symbol", + value_parser = validate_func_filename + )] + func_filename: String, + + /// Timeout for Radare2 analysis in seconds #[arg(long)] timeout: Option, #[arg(long, default_value = "false")] with_annotations: bool, + + #[arg(long, default_value = "false")] + keep_raw_bytes: bool, + + /// Toggle to retry previously aborted jobs due to extraction failures + #[arg(long, default_value = "false")] + retry_aborted: bool, + + /// Filter functions with less than this number of basic blocks (disabled by default) + #[arg(long)] + min_basic_blocks: Option, }, /// Generate single embeddings on the fly /// @@ -398,12 +456,33 @@ enum DedupSubCommands { } fn main() { - let env = Env::default() - .filter_or("LOG_LEVEL", "warn") - .write_style_or("LOG_STYLE", "always"); - - env_logger::init_from_env(env); let cli = Cli::parse(); + + // Set up the multi-progress for coordinating progress bars + let multi = MultiProgress::new(); + let multi_arc = Arc::new(multi); + + // Store the multi-progress instance globally + MULTI_PROGRESS + .set(multi_arc.clone()) + .expect("Failed to set global multi-progress"); + + // Build the logger with the specified log level + let logger = env_logger::Builder::from_env( + Env::default() + .filter_or("LOG_LEVEL", &cli.log_level) + .write_style_or("LOG_STYLE", "always"), + ) + .build(); + + let level = logger.filter(); + + // Set up the log bridge to prevent log messages from breaking progress bars + LogWrapper::new(multi_arc.as_ref().clone(), logger) + .try_init() + .expect("Failed to initialize log wrapper"); + + log::set_max_level(level); match &cli.command { #[cfg(feature = "goblin")] Commands::Info { path } => { @@ -666,7 +745,8 @@ fn main() { // if without metadata if !with_features & metadata_type.is_none() { debug!("Creating call graphs without any node features"); - file_paths_vec.par_iter().progress().for_each(|path| { + let pb = create_progress_bar(file_paths_vec.len() as u64); + file_paths_vec.par_iter().for_each(|path| { let suffix = graph_type.to_owned().to_string(); let full_output_path = get_save_file_path( &PathBuf::from(path), @@ -697,7 +777,9 @@ fn main() { full_output_path.to_string_lossy() ) } - }) + pb.inc(1); + }); + pb.finish(); } else { info!("Creating call graphs with node features"); debug!("Getting metadata file paths"); @@ -726,8 +808,10 @@ fn main() { .zip(metadata_paths_vec) .collect::>(); - combined_cgs_metadata.par_iter().progress().for_each( - |(filepath, metapath)| { + let pb = create_progress_bar(combined_cgs_metadata.len() as u64); + combined_cgs_metadata + .par_iter() + .for_each(|(filepath, metapath)| { let suffix = format!("{}-meta", graph_type.to_owned()); let full_output_path = get_save_file_path( &PathBuf::from(filepath), @@ -813,8 +897,9 @@ fn main() { full_output_path.to_string_lossy() ) } - }, - ); + pb.inc(1); + }); + pb.finish(); } } } @@ -1007,7 +1092,8 @@ fn main() { "{} files found. Beginning Processing.", file_paths_vec.len() ); - for file in file_paths_vec.iter().progress() { + let pb = create_progress_bar(file_paths_vec.len() as u64); + for file in file_paths_vec.iter() { let file = AGFJFile { functions: None, filename: PathBuf::from(file), @@ -1022,8 +1108,10 @@ fn main() { instruction_type, random_walk, *pairs, - ) + ); + pb.inc(1); } + pb.finish(); } } GenerateSubCommands::Tokeniser { @@ -1054,9 +1142,16 @@ fn main() { num_threads, debug, extended_analysis, + experimental_analysis, use_curl_pdb, + apply_relocations, + disable_pseudo_asm, + func_filename, timeout, with_annotations, + keep_raw_bytes, + retry_aborted, + min_basic_blocks, } => { info!("Creating extraction job with {} modes", modes.len()); if !output_dir.exists() { @@ -1064,16 +1159,36 @@ fn main() { exit(1) } + let analysis_mode; + // Make sure that extended analysis and experimental analysis are not both true + if *extended_analysis && *experimental_analysis { + error!("Extended analysis and experimental analysis cannot be both true"); + exit(1) + } else if *extended_analysis { + analysis_mode = "aaa"; + } else if *experimental_analysis { + analysis_mode = "aaaa"; + } else { + // Neither are true, so use the default analysis mode + analysis_mode = "aa"; + } + // Create a single extraction job with all modes let job = ExtractionJob::new( fpath, output_dir, modes, debug, - extended_analysis, + analysis_mode, use_curl_pdb, + apply_relocations, + disable_pseudo_asm, + func_filename, timeout, with_annotations, + keep_raw_bytes, + retry_aborted, + min_basic_blocks, ) .unwrap_or_else(|e| { error!("Failed to create extraction job: {}", e); @@ -1094,11 +1209,13 @@ fn main() { .build_global() .unwrap(); + let pb = create_progress_bar(job.files_to_be_processed.len() as u64); // Process all files in parallel, each file processes all modes with a single r2pipe - job.files_to_be_processed - .par_iter() - .progress() - .for_each(|path| path.process_all_modes()); + job.files_to_be_processed.par_iter().for_each(|path| { + path.process_all_modes(); + pb.inc(1); + }); + pb.finish(); } else if job.input_path_type == PathType::File { info!("Single file found"); @@ -1173,9 +1290,12 @@ fn main() { warn!("This only supports the Cisco Talos Binary Sim Dataset naming convention"); let corpus = EsilFuncStringCorpus::new(filename, output_path).unwrap(); - corpus.uniq_binaries.par_iter().progress().for_each(|name| { - corpus.dedup_subset(name, *print_stats, *just_stats, *just_hash_value) + let pb = create_progress_bar(corpus.uniq_binaries.len() as u64); + corpus.uniq_binaries.par_iter().for_each(|name| { + corpus.dedup_subset(name, *print_stats, *just_stats, *just_hash_value); + pb.inc(1); }); + pb.finish(); } }, } diff --git a/src/utils.rs b/src/utils.rs index 3ce29fc..9a5bfe8 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,3 +1,4 @@ +use regex::Regex; use std::fs::create_dir_all; use std::path::{Path, PathBuf}; use walkdir::WalkDir; @@ -74,6 +75,40 @@ pub fn get_save_file_path( } } +/// Validate the function filename template +/// +/// This function validates the function filename template to ensure it is a valid template. +/// The template must be a string containing `{symbol}` or `{address}`. +pub fn validate_func_filename(s: &str) -> Result { + if s == "symbol" || s == "address" { + Ok(s.to_string()) + } else if s.contains("{symbol}") || s.contains("{address}") { + Ok(s.to_string()) + } else { + Err( + "must be `symbol`, `address`, or a pattern containing `{symbol}`/`{address}`" + .to_string(), + ) + } +} + +/// Sanitize name for an output file +/// +/// This function sanitizes a file name by replacing non-valid characters with '_'. +pub fn sanitize_filename(name: &str) -> String { + // Replace non-valid characters with '_' + // Valid characters: letters, digits, '_', '-', and '.' + let re = Regex::new(r"[^\w.-]").unwrap(); + let sanitized_name = re.replace_all(name, "_").into_owned(); + + // This is a pretty dirty fix and may break things + if sanitized_name.len() > 100 { + sanitized_name[..50].to_string() + &sanitized_name[sanitized_name.len() - 50..] + } else { + sanitized_name + } +} + /// Get the JSON paths from a directory /// /// This function takes a path to a directory and traverses all @@ -129,6 +164,15 @@ pub fn parse_hex_escapes(s: String) -> Vec { bytes } +/// Function to return a default address value +/// +/// This function returns a default address of u64::MAX +/// This means that the address is not set as it is an impossible address to +/// reach, unless the binary is millions of terabytes. +pub fn get_default_addr() -> u64 { + u64::MAX +} + #[cfg(test)] mod tests { use super::*;