From c4fb537f376deaa7819f5aae54cce9be3a51df24 Mon Sep 17 00:00:00 2001 From: valbucci Date: Wed, 11 Jun 2025 10:51:29 +0100 Subject: [PATCH 01/42] Function attribute data types fom i64 to u64 --- src/afij.rs | 42 +++++++++++++++++++++--------------------- src/combos.rs | 10 +++++----- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/afij.rs b/src/afij.rs index eb43764..d2266c8 100644 --- a/src/afij.rs +++ b/src/afij.rs @@ -8,7 +8,7 @@ use serde_json::Value; pub struct AFIJFunctionInfo { 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 +23,21 @@ 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, pub minbound: u64, - pub maxbound: i128, + 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 +103,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 +130,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/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, From 6e5ae8d372a28591c6233b91948ec29d8f357dbd Mon Sep 17 00:00:00 2001 From: valbucci Date: Wed, 11 Jun 2025 10:53:09 +0100 Subject: [PATCH 02/42] Added option for custom function file naming --- src/extract.rs | 839 ++++++++++++++++++++++++------------------------- src/main.rs | 13 +- src/utils.rs | 14 + 3 files changed, 440 insertions(+), 426 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 8e9e668..4e80763 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1,8 +1,6 @@ use crate::afij::AFIJFunctionInfo; use crate::agcj::AGCJFunctionCallGraph; -use std::io; - use anyhow::anyhow; use anyhow::bail; use anyhow::Context; @@ -24,6 +22,7 @@ use sha1::Digest as Sha1Digest; use sha2::Digest as Sha2Digest; use sha2::Sha256; use std::collections::HashMap; +use std::sync::LazyLock; use std::env; use std::fs; use std::fs::File; @@ -39,7 +38,8 @@ 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, @@ -57,6 +57,29 @@ pub enum ExtractionJobType { 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::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::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 +87,15 @@ pub struct FileToBeProcessed { pub job_types: Vec, pub r2p_config: R2PipeConfig, pub with_annotations: bool, + pub func_filename_template: String, +} + +#[derive(Debug)] +pub struct FunctionToBeProcessed { + pub name: String, + pub addr: u64, + pub size: u64, + pub ninstrs: u64, } #[derive(Debug)] @@ -233,9 +265,9 @@ impl std::fmt::Display for AFLJFuncDetails { } } -impl From<(String, String, Vec, R2PipeConfig, bool)> for FileToBeProcessed { +impl From<(String, String, Vec, R2PipeConfig, bool, String)> for FileToBeProcessed { fn from( - orig: (String, String, Vec, R2PipeConfig, bool), + orig: (String, String, Vec, R2PipeConfig, bool, String), ) -> FileToBeProcessed { FileToBeProcessed { file_path: PathBuf::from(orig.0), @@ -243,6 +275,40 @@ impl From<(String, String, Vec, R2PipeConfig, bool)> for File job_types: orig.2, r2p_config: orig.3, with_annotations: orig.4, + func_filename_template: orig.5, + } + } +} + +impl From<(String, u64, u64, u64)> for FunctionToBeProcessed { + fn from(orig: (String, u64, u64, u64)) -> Self { + FunctionToBeProcessed { + name: orig.0, + addr: orig.1, + size: orig.2, + ninstrs: orig.3, + } + } +} + +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 as u64, + } + } +} + +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, } } } @@ -470,6 +536,7 @@ pub struct BinaryInfo { pub checksums: Option, // Populated manually with itj. } + impl ExtractionJob { pub fn new( input_path: &PathBuf, @@ -478,52 +545,16 @@ impl ExtractionJob { debug: &bool, extended_analysis: &bool, use_curl_pdb: &bool, + func_filename_template: &str, timeout: &Option, with_annotations: &bool, ) -> 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 fpath_md = fs::metadata(bin_path).unwrap(); - if fpath_md.is_file() { - PathType::File - } else if fpath_md.is_dir() { - PathType::Dir - } else { - PathType::Unk - } - } - - // 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), - } - } let mut job_types = vec![]; let mut extraction_job_types = vec![]; for mode in modes { - let job_type = extraction_job_matcher(mode)?; + let job_type = Self::extraction_job_matcher(mode)?; job_types.push((job_type, mode.clone())); extraction_job_types.push(job_type); // Store just the job type @@ -542,78 +573,104 @@ impl ExtractionJob { 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, - }; + let p_type = Self::get_path_type(input_path); - 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), + 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, + func_filename_template: func_filename_template.to_string(), + }; + + 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, + with_annotations: *with_annotations, + func_filename_template: func_filename_template.to_string(), + }) + .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 { + ExtractionJobType::FunctionBytes => None, // Output is a directory + _ => Some("json"), } } + /// 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 +688,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 +699,250 @@ impl ExtractionJob { paths.push(entry.to_string_lossy().to_string()); } } - paths } } -impl FileToBeProcessed { - pub fn get_output_filename(&self, job_type_suffix: &str) -> String { - let mut fp_filename = self - .file_path - .file_name() - .expect("Unable to get filename") - .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" +impl FunctionToBeProcessed { + + // Helper functions + fn sanitize_name(name: &str) -> String { + // Replace non-valid characters with '_' + // Valid characters: letters, digits, '_', '-', and '.' + let re = Regex::new(r"[^\w.-]").unwrap(); + re.replace_all(name, "_").into_owned() + } + + fn shorten_name(name: &str) -> String { + if name.len() > 100 { + name[..50].to_string() + &name[name.len() - 50..] } else { - fp_filename + "_" + job_type_suffix + ".json" + name.to_string() + } + } + + fn write_to_bin(&self, r2p: &mut R2Pipe, output_dirpath: &PathBuf, filename_template: &str) -> Result<()> { + let func_bytes = self.get_bytes(r2p)?; + let output_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); + std::fs::write(&output_filepath, func_bytes.bytes).with_context(|| format!("Failed to write function bytes to file: {:?}", output_filepath))?; + Ok(()) + } + + // Getters + fn get_hex_address(&self) -> String { + format!("{:x}", self.addr).trim_start_matches("0x").to_string() + } + + fn get_func_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), }; - fp_filename + + func_filename = Self::sanitize_name(&func_filename); + func_filename = Self::shorten_name(&func_filename); + if ["symbol", "address"].contains(&template) { + // Add an extension only if the user did not specify a custom template + func_filename = func_filename + "." + ext; + } + func_filename } - pub fn get_output_filepath(&self, job_type_suffix: &str) -> PathBuf { - let fp_filename = self.get_output_filename(job_type_suffix); + pub fn get_output_filepath(&self, output_dirpath: &PathBuf, filename_template: &str, ext: &str) -> PathBuf { + let output_filename = self.get_func_filename(filename_template, ext); let mut output_filepath = PathBuf::new(); - output_filepath.push(self.output_path.clone()); - output_filepath.push(fp_filename); + output_filepath.push(output_dirpath); + output_filepath.push(output_filename); + output_filepath + } + + fn get_bytes(&self, r2p: &mut R2Pipe) -> Result { + FileToBeProcessed::go_to_address(r2p, self.addr)?; + let mut function_bytes = r2p.cmd(format!("p8 {}", self.size).as_str())?; + function_bytes = function_bytes.trim().to_string(); + let decoded_bytes = hex::decode(&function_bytes) + .context("Failed to decode hex bytes")?; + + Ok(FuncBytes { + bytes: decoded_bytes + }) + } + + fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { + info!("Getting the basic block information for function @ {}", self.addr); + FileToBeProcessed::go_to_address(r2p, self.addr)?; + + let json = r2p.cmd("afbj").context("Command afbj failed")?; + // 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))?; + + // 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); + } + } + } + } + + // Deserialize JSON into a BasicBlockInfo struct + let bb_info: BasicBlockInfo = + serde_json::from_value(value.clone()).with_context(|| { + format!( + "Unable to convert {:?} into a BasicBlockInfo struct!", + value + ) + })?; + Ok(bb_info) + } + + fn get_local_variable_xref_details(&self, r2p: &mut R2Pipe) -> Result { + info!("Getting local variable xref details for function @ {}", self.addr); + FileToBeProcessed::go_to_address(r2p, self.addr)?; + let json = r2p.cmd("axvj").context("Command axvj failed")?; + + let local_variable_xrefs: LocalVariableXrefs = serde_json::from_str(&json) + .with_context(|| format!("Unable to convert {:?} to LocalVariableXrefs struct!", json))?; + Ok(local_variable_xrefs) + } + + fn get_xref_details(&self, r2p: &mut R2Pipe) -> Result, Error> { + info!("Getting xref details for function @ {}", self.addr); + FileToBeProcessed::go_to_address(r2p, self.addr)?; + let json = r2p.cmd("axffj").context("Command axffj failed")?; + let mut json_obj: Vec = serde_json::from_str(&json) + .with_context(|| format!("Unable to convert {:?} to FunctionXrefDetails struct!", json))?; + + // 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 { + FileToBeProcessed::go_to_address(r2p, self.addr)?; + let json = r2p.cmd("pdgj").context("Command pdgj failed")?; + + if with_annotations { + let json_obj: DecompJSON = serde_json::from_str(&json) + .with_context(|| format!("Unable to convert {:?} to DecompJSON struct!", 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) + } + } +} + +impl FileToBeProcessed { + /// Returns the name of the input binary file + 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()) + } + + /// 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 { + output_filename = output_filename + "_" + job_type_suffix + &ext_str; + }; - output_filepath.clone() + Ok(output_filename) + } + + 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) + } + + 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, job_type_suffix), ExtractionJobType::RegisterBehaviour => { - self.extract_register_behaviour(r2p, tmp_job_type_suffix) + self.extract_register_behaviour(r2p, job_type_suffix) } ExtractionJobType::FunctionXrefs => { - self.extract_function_xrefs(r2p, tmp_job_type_suffix) + self.extract_function_xrefs(r2p, job_type_suffix) } - ExtractionJobType::CFG => self.extract_func_cfgs(r2p, tmp_job_type_suffix), + ExtractionJobType::CFG => self.extract_func_cfgs(r2p, job_type_suffix), ExtractionJobType::CallGraphs => { - self.extract_function_call_graphs(r2p, tmp_job_type_suffix) + self.extract_function_call_graphs(r2p, job_type_suffix) } - ExtractionJobType::FuncInfo => self.extract_function_info(r2p, tmp_job_type_suffix), + ExtractionJobType::FuncInfo => self.extract_function_info(r2p, job_type_suffix), ExtractionJobType::FunctionVariables => { - self.extract_function_variables(r2p, tmp_job_type_suffix) + self.extract_function_variables(r2p, job_type_suffix) } ExtractionJobType::Decompilation => { - self.extract_decompilation(r2p, tmp_job_type_suffix) + self.extract_decompilation(r2p, job_type_suffix) } - ExtractionJobType::PCodeFunc => self.extract_pcode_function(r2p, tmp_job_type_suffix), - ExtractionJobType::PCodeBB => self.extract_pcode_basic_block(r2p, tmp_job_type_suffix), + ExtractionJobType::PCodeFunc => self.extract_pcode_function(r2p, job_type_suffix), + ExtractionJobType::PCodeBB => self.extract_pcode_basic_block(r2p, job_type_suffix), ExtractionJobType::LocalVariableXrefs => { - self.extract_local_variable_xrefs(r2p, tmp_job_type_suffix) + self.extract_local_variable_xrefs(r2p, job_type_suffix) } ExtractionJobType::GlobalStrings => { - self.extract_global_strings(r2p, tmp_job_type_suffix) + self.extract_global_strings(r2p, job_type_suffix) } ExtractionJobType::FunctionZignatures => { - self.extract_function_zignatures(r2p, tmp_job_type_suffix) + self.extract_function_zignatures(r2p, job_type_suffix) } ExtractionJobType::FunctionBytes => { - self.extract_function_bytes(r2p, tmp_job_type_suffix) + self.extract_function_bytes(r2p, job_type_suffix) } }?; // 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 {:?}: {}", @@ -750,8 +974,8 @@ 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); + let job_type_suffix = ExtractionJob::get_job_type_suffix(job_type); + let output_path = self.get_output_filepath(&job_type_suffix).unwrap(); if Path::new(&output_path).exists() { warn!( "Skipping {:?} job for {:?}: already processed at {:?}.", @@ -786,26 +1010,6 @@ 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<()> { info!("Starting binary information extraction"); let bininfo_json = r2p @@ -990,17 +1194,17 @@ impl FileToBeProcessed { pub fn extract_function_xrefs(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { let function_details = self.get_function_name_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) + for function_info in function_details { + let function = FunctionToBeProcessed::from(function_info); + 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"); + info!("All functions processed! Writing extracted data to file"); self.write_to_json(&json!(function_xrefs), job_type_suffix)?; Ok(()) } @@ -1010,9 +1214,13 @@ impl FileToBeProcessed { let function_details = self.get_function_name_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_info in function_details { + let function = FunctionToBeProcessed::from(function_info); + 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."); @@ -1027,7 +1235,7 @@ impl FileToBeProcessed { let mut function_pcode = Vec::new(); for function in function_details.iter() { - let ret = self.get_ghidra_pcode_function(function.offset, function.ninstrs, r2p); + let ret = self.get_ghidra_pcode(function.offset, function.ninstrs, r2p); let formatted_obj = PCodeJSONWithFuncName { function_name: function.name.clone(), @@ -1051,23 +1259,18 @@ impl FileToBeProcessed { let function_details = self.get_function_name_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) + for function_info in function_details { + let function = FunctionToBeProcessed::from(function_info); + let bb_info = function.get_basic_block_info(r2p) .with_context(|| { format!( - "Unable to get basic block addresses in {:?} at offset {:?}", - self.file_path, function.offset + "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() { - let ret = self - .get_ghidra_pcode_function( - bb.addr, - bb.ninstr.try_into().unwrap(), // Convert u64 to i64 - r2p, - ) + for bb in bb_info.iter() { + let ret = self.get_ghidra_pcode(bb.addr, bb.ninstr.try_into().unwrap(), r2p) .with_context(|| { format!( "Basic block decompilation failed in {:?} at offset {:?}", @@ -1103,8 +1306,9 @@ impl FileToBeProcessed { let function_details = self.get_function_name_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_info in function_details { + let function = FunctionToBeProcessed::from(function_info); + 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."); @@ -1152,27 +1356,18 @@ impl FileToBeProcessed { pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { info!("Starting function bytes extraction"); let function_details = self.get_function_name_list(r2p)?; + let output_dirpath = self.get_tmp_output_filepath(&job_type_suffix)?; - 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."); + if !output_dirpath.is_dir() { + std::fs::create_dir_all(&output_dirpath) + .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; + } + + for function_info in function_details { + let function = FunctionToBeProcessed::from(function_info); + debug!("Function Name: {} Address: {} Size: {}", + function.name, function.addr, function.size); + function.write_to_bin(r2p, &output_dirpath, &self.func_filename_template)?; } info!("Function bytes successfully extracted"); Ok(()) @@ -1225,35 +1420,13 @@ 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); + ) -> Result { + Self::go_to_address(r2p, address)?; let pcode_ret = r2p.cmd(format!("pdgsd {}", num_instructons).as_str())?; let lines = pcode_ret.lines(); let mut asm_ins = Vec::new(); @@ -1273,33 +1446,6 @@ 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( &self, r2p: &mut R2Pipe, @@ -1314,97 +1460,6 @@ impl FileToBeProcessed { Ok(json_obj) } - fn get_basic_block_addresses( - &self, - function_addr: u64, - r2p: &mut R2Pipe, - ) -> Result { - info!( - "Getting the basic block information for function @ {}", - function_addr - ); - 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))?; - - // 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); - } - } - } - } - - // 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) - } - - 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) - } else { - Err(json.unwrap_err()) - } - } - - 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. @@ -1421,20 +1476,6 @@ impl FileToBeProcessed { 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() - } - - 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()) - } - fn write_to_json(&self, json_obj: &Value, job_type_suffix: String) -> Result<()> { let mut fp_filename = self.get_file_name()?; @@ -1457,66 +1498,14 @@ impl FileToBeProcessed { 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); - - // 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(()); - } - - 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))?; - + /// Seeks to the function address + fn go_to_address(r2p: &mut R2Pipe, address: u64) -> Result<(), Error> { + r2p.cmd(format!("s {}", address).as_str()) + .with_context(|| format!("Failed to seek address {:x}", address))?; 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); diff --git a/src/main.rs b/src/main.rs index c30b533..42a7130 100644 --- a/src/main.rs +++ b/src/main.rs @@ -59,7 +59,7 @@ 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; @@ -304,6 +304,15 @@ enum Commands { #[arg(long, default_value = "false")] use_curl_pdb: bool, + /// Name function data files using symbol, address, or custom template + #[arg( + long, + value_name = "TEMPLATE", + default_value = "symbol", + value_parser = validate_func_filename + )] + func_filename: String, + #[arg(long)] timeout: Option, @@ -1055,6 +1064,7 @@ fn main() { debug, extended_analysis, use_curl_pdb, + func_filename, timeout, with_annotations, } => { @@ -1072,6 +1082,7 @@ fn main() { debug, extended_analysis, use_curl_pdb, + func_filename, timeout, with_annotations, ) diff --git a/src/utils.rs b/src/utils.rs index 3ce29fc..3690878 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -74,6 +74,20 @@ 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()) + } +} + /// Get the JSON paths from a directory /// /// This function takes a path to a directory and traverses all From 887c7df36224a85d11b02f330b040c41b825f891 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 12 Jun 2025 16:40:29 +0100 Subject: [PATCH 03/42] Centralised function name sanitisation (extract+generate) --- src/agcj.rs | 29 +++++++---------------------- src/extract.rs | 26 ++++---------------------- src/utils.rs | 18 ++++++++++++++++++ 3 files changed, 29 insertions(+), 44 deletions(-) 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/extract.rs b/src/extract.rs index 4e80763..9c74417 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1,5 +1,6 @@ use crate::afij::AFIJFunctionInfo; use crate::agcj::AGCJFunctionCallGraph; +use crate::utils::sanitize_filename; use anyhow::anyhow; use anyhow::bail; @@ -15,7 +16,6 @@ use serde_json; use glob::glob; use md5; -use regex::Regex; use serde_json::{json, Deserializer, Value}; use sha1; use sha1::Digest as Sha1Digest; @@ -704,23 +704,6 @@ impl ExtractionJob { } impl FunctionToBeProcessed { - - // Helper functions - fn sanitize_name(name: &str) -> String { - // Replace non-valid characters with '_' - // Valid characters: letters, digits, '_', '-', and '.' - let re = Regex::new(r"[^\w.-]").unwrap(); - re.replace_all(name, "_").into_owned() - } - - fn shorten_name(name: &str) -> String { - if name.len() > 100 { - name[..50].to_string() + &name[name.len() - 50..] - } else { - name.to_string() - } - } - fn write_to_bin(&self, r2p: &mut R2Pipe, output_dirpath: &PathBuf, filename_template: &str) -> Result<()> { let func_bytes = self.get_bytes(r2p)?; let output_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); @@ -733,7 +716,7 @@ impl FunctionToBeProcessed { format!("{:x}", self.addr).trim_start_matches("0x").to_string() } - fn get_func_filename(&self, template: &str, ext: &str) -> 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(), @@ -742,8 +725,7 @@ impl FunctionToBeProcessed { .replace("{ext}", ext), }; - func_filename = Self::sanitize_name(&func_filename); - func_filename = Self::shorten_name(&func_filename); + func_filename = sanitize_filename(&func_filename); if ["symbol", "address"].contains(&template) { // Add an extension only if the user did not specify a custom template func_filename = func_filename + "." + ext; @@ -752,7 +734,7 @@ impl FunctionToBeProcessed { } pub fn get_output_filepath(&self, output_dirpath: &PathBuf, filename_template: &str, ext: &str) -> PathBuf { - let output_filename = self.get_func_filename(filename_template, ext); + let output_filename = self.get_output_filename(filename_template, ext); let mut output_filepath = PathBuf::new(); output_filepath.push(output_dirpath); diff --git a/src/utils.rs b/src/utils.rs index 3690878..954463c 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,5 +1,6 @@ use std::fs::create_dir_all; use std::path::{Path, PathBuf}; +use regex::Regex; use walkdir::WalkDir; /// Formats a save file path @@ -88,6 +89,23 @@ pub fn validate_func_filename(s: &str) -> Result { } } +/// 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 From 3e759137bccdaa42dbe8a2505fa648b45a82a85d Mon Sep 17 00:00:00 2001 From: valbucci Date: Wed, 9 Jul 2025 12:59:53 +0100 Subject: [PATCH 04/42] Changed log level for extraction skipping from warn to debug --- src/extract.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/extract.rs b/src/extract.rs index 9c74417..6269c7d 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -959,7 +959,7 @@ impl FileToBeProcessed { let job_type_suffix = ExtractionJob::get_job_type_suffix(job_type); let output_path = self.get_output_filepath(&job_type_suffix).unwrap(); if Path::new(&output_path).exists() { - warn!( + debug!( "Skipping {:?} job for {:?}: already processed at {:?}.", job_type_suffix, self.file_path, output_path ); From 2bc7327d134c534f03896597fd393152434ba104 Mon Sep 17 00:00:00 2001 From: valbucci Date: Wed, 9 Jul 2025 13:45:10 +0100 Subject: [PATCH 05/42] Fix tmp path bug in 6e5ae8d; Use output path over job suffix in extract functions --- src/extract.rs | 345 +++++++++++++++++++++++++++---------------------- src/utils.rs | 7 +- 2 files changed, 198 insertions(+), 154 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 6269c7d..8d3350d 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -22,13 +22,13 @@ use sha1::Digest as Sha1Digest; use sha2::Digest as Sha2Digest; use sha2::Sha256; use std::collections::HashMap; -use std::sync::LazyLock; use std::env; use std::fs; use std::fs::File; use std::io::{BufReader, Read}; use std::path::{Path, PathBuf}; use std::string::String; +use std::sync::LazyLock; use walkdir::WalkDir; #[derive(PartialEq, Debug)] @@ -57,28 +57,33 @@ pub enum ExtractionJobType { 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::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::FunctionZignatures, "zigs"), - ]) -}); - -static SUFFIX_TO_JOB_TYPE: LazyLock> = LazyLock::new(|| { - JOB_TYPE_TO_SUFFIX.iter().map(|(job_type, suffix)| (*suffix, *job_type)).collect() -}); +static JOB_TYPE_TO_SUFFIX: LazyLock> = + LazyLock::new(|| { + HashMap::from([ + (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::FunctionBytes, "bytes"), + (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 { @@ -265,9 +270,25 @@ impl std::fmt::Display for AFLJFuncDetails { } } -impl From<(String, String, Vec, R2PipeConfig, bool, String)> for FileToBeProcessed { +impl + From<( + String, + String, + Vec, + R2PipeConfig, + bool, + String, + )> for FileToBeProcessed +{ fn from( - orig: (String, String, Vec, R2PipeConfig, bool, String), + orig: ( + String, + String, + Vec, + R2PipeConfig, + bool, + String, + ), ) -> FileToBeProcessed { FileToBeProcessed { file_path: PathBuf::from(orig.0), @@ -536,7 +557,6 @@ pub struct BinaryInfo { pub checksums: Option, // Populated manually with itj. } - impl ExtractionJob { pub fn new( input_path: &PathBuf, @@ -549,7 +569,6 @@ impl ExtractionJob { timeout: &Option, with_annotations: &bool, ) -> Result { - let mut job_types = vec![]; let mut extraction_job_types = vec![]; @@ -632,7 +651,6 @@ impl 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 @@ -654,18 +672,22 @@ impl ExtractionJob { /// 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)) + .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() + 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 { - ExtractionJobType::FunctionBytes => None, // Output is a directory + ExtractionJobType::FunctionBytes => None, // Output is a directory _ => Some("json"), } } @@ -704,25 +726,38 @@ impl ExtractionJob { } impl FunctionToBeProcessed { - fn write_to_bin(&self, r2p: &mut R2Pipe, output_dirpath: &PathBuf, filename_template: &str) -> Result<()> { + fn write_to_bin( + &self, + r2p: &mut R2Pipe, + output_dirpath: &PathBuf, + filename_template: &str, + ) -> Result<()> { let func_bytes = self.get_bytes(r2p)?; let output_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); - std::fs::write(&output_filepath, func_bytes.bytes).with_context(|| format!("Failed to write function bytes to file: {:?}", output_filepath))?; + std::fs::write(&output_filepath, func_bytes.bytes).with_context(|| { + format!( + "Failed to write function bytes to file: {:?}", + output_filepath + ) + })?; Ok(()) } - // Getters + // Getters fn get_hex_address(&self) -> String { - format!("{:x}", self.addr).trim_start_matches("0x").to_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), + _ => template + .replace("{symbol}", &self.name) + .replace("{address}", &self.get_hex_address()) + .replace("{ext}", ext), }; func_filename = sanitize_filename(&func_filename); @@ -733,7 +768,12 @@ impl FunctionToBeProcessed { func_filename } - pub fn get_output_filepath(&self, output_dirpath: &PathBuf, filename_template: &str, ext: &str) -> PathBuf { + 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(); @@ -746,23 +786,25 @@ impl FunctionToBeProcessed { FileToBeProcessed::go_to_address(r2p, self.addr)?; let mut function_bytes = r2p.cmd(format!("p8 {}", self.size).as_str())?; function_bytes = function_bytes.trim().to_string(); - let decoded_bytes = hex::decode(&function_bytes) - .context("Failed to decode hex bytes")?; + let decoded_bytes = hex::decode(&function_bytes).context("Failed to decode hex bytes")?; Ok(FuncBytes { - bytes: decoded_bytes + bytes: decoded_bytes, }) } fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { - info!("Getting the basic block information for function @ {}", self.addr); + info!( + "Getting the basic block information for function @ {}", + self.addr + ); FileToBeProcessed::go_to_address(r2p, self.addr)?; let json = r2p.cmd("afbj").context("Command afbj failed")?; // 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))?; - + // 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() { @@ -774,25 +816,32 @@ impl FunctionToBeProcessed { } } } - + // Deserialize JSON into a BasicBlockInfo struct - let bb_info: BasicBlockInfo = - serde_json::from_value(value.clone()).with_context(|| { - format!( - "Unable to convert {:?} into a BasicBlockInfo struct!", - value - ) - })?; + let bb_info: BasicBlockInfo = serde_json::from_value(value.clone()).with_context(|| { + format!( + "Unable to convert {:?} into a BasicBlockInfo struct!", + value + ) + })?; Ok(bb_info) } - fn get_local_variable_xref_details(&self, r2p: &mut R2Pipe) -> Result { - info!("Getting local variable xref details for function @ {}", self.addr); + fn get_local_variable_xref_details( + &self, + r2p: &mut R2Pipe, + ) -> Result { + info!( + "Getting local variable xref details for function @ {}", + self.addr + ); FileToBeProcessed::go_to_address(r2p, self.addr)?; let json = r2p.cmd("axvj").context("Command axvj failed")?; - let local_variable_xrefs: LocalVariableXrefs = serde_json::from_str(&json) - .with_context(|| format!("Unable to convert {:?} to LocalVariableXrefs struct!", json))?; + let local_variable_xrefs: LocalVariableXrefs = + serde_json::from_str(&json).with_context(|| { + format!("Unable to convert {:?} to LocalVariableXrefs struct!", json) + })?; Ok(local_variable_xrefs) } @@ -800,9 +849,14 @@ impl FunctionToBeProcessed { info!("Getting xref details for function @ {}", self.addr); FileToBeProcessed::go_to_address(r2p, self.addr)?; let json = r2p.cmd("axffj").context("Command axffj failed")?; - let mut json_obj: Vec = serde_json::from_str(&json) - .with_context(|| format!("Unable to convert {:?} to FunctionXrefDetails struct!", json))?; - + let mut json_obj: Vec = + serde_json::from_str(&json).with_context(|| { + format!( + "Unable to convert {:?} to FunctionXrefDetails struct!", + json + ) + })?; + // 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() { @@ -818,7 +872,11 @@ impl FunctionToBeProcessed { Ok(json_obj) } - fn get_ghidra_decomp(&self, r2p: &mut R2Pipe, with_annotations: bool) -> Result { + fn get_ghidra_decomp( + &self, + r2p: &mut R2Pipe, + with_annotations: bool, + ) -> Result { FileToBeProcessed::go_to_address(r2p, self.addr)?; let json = r2p.cmd("pdgj").context("Command pdgj failed")?; @@ -853,12 +911,13 @@ impl FileToBeProcessed { /// 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 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 { + + if job_type == ExtractionJobType::Decompilation && self.with_annotations { output_filename = output_filename + "_" + job_type_suffix + "_annotations" + &ext_str; } else { output_filename = output_filename + "_" + job_type_suffix + &ext_str; @@ -875,8 +934,11 @@ impl FileToBeProcessed { } 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 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) } @@ -890,38 +952,30 @@ impl FileToBeProcessed { let tmp_output_path = self.get_tmp_output_filepath(&job_type_suffix)?; match job_type { - ExtractionJobType::BinInfo => self.extract_binary_info(r2p, job_type_suffix), + ExtractionJobType::BinInfo => self.extract_binary_info(r2p, &tmp_output_path), ExtractionJobType::RegisterBehaviour => { - self.extract_register_behaviour(r2p, job_type_suffix) - } - ExtractionJobType::FunctionXrefs => { - self.extract_function_xrefs(r2p, job_type_suffix) + self.extract_register_behaviour(r2p, &tmp_output_path) } - ExtractionJobType::CFG => self.extract_func_cfgs(r2p, job_type_suffix), + ExtractionJobType::FunctionXrefs => self.extract_function_xrefs(r2p, &tmp_output_path), + ExtractionJobType::CFG => self.extract_func_cfgs(r2p, &tmp_output_path), ExtractionJobType::CallGraphs => { - self.extract_function_call_graphs(r2p, job_type_suffix) + self.extract_function_call_graphs(r2p, &tmp_output_path) } - ExtractionJobType::FuncInfo => self.extract_function_info(r2p, job_type_suffix), + ExtractionJobType::FuncInfo => self.extract_function_info(r2p, &tmp_output_path), ExtractionJobType::FunctionVariables => { - self.extract_function_variables(r2p, job_type_suffix) + self.extract_function_variables(r2p, &tmp_output_path) } - ExtractionJobType::Decompilation => { - self.extract_decompilation(r2p, job_type_suffix) - } - ExtractionJobType::PCodeFunc => self.extract_pcode_function(r2p, job_type_suffix), - ExtractionJobType::PCodeBB => self.extract_pcode_basic_block(r2p, 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, job_type_suffix) - } - ExtractionJobType::GlobalStrings => { - self.extract_global_strings(r2p, 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, job_type_suffix) - } - ExtractionJobType::FunctionBytes => { - self.extract_function_bytes(r2p, job_type_suffix) + self.extract_function_zignatures(r2p, &tmp_output_path) } + ExtractionJobType::FunctionBytes => self.extract_function_bytes(r2p, &tmp_output_path), }?; // Apply final output file name when extraction is done @@ -992,7 +1046,7 @@ impl FileToBeProcessed { } } - 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") @@ -1048,14 +1102,14 @@ 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 mut register_behaviour_vec: HashMap = HashMap::new(); @@ -1077,14 +1131,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 @@ -1094,17 +1148,17 @@ 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<()> { + pub fn extract_function_info(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting function metdata extraction"); let function_details: Vec = self.get_function_name_list(r2p)?; info!("Writing extracted data to file"); let json = json!(function_details); - self.write_to_json(&json, job_type_suffix) + self.write_to_json(&json, output_path) .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; Ok(()) } @@ -1112,7 +1166,7 @@ impl FileToBeProcessed { 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)?; @@ -1133,11 +1187,11 @@ 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<()> { + pub fn extract_func_cfgs(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Executing agfj @@f on {:?}", self.file_path); let json_raw = r2p.cmd("agfj @@f").with_context(|| { @@ -1158,7 +1212,7 @@ impl FileToBeProcessed { self.file_path )); } else { - self.write_to_json(&json, job_type_suffix)?; + self.write_to_json(&json, output_path)?; } } Err(e) => { @@ -1173,45 +1227,48 @@ impl FileToBeProcessed { Ok(()) } - pub fn extract_function_xrefs(&self, r2p: &mut R2Pipe, job_type_suffix: String) -> Result<()> { + pub fn extract_function_xrefs(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { let function_details = self.get_function_name_list(r2p)?; let mut function_xrefs: HashMap> = HashMap::new(); info!("Extracting xrefs for each function"); for function_info in function_details { let function = FunctionToBeProcessed::from(function_info); - let ret = function.get_xref_details(r2p) - .with_context(|| { - format!("Unable to get function xrefs from {:?}", self.file_path) - })?; + 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! Writing extracted data to file"); - self.write_to_json(&json!(function_xrefs), job_type_suffix)?; + 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 mut function_decomp: HashMap = HashMap::new(); for function_info in function_details { let function = FunctionToBeProcessed::from(function_info); - let ret = function.get_ghidra_decomp(r2p, self.with_annotations) + let ret = function + .get_ghidra_decomp(r2p, self.with_annotations) .with_context(|| { - format!("Unable to get decompilation for {:?} @ {:?}", self.file_path, function.addr) + 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 mut function_pcode = Vec::new(); @@ -1228,31 +1285,27 @@ 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 mut function_pcode = Vec::new(); for function_info in function_details { let function = FunctionToBeProcessed::from(function_info); - 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 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_info.iter() { - let ret = self.get_ghidra_pcode(bb.addr, bb.ninstr.try_into().unwrap(), r2p) + let ret = self + .get_ghidra_pcode(bb.addr, bb.ninstr.try_into().unwrap(), r2p) .with_context(|| { format!( "Basic block decompilation failed in {:?} at offset {:?}", @@ -1275,14 +1328,14 @@ 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)?; @@ -1296,11 +1349,11 @@ impl FileToBeProcessed { 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") @@ -1310,14 +1363,14 @@ impl FileToBeProcessed { 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 @@ -1331,24 +1384,25 @@ 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) -> Result<()> { info!("Starting function bytes extraction"); let function_details = self.get_function_name_list(r2p)?; - let output_dirpath = self.get_tmp_output_filepath(&job_type_suffix)?; if !output_dirpath.is_dir() { std::fs::create_dir_all(&output_dirpath) .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; } - + for function_info in function_details { let function = FunctionToBeProcessed::from(function_info); - debug!("Function Name: {} Address: {} Size: {}", - function.name, function.addr, function.size); + debug!( + "Function Name: {} Address: {} Size: {}", + function.name, function.addr, function.size + ); function.write_to_bin(r2p, &output_dirpath, &self.func_filename_template)?; } info!("Function bytes successfully extracted"); @@ -1458,20 +1512,7 @@ impl FileToBeProcessed { json_objects.map(Value::Array) } - 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); - + fn write_to_json(&self, json_obj: &Value, output_filepath: &PathBuf) -> Result<()> { let file = File::create(&output_filepath) .with_context(|| format!("Unable to create file {:?}", output_filepath))?; serde_json::to_writer(&file, &json_obj) diff --git a/src/utils.rs b/src/utils.rs index 954463c..b4a1e8d 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,6 +1,6 @@ +use regex::Regex; use std::fs::create_dir_all; use std::path::{Path, PathBuf}; -use regex::Regex; use walkdir::WalkDir; /// Formats a save file path @@ -85,7 +85,10 @@ pub fn validate_func_filename(s: &str) -> Result { } 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()) + Err( + "must be `symbol`, `address`, or a pattern containing `{symbol}`/`{address}`" + .to_string(), + ) } } From 9d36bf7db5bad2585ab88f9ac2b19098c748dbce Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 10 Jul 2025 13:31:33 +0100 Subject: [PATCH 06/42] Added error logging and option to skip extraction if previously aborted --- src/extract.rs | 40 ++++++++++++++++++++++++++++++++-------- src/main.rs | 6 ++++++ 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 8d3350d..6898e55 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -26,7 +26,7 @@ use std::env; use std::fs; use std::fs::File; use std::io::{BufReader, Read}; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::string::String; use std::sync::LazyLock; use walkdir::WalkDir; @@ -92,6 +92,7 @@ pub struct FileToBeProcessed { pub job_types: Vec, pub r2p_config: R2PipeConfig, pub with_annotations: bool, + pub retry_aborted: bool, pub func_filename_template: String, } @@ -277,6 +278,7 @@ impl Vec, R2PipeConfig, bool, + bool, String, )> for FileToBeProcessed { @@ -287,6 +289,7 @@ impl Vec, R2PipeConfig, bool, + bool, String, ), ) -> FileToBeProcessed { @@ -296,7 +299,8 @@ impl job_types: orig.2, r2p_config: orig.3, with_annotations: orig.4, - func_filename_template: orig.5, + retry_aborted: orig.5, + func_filename_template: orig.6, } } } @@ -568,6 +572,7 @@ impl ExtractionJob { func_filename_template: &str, timeout: &Option, with_annotations: &bool, + retry_aborted: &bool, ) -> Result { let mut job_types = vec![]; let mut extraction_job_types = vec![]; @@ -604,6 +609,7 @@ impl ExtractionJob { job_types: extraction_job_types, // Use the vector of just ExtractionJobType r2p_config: r2_handle_config, with_annotations: *with_annotations, + retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), }; @@ -635,6 +641,7 @@ impl ExtractionJob { job_types: extraction_job_types.clone(), r2p_config: r2_handle_config, with_annotations: *with_annotations, + retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), }) .collect(); @@ -1011,13 +1018,27 @@ impl FileToBeProcessed { // Check if the extracted data file already exists let job_type_suffix = ExtractionJob::get_job_type_suffix(job_type); - let output_path = self.get_output_filepath(&job_type_suffix).unwrap(); - if Path::new(&output_path).exists() { + 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 = output_path.with_extension("error.log"); + + if output_path.exists() { debug!( "Skipping {:?} job for {:?}: already processed at {:?}.", job_type_suffix, self.file_path, output_path ); continue; + } else if error_path.exists() && !self.retry_aborted { + debug!( + "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. @@ -1032,10 +1053,13 @@ impl FileToBeProcessed { "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(); + } } } diff --git a/src/main.rs b/src/main.rs index 42a7130..a612e60 100644 --- a/src/main.rs +++ b/src/main.rs @@ -318,6 +318,10 @@ enum Commands { #[arg(long, default_value = "false")] with_annotations: bool, + + /// Toggle to retry previously aborted jobs due to extraction failures + #[arg(long, default_value = "false")] + retry_aborted: bool, }, /// Generate single embeddings on the fly /// @@ -1067,6 +1071,7 @@ fn main() { func_filename, timeout, with_annotations, + retry_aborted, } => { info!("Creating extraction job with {} modes", modes.len()); if !output_dir.exists() { @@ -1085,6 +1090,7 @@ fn main() { func_filename, timeout, with_annotations, + retry_aborted, ) .unwrap_or_else(|e| { error!("Failed to create extraction job: {}", e); From c91b05175f364b1f756e21053a04787770900744 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 10 Jul 2025 19:08:01 +0100 Subject: [PATCH 07/42] Added --log-level option to set verbosity --- src/main.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/main.rs b/src/main.rs index a612e60..74399ed 100644 --- a/src/main.rs +++ b/src/main.rs @@ -93,6 +93,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, } @@ -411,12 +415,13 @@ enum DedupSubCommands { } fn main() { + let cli = Cli::parse(); + let env = Env::default() - .filter_or("LOG_LEVEL", "warn") + .filter_or("LOG_LEVEL", &cli.log_level) .write_style_or("LOG_STYLE", "always"); env_logger::init_from_env(env); - let cli = Cli::parse(); match &cli.command { #[cfg(feature = "goblin")] Commands::Info { path } => { From 03e4ac77f10583d9801ee0025aa0669ad7f3fd5d Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 10 Jul 2025 19:51:07 +0100 Subject: [PATCH 08/42] Fixed mix up between log lines and progress bar --- Cargo.toml | 1 + src/main.rs | 81 +++++++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 67 insertions(+), 15 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 8c43039..217b035 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ 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-log-bridge = "0.2.3" itertools = "0.10.5" prettytable = "0.10.0" serde-aux = "4.2.0" diff --git a/src/main.rs b/src/main.rs index 74399ed..de5ce2a 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; @@ -64,6 +66,22 @@ 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, @@ -417,11 +435,28 @@ enum DedupSubCommands { fn main() { let cli = Cli::parse(); - let env = Env::default() - .filter_or("LOG_LEVEL", &cli.log_level) - .write_style_or("LOG_STYLE", "always"); - - env_logger::init_from_env(env); + // 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 } => { @@ -684,7 +719,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), @@ -715,7 +751,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"); @@ -744,7 +782,8 @@ fn main() { .zip(metadata_paths_vec) .collect::>(); - combined_cgs_metadata.par_iter().progress().for_each( + 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( @@ -831,8 +870,10 @@ fn main() { full_output_path.to_string_lossy() ) } + pb.inc(1); }, ); + pb.finish(); } } } @@ -1025,7 +1066,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), @@ -1040,8 +1082,10 @@ fn main() { instruction_type, random_walk, *pairs, - ) + ); + pb.inc(1); } + pb.finish(); } } GenerateSubCommands::Tokeniser { @@ -1116,11 +1160,15 @@ 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()); + .for_each(|path| { + path.process_all_modes(); + pb.inc(1); + }); + pb.finish(); } else if job.input_path_type == PathType::File { info!("Single file found"); @@ -1195,9 +1243,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(); } }, } From 63369f66866f17bc8f3b88799b6b2f76096f6e53 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 10 Jul 2025 20:16:55 +0100 Subject: [PATCH 09/42] Bumped indicatif to 0.17.11 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 217b035..1764da6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ 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.17.11", features = ["rayon"]} indicatif-log-bridge = "0.2.3" itertools = "0.10.5" prettytable = "0.10.0" From ebba9893500617d93d61479faff337ea3bfd241a Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 11 Sep 2025 20:47:10 +0100 Subject: [PATCH 10/42] Added 'addr' alias for AFIJFunctionInfo 'offset' --- src/afij.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/afij.rs b/src/afij.rs index d2266c8..c90d57e 100644 --- a/src/afij.rs +++ b/src/afij.rs @@ -6,6 +6,7 @@ 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: u64, From 21eb2df11bd5fb89c9b8e8be87e322dd850a19f6 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 11 Sep 2025 21:25:11 +0100 Subject: [PATCH 11/42] Bumped indicatif to 0.18.0 --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 1764da6..3e1cdf7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ 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.11", features = ["rayon"]} +indicatif = { version = "0.18.0", features = ["rayon"]} indicatif-log-bridge = "0.2.3" itertools = "0.10.5" prettytable = "0.10.0" From 9dfd5e5f3b18d0747e603a1dc80cfdf3db7695fb Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 11 Sep 2025 21:35:22 +0100 Subject: [PATCH 12/42] Added 'minaddr' and 'maxaddr' aliases --- src/afij.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/afij.rs b/src/afij.rs index c90d57e..0776329 100644 --- a/src/afij.rs +++ b/src/afij.rs @@ -28,7 +28,9 @@ pub struct AFIJFunctionInfo { pub edges: u64, pub ebbs: u64, pub signature: String, + #[serde(alias = "minaddr")] pub minbound: u64, + #[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 From 1e087cfb5e88d327753ec9bb2bec6e72975b776c Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 9 Oct 2025 11:01:52 +0100 Subject: [PATCH 13/42] Changed offset data types to u64 and added addr alias --- src/agfj.rs | 12 +++++++----- src/bb.rs | 33 ++++++++++++++++----------------- src/utils.rs | 10 ++++++++++ 3 files changed, 33 insertions(+), 22 deletions(-) diff --git a/src/agfj.rs b/src/agfj.rs index 90adcdf..4352eda 100644 --- a/src/agfj.rs +++ b/src/agfj.rs @@ -41,12 +41,13 @@ pub struct AGFJFunc { nargs: u64, ninstr: u64, nlocals: u64, + #[serde(alias = "addr")] 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; @@ -811,7 +813,7 @@ 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..c6dc387 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)] 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/utils.rs b/src/utils.rs index b4a1e8d..19caffd 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -164,6 +164,16 @@ 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::*; From 3f9e509f2b451508a9aa76571829c548a824dd22 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 9 Oct 2025 13:06:17 +0100 Subject: [PATCH 14/42] Implemented JSON stream function for @@f commands --- Cargo.toml | 2 +- src/extract.rs | 100 +++++++++++++++++++++++++++---------------------- 2 files changed, 56 insertions(+), 46 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3e1cdf7..f5cd8c0 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" diff --git a/src/extract.rs b/src/extract.rs index 6898e55..78e47a6 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -16,7 +16,8 @@ use serde_json; use glob::glob; use md5; -use serde_json::{json, Deserializer, Value}; +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; @@ -25,7 +26,7 @@ use std::collections::HashMap; use std::env; use std::fs; use std::fs::File; -use std::io::{BufReader, Read}; +use std::io::{BufReader, BufWriter, Read}; use std::path::PathBuf; use std::string::String; use std::sync::LazyLock; @@ -1218,36 +1219,12 @@ impl FileToBeProcessed { pub fn extract_func_cfgs(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Executing agfj @@f on {:?}", self.file_path); - let json_raw = r2p.cmd("agfj @@f").with_context(|| { - format!( - "Failed to extract control flow graph information from {:?}.", - self.file_path - ) - })?; + let json_raw = r2p + .cmd("agfj @@f") + .with_context(|| format!("Failed to extract CFGs from {:?}.", self.file_path))?; - 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, output_path)?; - } - } - Err(e) => { - return Err(anyhow::anyhow!( - "Unable to parse json for {:?}: {}: {}", - self.file_path, - json_raw, - e - )); - } - } + self.stream_write_to_json(&json_raw, output_path) + .with_context(|| format!("Failed to write CFGs to {:?}.", self.file_path))?; Ok(()) } @@ -1521,27 +1498,60 @@ impl FileToBeProcessed { } // 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| { - match result { - Ok(Value::Array(ref arr)) if arr.is_empty() => None, // skip empty arrays - other => Some(other), - } - }) - .collect(); - // Map the collected vector into a JSON array. - json_objects.map(Value::Array) - } + /// 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) .with_context(|| format!("Failed to write JSON to {:?}", output_filepath))?; + info!("JSON written to {:?}", output_filepath); + Ok(()) + } + + /// 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::>(); + + // 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"); + } + } + } + } + seq.end()?; + info!("JSON stream written to {:?}", output_filepath); Ok(()) } From b64edf715efdaba480aab7143359423adfc0c7df Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 9 Oct 2025 13:06:36 +0100 Subject: [PATCH 15/42] Cargo fmt --- src/agfj.rs | 5 ++++- src/main.rs | 44 ++++++++++++++++++++++++-------------------- src/utils.rs | 3 +-- 3 files changed, 29 insertions(+), 23 deletions(-) diff --git a/src/agfj.rs b/src/agfj.rs index 4352eda..20339f7 100644 --- a/src/agfj.rs +++ b/src/agfj.rs @@ -813,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, get_default_addr()); + 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/main.rs b/src/main.rs index de5ce2a..b43f3b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -71,7 +71,10 @@ static MULTI_PROGRESS: std::sync::OnceLock> = std::sync::Once /// Get the global multi-progress instance fn get_multi_progress() -> Arc { - MULTI_PROGRESS.get().expect("Multi-progress not initialized").clone() + MULTI_PROGRESS + .get() + .expect("Multi-progress not initialized") + .clone() } /// Create a progress bar that works with the log bridge @@ -434,28 +437,31 @@ enum DedupSubCommands { fn main() { 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"); - + 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(); - + .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")] @@ -783,8 +789,9 @@ fn main() { .collect::>(); let pb = create_progress_bar(combined_cgs_metadata.len() as u64); - combined_cgs_metadata.par_iter().for_each( - |(filepath, metapath)| { + 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), @@ -871,8 +878,7 @@ fn main() { ) } pb.inc(1); - }, - ); + }); pb.finish(); } } @@ -1162,12 +1168,10 @@ fn main() { 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() - .for_each(|path| { - path.process_all_modes(); - pb.inc(1); - }); + 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"); diff --git a/src/utils.rs b/src/utils.rs index 19caffd..9a5bfe8 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -164,11 +164,10 @@ 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 +/// 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 From 7d277862d455b96765cea44402fb95f59798ea43 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 9 Oct 2025 13:33:40 +0100 Subject: [PATCH 16/42] Enhanced bytes extraction with mask (p8fmj) --- src/extract.rs | 28 ++++++++++++---------------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 78e47a6..d765cdb 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -695,7 +695,7 @@ impl ExtractionJob { fn get_output_extension(job_type: &ExtractionJobType) -> Option<&str> { match job_type { - ExtractionJobType::FunctionBytes => None, // Output is a directory + // Add here if output is not a JSON file (e.g. txt file or directory) _ => Some("json"), } } @@ -792,7 +792,7 @@ impl FunctionToBeProcessed { fn get_bytes(&self, r2p: &mut R2Pipe) -> Result { FileToBeProcessed::go_to_address(r2p, self.addr)?; - let mut function_bytes = r2p.cmd(format!("p8 {}", self.size).as_str())?; + let mut function_bytes = r2p.cmd("p8f")?; function_bytes = function_bytes.trim().to_string(); let decoded_bytes = hex::decode(&function_bytes).context("Failed to decode hex bytes")?; @@ -1389,23 +1389,19 @@ impl FileToBeProcessed { Ok(()) } - pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, output_dirpath: &PathBuf) -> Result<()> { + pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting function bytes extraction"); - let function_details = self.get_function_name_list(r2p)?; - if !output_dirpath.is_dir() { - std::fs::create_dir_all(&output_dirpath) - .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; - } + let json_raw = r2p.cmd("p8fmj @@f").with_context(|| { + format!( + "Failed to extract function bytes from {:?}.", + self.file_path + ) + })?; + + self.stream_write_to_json(&json_raw, output_path) + .with_context(|| format!("Failed to write function bytes to {:?}.", self.file_path))?; - for function_info in function_details { - let function = FunctionToBeProcessed::from(function_info); - debug!( - "Function Name: {} Address: {} Size: {}", - function.name, function.addr, function.size - ); - function.write_to_bin(r2p, &output_dirpath, &self.func_filename_template)?; - } info!("Function bytes successfully extracted"); Ok(()) } From 3148d3f2cb6649fe3d2ac7d23ab6f32a4ab0c588 Mon Sep 17 00:00:00 2001 From: valbucci Date: Sat, 11 Oct 2025 13:23:10 +0100 Subject: [PATCH 17/42] Added bytes-masked extraction and restored function-separated bin files --- src/extract.rs | 77 ++++++++++++++++++++++++++++++++++++-------------- src/main.rs | 3 +- 2 files changed, 58 insertions(+), 22 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index d765cdb..b53f23e 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -55,6 +55,7 @@ pub enum ExtractionJobType { LocalVariableXrefs, GlobalStrings, FunctionBytes, + FunctionBytesMasked, FunctionZignatures, } @@ -74,6 +75,7 @@ static JOB_TYPE_TO_SUFFIX: LazyLock> = (ExtractionJobType::LocalVariableXrefs, "localvar-xrefs"), (ExtractionJobType::GlobalStrings, "strings"), (ExtractionJobType::FunctionBytes, "bytes"), + (ExtractionJobType::FunctionBytesMasked, "bytes-masked"), (ExtractionJobType::FunctionZignatures, "zigs"), ]) }); @@ -445,6 +447,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) @@ -696,6 +699,7 @@ impl ExtractionJob { fn get_output_extension(job_type: &ExtractionJobType) -> Option<&str> { match job_type { // Add here if output is not a JSON file (e.g. txt file or directory) + ExtractionJobType::FunctionBytes => None, // Output is a directory _ => Some("json"), } } @@ -739,15 +743,30 @@ impl FunctionToBeProcessed { r2p: &mut R2Pipe, output_dirpath: &PathBuf, filename_template: &str, + apply_mask: bool, ) -> Result<()> { - let func_bytes = self.get_bytes(r2p)?; - let output_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); - std::fs::write(&output_filepath, func_bytes.bytes).with_context(|| { + let func_bytes = self.get_bytes(r2p, apply_mask).context("Failed to get function bytes")?; + let bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); + let masked_bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); + + info!("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: {:?}", - output_filepath + bytes_filepath ) })?; + + if apply_mask { + let bytes_mask = func_bytes.mask.context("Failed to get function masked bytes")?; + info!("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(()) } @@ -790,15 +809,24 @@ impl FunctionToBeProcessed { output_filepath } - fn get_bytes(&self, r2p: &mut R2Pipe) -> Result { + fn get_bytes(&self, r2p: &mut R2Pipe, apply_mask: bool) -> Result { FileToBeProcessed::go_to_address(r2p, self.addr)?; - let mut function_bytes = r2p.cmd("p8f")?; - function_bytes = function_bytes.trim().to_string(); - let decoded_bytes = hex::decode(&function_bytes).context("Failed to decode hex bytes")?; + let mut function_bytes_and_mask = r2p.cmd("p8fm")?; + function_bytes_and_mask = function_bytes_and_mask.trim().to_string(); + let parts: Vec<&str> = function_bytes_and_mask.split(":").collect(); + let function_bytes = hex::decode(parts[0]).context("Failed to decode hex function bytes")?; + let mut masked_bytes: Option> = None; + + if apply_mask { + let bytes_mask = hex::decode(parts[1]).context("Failed to decode hex bytes mask")?; + let mut masked_bytes_tmp = vec![0; function_bytes.len()]; + for i in 0..function_bytes.len() { + masked_bytes_tmp[i] = function_bytes[i] & bytes_mask[i]; + } + masked_bytes = Some(masked_bytes_tmp); + } - Ok(FuncBytes { - bytes: decoded_bytes, - }) + Ok(FuncBytes { bytes: function_bytes, mask: masked_bytes }) } fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { @@ -983,7 +1011,8 @@ impl FileToBeProcessed { ExtractionJobType::FunctionZignatures => { self.extract_function_zignatures(r2p, &tmp_output_path) } - ExtractionJobType::FunctionBytes => self.extract_function_bytes(r2p, &tmp_output_path), + ExtractionJobType::FunctionBytes => 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 @@ -1389,18 +1418,24 @@ impl FileToBeProcessed { Ok(()) } - pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> 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)?; - let json_raw = r2p.cmd("p8fmj @@f").with_context(|| { - format!( - "Failed to extract function bytes from {:?}.", - self.file_path - ) - })?; + if !output_dirpath.is_dir() { + std::fs::create_dir_all(&output_dirpath) + .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; + } - self.stream_write_to_json(&json_raw, output_path) - .with_context(|| format!("Failed to write function bytes to {:?}.", self.file_path))?; + for function_info in function_details { + let function = FunctionToBeProcessed::from(function_info); + debug!( + "Function Name: {} Address: {} Size: {}", + function.name, function.addr, function.size + ); + function.write_to_bin(r2p, &output_dirpath, &self.func_filename_template, apply_mask)?; + } info!("Function bytes successfully extracted"); Ok(()) diff --git a/src/main.rs b/src/main.rs index b43f3b4..ab24a30 100644 --- a/src/main.rs +++ b/src/main.rs @@ -309,7 +309,7 @@ enum Commands { #[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" + "pcode-func", "pcode-bb", "localvar-xrefs", "strings", "bytes", "bytes-masked", "zigs" ]) .map(|s| s.parse::().unwrap()), num_args = 1.., @@ -330,6 +330,7 @@ enum Commands { use_curl_pdb: bool, /// Name function data files using symbol, address, or custom template + /// e.g. '{address}-{symbol}.{ext}' #[arg( long, value_name = "TEMPLATE", From 097faa9aad2da6ac5cf8e1f5a8061637c1b00c02 Mon Sep 17 00:00:00 2001 From: valbucci Date: Sun, 12 Oct 2025 15:28:18 +0100 Subject: [PATCH 18/42] Added func-cfg mode: same as CFG, but split into multiple files --- src/agfj.rs | 12 +-- src/bb.rs | 2 +- src/extract.rs | 202 +++++++++++++++++++++++++++++++++++++++++++------ src/main.rs | 2 +- 4 files changed, 187 insertions(+), 31 deletions(-) diff --git a/src/agfj.rs b/src/agfj.rs index 20339f7..289b46c 100644 --- a/src/agfj.rs +++ b/src/agfj.rs @@ -35,15 +35,15 @@ 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, - #[serde(alias = "addr")] - offset: u64, - size: Option, stack: u64, r#type: String, pub blocks: Vec, @@ -794,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()); diff --git a/src/bb.rs b/src/bb.rs index c6dc387..a397d40 100644 --- a/src/bb.rs +++ b/src/bb.rs @@ -105,7 +105,7 @@ pub struct Op { } #[serde_as] -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Clone, Debug)] pub struct ACFJBlock { #[serde(default = "get_default_addr", alias = "addr")] pub offset: u64, diff --git a/src/extract.rs b/src/extract.rs index b53f23e..fe9799f 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1,5 +1,6 @@ use crate::afij::AFIJFunctionInfo; use crate::agcj::AGCJFunctionCallGraph; +use crate::agfj::AGFJFunc; use crate::utils::sanitize_filename; use anyhow::anyhow; @@ -46,6 +47,7 @@ pub enum ExtractionJobType { RegisterBehaviour, FunctionXrefs, CFG, + FunctionCFG, // Like CFG, but in a separate file for each function CallGraphs, FuncInfo, FunctionVariables, @@ -66,6 +68,7 @@ static JOB_TYPE_TO_SUFFIX: LazyLock> = (ExtractionJobType::RegisterBehaviour, "reg"), (ExtractionJobType::FunctionXrefs, "func-xrefs"), (ExtractionJobType::CFG, "cfg"), + (ExtractionJobType::FunctionCFG, "func-cfg"), (ExtractionJobType::CallGraphs, "cg"), (ExtractionJobType::FuncInfo, "finfo"), (ExtractionJobType::FunctionVariables, "fvars"), @@ -341,6 +344,17 @@ impl From for FunctionToBeProcessed { } } +impl From for FunctionToBeProcessed { + fn from(func: AGFJFunc) -> Self { + FunctionToBeProcessed { + name: func.name, + addr: func.offset, + size: func.size, + ninstrs: func.ninstr, + } + } +} + // Structs for pdgj - Ghidra Decomp JSON output #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct DecompJSON { @@ -698,8 +712,9 @@ impl ExtractionJob { fn get_output_extension(job_type: &ExtractionJobType) -> Option<&str> { match job_type { - // Add here if output is not a JSON file (e.g. txt file or directory) - ExtractionJobType::FunctionBytes => None, // Output is a directory + // Add here if output is not a JSON file (e.g. None for a directory) + ExtractionJobType::FunctionBytes => None, + ExtractionJobType::FunctionCFG => None, _ => Some("json"), } } @@ -738,6 +753,8 @@ impl ExtractionJob { } 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, @@ -745,11 +762,14 @@ impl FunctionToBeProcessed { filename_template: &str, apply_mask: bool, ) -> Result<()> { - let func_bytes = self.get_bytes(r2p, apply_mask).context("Failed to get function bytes")?; + let func_bytes = self + .get_bytes(r2p, apply_mask) + .context("Failed to get function bytes")?; let bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); - let masked_bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); - - info!("Writing function bytes to file: {:?}", bytes_filepath); + let masked_bytes_filepath = + self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); + + 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: {:?}", @@ -758,8 +778,13 @@ impl FunctionToBeProcessed { })?; if apply_mask { - let bytes_mask = func_bytes.mask.context("Failed to get function masked bytes")?; - info!("Writing function masked bytes to file: {:?}", masked_bytes_filepath); + let bytes_mask = func_bytes + .mask + .context("Failed to get function masked bytes")?; + 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: {:?}", @@ -770,6 +795,42 @@ impl FunctionToBeProcessed { 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 fn get_hex_address(&self) -> String { format!("{:x}", self.addr) @@ -814,7 +875,8 @@ impl FunctionToBeProcessed { let mut function_bytes_and_mask = r2p.cmd("p8fm")?; function_bytes_and_mask = function_bytes_and_mask.trim().to_string(); let parts: Vec<&str> = function_bytes_and_mask.split(":").collect(); - let function_bytes = hex::decode(parts[0]).context("Failed to decode hex function bytes")?; + let function_bytes = + hex::decode(parts[0]).context("Failed to decode hex function bytes")?; let mut masked_bytes: Option> = None; if apply_mask { @@ -826,7 +888,10 @@ impl FunctionToBeProcessed { masked_bytes = Some(masked_bytes_tmp); } - Ok(FuncBytes { bytes: function_bytes, mask: masked_bytes }) + Ok(FuncBytes { + bytes: function_bytes, + mask: masked_bytes, + }) } fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { @@ -934,6 +999,29 @@ impl FunctionToBeProcessed { Ok(parsed_obj) } } + + fn get_cfg(&self, r2p: &mut R2Pipe) -> Result { + FileToBeProcessed::go_to_address(r2p, self.addr)?; + debug!("Getting CFG for function @ {}", self.addr); + // let json_str = r2p.cmd("agfj").context("Command agfj failed")?; + let json = r2p.cmdj("agfj").context("Command agfj failed")?; + // let cfg: AGFJFunc = serde_json::from_str(&json_str) + // .with_context(|| format!("Unable to convert {:?} to AGFJFunc struct!", json_str))?; + + // 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 { @@ -993,7 +1081,8 @@ impl FileToBeProcessed { self.extract_register_behaviour(r2p, &tmp_output_path) } ExtractionJobType::FunctionXrefs => self.extract_function_xrefs(r2p, &tmp_output_path), - ExtractionJobType::CFG => self.extract_func_cfgs(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_output_path) } @@ -1011,8 +1100,12 @@ impl FileToBeProcessed { ExtractionJobType::FunctionZignatures => { self.extract_function_zignatures(r2p, &tmp_output_path) } - ExtractionJobType::FunctionBytes => self.extract_function_bytes(r2p, &tmp_output_path, false), - ExtractionJobType::FunctionBytesMasked => self.extract_function_bytes(r2p, &tmp_output_path, true), + ExtractionJobType::FunctionBytes => { + 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 @@ -1245,15 +1338,68 @@ impl FileToBeProcessed { Ok(()) } - pub fn extract_func_cfgs(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { - info!("Executing agfj @@f on {:?}", self.file_path); + 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))?; + 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))?; + 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))?; + } + + // Extract the CFGs for each function + for function in self.get_function_name_list(r2p)? { + let function = FunctionToBeProcessed::from(function); + debug!( + "Extracting CFG for function {:?} @ {:?}", + function.name, function.addr + ); + match function.write_cfg_to_json(r2p, &output_dirpath, &self.func_filename_template) + { + Ok(()) => { + debug!( + "Successfully extracted CFG for function {:?} @ {:?}", + function.name, function.addr + ); + } + Err(e) => { + let error_path = function.get_output_filepath( + &output_dirpath, + &self.func_filename_template, + "error.log", + ); + error!( + "Failed to extract CFG for function {:?} @ {:?}: {}", + 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 { + info!("Error stored at {:?}", error_path); + } + continue; + } + } + } + } Ok(()) } @@ -1418,9 +1564,14 @@ impl FileToBeProcessed { Ok(()) } - pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, output_dirpath: &PathBuf, apply_mask: bool) -> 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)?; if !output_dirpath.is_dir() { @@ -1434,7 +1585,12 @@ impl FileToBeProcessed { "Function Name: {} Address: {} Size: {}", function.name, function.addr, function.size ); - function.write_to_bin(r2p, &output_dirpath, &self.func_filename_template, apply_mask)?; + function.write_to_bin( + r2p, + &output_dirpath, + &self.func_filename_template, + apply_mask, + )?; } info!("Function bytes successfully extracted"); diff --git a/src/main.rs b/src/main.rs index ab24a30..0d36fa7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -308,7 +308,7 @@ 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", + "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()), From d985eccbce41268d03ea8ddb0660957bf0c8cb62 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 13 Oct 2025 11:54:02 +0100 Subject: [PATCH 19/42] Function-wise error handling in bytes extraction --- src/extract.rs | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/src/extract.rs b/src/extract.rs index b53f23e..b8a0d90 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1434,7 +1434,36 @@ impl FileToBeProcessed { "Function Name: {} Address: {} Size: {}", function.name, function.addr, function.size ); - function.write_to_bin(r2p, &output_dirpath, &self.func_filename_template, apply_mask)?; + match function.write_to_bin( + r2p, + &output_dirpath, + &self.func_filename_template, + apply_mask, + ) { + Ok(()) => { + debug!( + "Successfully extracted bytes for function {:?} @ {:?}", + function.name, function.addr + ); + } + Err(e) => { + let error_path = function.get_output_filepath( + output_dirpath, + &self.func_filename_template, + "error.log", + ); + error!( + "Failed to extract bytes for function {:?} @ {:?}: {}", + 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 { + info!("Error stored at {:?}", error_path); + } + continue; + } + } } info!("Function bytes successfully extracted"); From 1b2f8112d8f755aed37d46fa31149859cb3537fb Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 13 Oct 2025 12:19:33 +0100 Subject: [PATCH 20/42] Fixed directory naming bug for bytes-masked mode --- src/extract.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index b8a0d90..e2d9d00 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -698,8 +698,9 @@ impl ExtractionJob { fn get_output_extension(job_type: &ExtractionJobType) -> Option<&str> { match job_type { - // Add here if output is not a JSON file (e.g. txt file or directory) - ExtractionJobType::FunctionBytes => None, // Output is a directory + // Add here if output is not a JSON file (None if a directory) + ExtractionJobType::FunctionBytes => None, + ExtractionJobType::FunctionBytesMasked => None, _ => Some("json"), } } From 05296c192a25f027d3013b8f9cdaa01360d723ae Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 13 Oct 2025 12:40:32 +0100 Subject: [PATCH 21/42] Skip extract log level changed to info --- src/extract.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index cb49056..fdfc6e6 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1152,13 +1152,13 @@ impl FileToBeProcessed { let error_path = output_path.with_extension("error.log"); if output_path.exists() { - debug!( + info!( "Skipping {:?} job for {:?}: already processed at {:?}.", job_type_suffix, self.file_path, output_path ); continue; } else if error_path.exists() && !self.retry_aborted { - debug!( + info!( "Skipping {:?} job for {:?}: already processed and failed. Error log at {:?}.", job_type_suffix, self.file_path, error_path ); From 7a848ea444f20188b72eba50e1fad02e30a8de50 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 13 Oct 2025 21:34:55 +0100 Subject: [PATCH 22/42] Added more control to radare2 spawn options --- src/extract.rs | 112 ++++++++++++++++++++++++++++++------------------- src/main.rs | 33 ++++++++++++++- 2 files changed, 102 insertions(+), 43 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index fdfc6e6..63c0763 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -119,11 +119,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, } @@ -585,8 +587,10 @@ 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, @@ -610,7 +614,9 @@ impl ExtractionJob { 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, }; @@ -657,7 +663,7 @@ impl ExtractionJob { file_path: PathBuf::from(f), output_path: output_path.to_owned(), job_types: extraction_job_types.clone(), - r2p_config: r2_handle_config, + r2p_config: r2_handle_config.clone(), with_annotations: *with_annotations, retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), @@ -1798,37 +1804,63 @@ impl FileToBeProcessed { 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"), None => R2Pipe::spawn(self.file_path.to_str().unwrap(), Some(opts)) .expect("Failed to spawn new R2Pipe"), }; - 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() { @@ -1848,22 +1880,18 @@ impl FileToBeProcessed { } 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() - ); - }; + 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) + .expect(&format!("Unable to complete analysis! ({analysis_mode})")); + debug!( + "'{}' r2 command complete for {}", + analysis_mode, + self.file_path.display() + ); } } diff --git a/src/main.rs b/src/main.rs index 0d36fa7..f22c110 100644 --- a/src/main.rs +++ b/src/main.rs @@ -326,9 +326,20 @@ 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( @@ -339,6 +350,7 @@ enum Commands { )] func_filename: String, + /// Timeout for Radare2 analysis in seconds #[arg(long)] timeout: Option, @@ -1123,7 +1135,10 @@ fn main() { num_threads, debug, extended_analysis, + experimental_analysis, use_curl_pdb, + apply_relocations, + disable_pseudo_asm, func_filename, timeout, with_annotations, @@ -1135,14 +1150,30 @@ 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, From b34a16d5fb9ae05b6e0f91c174e22260403afc54 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 13 Oct 2025 21:53:12 +0100 Subject: [PATCH 23/42] Improved error handling of bytes mask extraction --- src/extract.rs | 61 +++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 50 insertions(+), 11 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index e2d9d00..ce85aaf 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -746,10 +746,13 @@ impl FunctionToBeProcessed { filename_template: &str, apply_mask: bool, ) -> Result<()> { - let func_bytes = self.get_bytes(r2p, apply_mask).context("Failed to get function bytes")?; + let func_bytes = self + .get_bytes(r2p, apply_mask) + .context(format!("Failed to get function bytes for {:?} @ {:?}", self.name, self.addr))?; let bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); - let masked_bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); - + let masked_bytes_filepath = + self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); + info!("Writing function bytes to file: {:?}", bytes_filepath); std::fs::write(&bytes_filepath, func_bytes.bytes).with_context(|| { format!( @@ -759,8 +762,13 @@ impl FunctionToBeProcessed { })?; if apply_mask { - let bytes_mask = func_bytes.mask.context("Failed to get function masked bytes")?; - info!("Writing function masked bytes to file: {:?}", masked_bytes_filepath); + let bytes_mask = func_bytes + .mask + .context(format!("Failed to get function masked bytes for {:?} @ {:?}", self.name, self.addr))?; + info!( + "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: {:?}", @@ -815,11 +823,30 @@ impl FunctionToBeProcessed { let mut function_bytes_and_mask = r2p.cmd("p8fm")?; function_bytes_and_mask = function_bytes_and_mask.trim().to_string(); let parts: Vec<&str> = function_bytes_and_mask.split(":").collect(); - let function_bytes = hex::decode(parts[0]).context("Failed to decode hex function bytes")?; + + if parts.len() < 2 { + return Err(anyhow::anyhow!( + "Invalid p8fm output format: expected 'bytes:mask' but got '{}'", + function_bytes_and_mask + )); + } + + let function_bytes = + hex::decode(parts[0]).context("Failed to decode hex function bytes")?; let mut masked_bytes: Option> = None; if apply_mask { let bytes_mask = hex::decode(parts[1]).context("Failed to decode hex bytes mask")?; + + // Ensure function_bytes and bytes_mask have the same length + if function_bytes.len() != bytes_mask.len() { + return Err(anyhow::anyhow!( + "Function bytes length ({}) and mask length ({}) do not match", + function_bytes.len(), + bytes_mask.len() + )); + } + let mut masked_bytes_tmp = vec![0; function_bytes.len()]; for i in 0..function_bytes.len() { masked_bytes_tmp[i] = function_bytes[i] & bytes_mask[i]; @@ -827,7 +854,10 @@ impl FunctionToBeProcessed { masked_bytes = Some(masked_bytes_tmp); } - Ok(FuncBytes { bytes: function_bytes, mask: masked_bytes }) + Ok(FuncBytes { + bytes: function_bytes, + mask: masked_bytes, + }) } fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { @@ -1012,8 +1042,12 @@ impl FileToBeProcessed { ExtractionJobType::FunctionZignatures => { self.extract_function_zignatures(r2p, &tmp_output_path) } - ExtractionJobType::FunctionBytes => self.extract_function_bytes(r2p, &tmp_output_path, false), - ExtractionJobType::FunctionBytesMasked => self.extract_function_bytes(r2p, &tmp_output_path, true), + ExtractionJobType::FunctionBytes => { + 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 @@ -1419,9 +1453,14 @@ impl FileToBeProcessed { Ok(()) } - pub fn extract_function_bytes(&self, r2p: &mut R2Pipe, output_dirpath: &PathBuf, apply_mask: bool) -> 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)?; if !output_dirpath.is_dir() { From eba79ee0c7691338bdfcc4cc8ccc718888dce980 Mon Sep 17 00:00:00 2001 From: valbucci Date: Tue, 14 Oct 2025 12:05:11 +0100 Subject: [PATCH 24/42] Refactored bytes and bytes-masked extraction --- src/extract.rs | 69 ++++++++++++++++++++++++++++++++++---------------- 1 file changed, 47 insertions(+), 22 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index ce85aaf..d670bef 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -748,7 +748,7 @@ impl FunctionToBeProcessed { ) -> Result<()> { let func_bytes = self .get_bytes(r2p, apply_mask) - .context(format!("Failed to get function bytes for {:?} @ {:?}", self.name, self.addr))?; + .map_err(|e| anyhow::anyhow!("Failed to get bytes: {}", e))?; let bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); let masked_bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); @@ -764,7 +764,7 @@ impl FunctionToBeProcessed { if apply_mask { let bytes_mask = func_bytes .mask - .context(format!("Failed to get function masked bytes for {:?} @ {:?}", self.name, self.addr))?; + .context("Masked bytes missing from get_bytes output")?; info!( "Writing function masked bytes to file: {:?}", masked_bytes_filepath @@ -819,39 +819,64 @@ impl FunctionToBeProcessed { } fn get_bytes(&self, r2p: &mut R2Pipe, apply_mask: bool) -> Result { - FileToBeProcessed::go_to_address(r2p, self.addr)?; - let mut function_bytes_and_mask = r2p.cmd("p8fm")?; - function_bytes_and_mask = function_bytes_and_mask.trim().to_string(); - let parts: Vec<&str> = function_bytes_and_mask.split(":").collect(); - - if parts.len() < 2 { - return Err(anyhow::anyhow!( - "Invalid p8fm output format: expected 'bytes:mask' but got '{}'", - function_bytes_and_mask - )); - } - - let function_bytes = - hex::decode(parts[0]).context("Failed to decode hex function bytes")?; + let cmd_str; + let function_bytes; + let function_mask; + let function_bytes_and_mask; let mut masked_bytes: Option> = None; if apply_mask { - let bytes_mask = hex::decode(parts[1]).context("Failed to decode hex bytes mask")?; + cmd_str = format!("p8fm @ {}", self.addr); + function_bytes_and_mask = r2p + .cmd(cmd_str.as_str()) + .with_context(|| format!("Failed to execute `{}`", cmd_str))? + .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() != bytes_mask.len() { + if function_bytes.len() != function_mask.len() { return Err(anyhow::anyhow!( - "Function bytes length ({}) and mask length ({}) do not match", + "Function bytes length ({}) and mask length ({}) do not match.\n\ + Output from `{}`: '{}'", function_bytes.len(), - bytes_mask.len() + function_mask.len(), + cmd_str, + function_bytes_and_mask )); } let mut masked_bytes_tmp = vec![0; function_bytes.len()]; for i in 0..function_bytes.len() { - masked_bytes_tmp[i] = function_bytes[i] & bytes_mask[i]; + masked_bytes_tmp[i] = function_bytes[i] & function_mask[i]; } masked_bytes = Some(masked_bytes_tmp); + } else { + cmd_str = format!("p8f @ {}", self.addr); + let function_bytes_str = r2p + .cmd(cmd_str.as_str()) + .with_context(|| format!("Failed to execute `{}`", cmd_str))? + .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 { @@ -1499,7 +1524,7 @@ impl FileToBeProcessed { if let Err(write_err) = std::fs::write(&error_path, e.to_string()) { error!("Failed to write error to {:?}: {}", error_path, write_err); } else { - info!("Error stored at {:?}", error_path); + info!("Error stored in {:?}", error_path); } continue; } From ced68cfd60d7eb1f900c19149f51894f33c5ba7a Mon Sep 17 00:00:00 2001 From: valbucci Date: Tue, 14 Oct 2025 12:22:23 +0100 Subject: [PATCH 25/42] Added count of successful function bytes extraction --- src/extract.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index d670bef..546c814 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1487,12 +1487,13 @@ impl FileToBeProcessed { info!("Starting function bytes extraction"); let function_details = self.get_function_name_list(r2p)?; - + let functions_count = function_details.len(); if !output_dirpath.is_dir() { std::fs::create_dir_all(&output_dirpath) .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; } + let mut success_count: u32 = 0; for function_info in function_details { let function = FunctionToBeProcessed::from(function_info); debug!( @@ -1510,6 +1511,7 @@ impl FileToBeProcessed { "Successfully extracted bytes for function {:?} @ {:?}", function.name, function.addr ); + success_count += 1; } Err(e) => { let error_path = function.get_output_filepath( @@ -1531,8 +1533,13 @@ impl FileToBeProcessed { } } - info!("Function bytes successfully extracted"); - Ok(()) + let file_name = self.get_file_name()?; + if success_count > 0 { + info!("Bytes extracted for {}/{} functions in {:?}", success_count, functions_count, file_name); + Ok(()) + } else { + Err(anyhow::anyhow!("Failed to extract bytes for any function in {:?}", file_name)) + } } fn get_checksums(&self) -> Result { From 70fd2e5ea889c2f125e6dfbbddc0cead792e0932 Mon Sep 17 00:00:00 2001 From: valbucci Date: Tue, 14 Oct 2025 14:33:10 +0100 Subject: [PATCH 26/42] Cargo fmt --- src/extract.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 6c6c3f7..ef6f44f 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1671,10 +1671,16 @@ impl FileToBeProcessed { let file_name = self.get_file_name()?; if success_count > 0 { - info!("Bytes extracted for {}/{} functions in {:?}", success_count, functions_count, file_name); + info!( + "Bytes extracted for {}/{} functions in {:?}", + success_count, functions_count, file_name + ); Ok(()) } else { - Err(anyhow::anyhow!("Failed to extract bytes for any function in {:?}", file_name)) + Err(anyhow::anyhow!( + "Failed to extract bytes for any function in {:?}", + file_name + )) } } From f3f494d610f903807e07383ca2d05541e6d772d0 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 16 Oct 2025 12:47:35 +0100 Subject: [PATCH 27/42] Better handling for broken R2Pipe --- src/extract.rs | 82 ++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 70 insertions(+), 12 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index ef6f44f..2d85251 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1215,11 +1215,18 @@ impl FileToBeProcessed { } // 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!( @@ -1854,7 +1861,7 @@ 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"); @@ -1912,9 +1919,10 @@ impl FileToBeProcessed { 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 self.r2p_config.use_curl_pdb { @@ -1926,16 +1934,19 @@ 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) { + fn analyse_r2_pipe(&self, r2p: &mut R2Pipe) -> Result<()> { let analysis_mode = self.r2p_config.analysis_mode.as_str(); debug!( "Executing '{}' r2 command for {}", @@ -1943,11 +1954,58 @@ impl FileToBeProcessed { self.file_path.display() ); r2p.cmd(analysis_mode) - .expect(&format!("Unable to complete analysis! ({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 + ); } } From b750aa84e9534980a341a031becbdabf6574b0aa Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 16 Oct 2025 18:30:36 +0100 Subject: [PATCH 28/42] Removed go_to_address for improved efficiency --- src/extract.rs | 127 +++++++++++++++++++++---------------------------- 1 file changed, 55 insertions(+), 72 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 2d85251..1584ae2 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -406,6 +406,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, @@ -417,6 +426,7 @@ pub struct BasicBlockMetadataEntry { pub outputs: u64, pub ninstr: u64, pub instrs: Vec, + #[serde(deserialize_with = "deserialize_bool_from_int")] pub traced: bool, } @@ -886,9 +896,13 @@ impl FunctionToBeProcessed { 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.as_str()) - .with_context(|| format!("Failed to execute `{}`", cmd_str))? + .cmd(&cmd_str) + .context("Failed to execute `{}`")? .trim() .to_string(); let parts: Vec<&str> = function_bytes_and_mask.split(":").collect(); @@ -908,6 +922,7 @@ impl FunctionToBeProcessed { // Ensure function_bytes and bytes_mask have the same length if function_bytes.len() != function_mask.len() { + // TODO: Iteratively identify and fix missing bytes in the mask return Err(anyhow::anyhow!( "Function bytes length ({}) and mask length ({}) do not match.\n\ Output from `{}`: '{}'", @@ -925,9 +940,10 @@ impl FunctionToBeProcessed { 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.as_str()) - .with_context(|| format!("Failed to execute `{}`", cmd_str))? + .cmd(&cmd_str) + .context("Failed to execute `{}`")? .trim() .to_string(); function_bytes = hex::decode(function_bytes_str.clone()).with_context(|| { @@ -945,36 +961,12 @@ impl FunctionToBeProcessed { } fn get_basic_block_info(&self, r2p: &mut R2Pipe) -> Result { - info!( - "Getting the basic block information for function @ {}", - self.addr - ); - FileToBeProcessed::go_to_address(r2p, self.addr)?; - - let json = r2p.cmd("afbj").context("Command afbj failed")?; - // 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))?; - - // 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); - } - } - } - } + 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")?; - // Deserialize JSON into a BasicBlockInfo struct - let bb_info: BasicBlockInfo = serde_json::from_value(value.clone()).with_context(|| { - format!( - "Unable to convert {:?} into a BasicBlockInfo struct!", - value - ) - })?; + let bb_info: BasicBlockInfo = + serde_json::from_value(value).context("Unable to convert to BasicBlockInfo struct!")?; Ok(bb_info) } @@ -986,27 +978,25 @@ impl FunctionToBeProcessed { "Getting local variable xref details for function @ {}", self.addr ); - FileToBeProcessed::go_to_address(r2p, self.addr)?; - let json = r2p.cmd("axvj").context("Command axvj failed")?; + 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_str(&json).with_context(|| { - format!("Unable to convert {:?} to LocalVariableXrefs struct!", json) - })?; + 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); - FileToBeProcessed::go_to_address(r2p, self.addr)?; - let json = r2p.cmd("axffj").context("Command axffj failed")?; - let mut json_obj: Vec = - serde_json::from_str(&json).with_context(|| { - format!( - "Unable to convert {:?} to FunctionXrefDetails struct!", - json - ) - })?; + 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. @@ -1028,19 +1018,20 @@ impl FunctionToBeProcessed { r2p: &mut R2Pipe, with_annotations: bool, ) -> Result { - FileToBeProcessed::go_to_address(r2p, self.addr)?; - let json = r2p.cmd("pdgj").context("Command pdgj failed")?; + 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_str(&json) - .with_context(|| format!("Unable to convert {:?} to DecompJSON struct!", json))?; + 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_str(&json) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; + let json_obj: Value = + serde_json::from_value(json).context("Unable to convert to JSON object!")?; let parsed_code = json_obj["code"] .as_str() - .with_context(|| format!("Unable to get code from {:?}!", json))? + .context("Unable to get code from JSON object!")? .to_string(); let parsed_obj = DecompJSON { code: parsed_code, @@ -1051,12 +1042,9 @@ impl FunctionToBeProcessed { } fn get_cfg(&self, r2p: &mut R2Pipe) -> Result { - FileToBeProcessed::go_to_address(r2p, self.addr)?; - debug!("Getting CFG for function @ {}", self.addr); - // let json_str = r2p.cmd("agfj").context("Command agfj failed")?; - let json = r2p.cmdj("agfj").context("Command agfj failed")?; - // let cfg: AGFJFunc = serde_json::from_str(&json_str) - // .with_context(|| format!("Unable to convert {:?} to AGFJFunc struct!", json_str))?; + 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()) @@ -1744,12 +1732,12 @@ impl FileToBeProcessed { num_instructons: u64, r2p: &mut R2Pipe, ) -> Result { - Self::go_to_address(r2p, address)?; - let pcode_ret = r2p.cmd(format!("pdgsd {}", num_instructons).as_str())?; + 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()); @@ -1785,7 +1773,9 @@ impl FileToBeProcessed { 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(()) @@ -1836,13 +1826,6 @@ impl FileToBeProcessed { Ok(()) } - /// Seeks to the function address - fn go_to_address(r2p: &mut R2Pipe, address: u64) -> Result<(), Error> { - r2p.cmd(format!("s {}", address).as_str()) - .with_context(|| format!("Failed to seek address {:x}", address))?; - Ok(()) - } - /// Downloads the PDB file if available fn handle_symbols_pdb(&self, r2p: &mut R2Pipe) -> Result<(), Error> { // Download symbols if available From d48bfeb0c5647121a0737f8a57bd139015e8c883 Mon Sep 17 00:00:00 2001 From: valbucci Date: Thu, 16 Oct 2025 20:15:50 +0100 Subject: [PATCH 29/42] Lazy initialisation of functions list --- Cargo.toml | 2 + src/extract.rs | 132 +++++++++++++++++++++++++++++++++++-------------- 2 files changed, 96 insertions(+), 38 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f5cd8c0..7f0a364 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,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/extract.rs b/src/extract.rs index 1584ae2..ab16875 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -31,6 +31,7 @@ use std::io::{BufReader, BufWriter, Read}; use std::path::PathBuf; use std::string::String; use std::sync::LazyLock; +use once_cell::sync::OnceCell; use walkdir::WalkDir; #[derive(PartialEq, Debug)] @@ -100,14 +101,16 @@ pub struct FileToBeProcessed { pub with_annotations: bool, pub retry_aborted: bool, pub func_filename_template: String, + pub function_list: OnceCell>, } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct FunctionToBeProcessed { pub name: String, pub addr: u64, pub size: u64, pub ninstrs: u64, + pub nblocks: u64, } #[derive(Debug)] @@ -143,6 +146,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, @@ -150,6 +154,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, @@ -158,15 +164,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, + #[serde(alias = "minaddr")] pub minbound: i64, + #[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, @@ -309,17 +326,19 @@ impl with_annotations: orig.4, retry_aborted: orig.5, func_filename_template: orig.6, + function_list: OnceCell::new(), } } } -impl From<(String, u64, u64, u64)> for FunctionToBeProcessed { - fn from(orig: (String, u64, u64, u64)) -> Self { +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, } } } @@ -330,7 +349,8 @@ impl From for FunctionToBeProcessed { name: func_info.name, addr: func_info.offset, size: func_info.size, - ninstrs: func_info.ninstrs as u64, + ninstrs: func_info.ninstrs, + nblocks: func_info.nbbs, } } } @@ -342,6 +362,7 @@ impl From for FunctionToBeProcessed { addr: func_details.offset, size: func_details.size, ninstrs: func_details.ninstrs, + nblocks: func_details.nbbs, } } } @@ -353,6 +374,7 @@ impl From for FunctionToBeProcessed { addr: func.offset, size: func.size, ninstrs: func.ninstr, + nblocks: func.blocks.len() as u64, } } } @@ -645,6 +667,7 @@ impl ExtractionJob { with_annotations: *with_annotations, retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), + function_list: OnceCell::new(), }; Ok(ExtractionJob { @@ -677,6 +700,7 @@ impl ExtractionJob { with_annotations: *with_annotations, retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), + function_list: OnceCell::new(), }) .collect(); @@ -1157,6 +1181,34 @@ impl FileToBeProcessed { Ok(()) } + pub fn write_function_index(&self, functions: &Vec, output_dirpath: &PathBuf, ext: &str) -> Result<()> { + let file = File::create(output_dirpath.join(".func-index.csv"))?; + 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(()) + } + + pub fn get_function_list(&self, r2p: &mut R2Pipe) -> Result<&[FunctionToBeProcessed]> { + self.function_list + .get_or_try_init(|| self.setup_function_list(r2p)) + .map(|v| v.as_slice()) + } + pub fn process_all_modes(&self) { info!( "Starting extraction for {} job types on {:?}", @@ -1303,7 +1355,7 @@ impl FileToBeProcessed { r2p: &mut R2Pipe, 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() { @@ -1346,10 +1398,7 @@ impl FileToBeProcessed { pub fn extract_function_info(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting function metdata extraction"); - let function_details: Vec = self.get_function_name_list(r2p)?; - - info!("Writing extracted data to file"); - let json = json!(function_details); + let json = r2p.cmdj("aflj").with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; self.write_to_json(&json, output_path) .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; Ok(()) @@ -1361,7 +1410,7 @@ impl FileToBeProcessed { 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() { @@ -1411,8 +1460,7 @@ impl FileToBeProcessed { } // Extract the CFGs for each function - for function in self.get_function_name_list(r2p)? { - let function = FunctionToBeProcessed::from(function); + for function in self.get_function_list(r2p)? { debug!( "Extracting CFG for function {:?} @ {:?}", function.name, function.addr @@ -1449,12 +1497,11 @@ impl FileToBeProcessed { } pub fn extract_function_xrefs(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let mut function_xrefs: HashMap> = HashMap::new(); info!("Extracting xrefs for each function"); - for function_info in function_details { - let function = FunctionToBeProcessed::from(function_info); + for function in function_details { let ret = function.get_xref_details(r2p).with_context(|| { format!("Unable to get function xrefs from {:?}", self.file_path) })?; @@ -1467,11 +1514,10 @@ impl FileToBeProcessed { 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_info in function_details { - let function = FunctionToBeProcessed::from(function_info); + for function in function_details { let ret = function .get_ghidra_decomp(r2p, self.with_annotations) .with_context(|| { @@ -1491,11 +1537,11 @@ impl FileToBeProcessed { 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.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(), @@ -1512,11 +1558,10 @@ impl FileToBeProcessed { 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_info in function_details { - let function = FunctionToBeProcessed::from(function_info); + for function in function_details { let bb_info = function.get_basic_block_info(r2p).with_context(|| { format!( "Unable to get basic block addresses in {:?} @ {:?}", @@ -1559,11 +1604,10 @@ impl FileToBeProcessed { 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_info in function_details { - let function = FunctionToBeProcessed::from(function_info); + for function in function_details { let ret = function.get_local_variable_xref_details(r2p)?; function_local_variable_xrefs.insert(function.name.clone(), ret); } @@ -1617,7 +1661,7 @@ impl FileToBeProcessed { ) -> Result<()> { info!("Starting function bytes extraction"); - let function_details = self.get_function_name_list(r2p)?; + let function_details = self.get_function_list(r2p)?; let functions_count = function_details.len(); if !output_dirpath.is_dir() { std::fs::create_dir_all(&output_dirpath) @@ -1625,8 +1669,7 @@ impl FileToBeProcessed { } let mut success_count: u32 = 0; - for function_info in function_details { - let function = FunctionToBeProcessed::from(function_info); + for function in function_details { debug!( "Function Name: {} Address: {} Size: {}", function.name, function.addr, function.size @@ -1752,18 +1795,31 @@ impl FileToBeProcessed { }) } - fn get_function_name_list( + fn setup_function_list( &self, r2p: &mut R2Pipe, - ) -> Result, anyhow::Error> { - info!("Getting function information from binary"); + ) -> Result> { + info!("Setting up function list for {:?} ...", self.get_file_name()); let json = r2p - .cmd("aflj") + .cmdj("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) + + let array = match json { + serde_json::Value::Array(arr) => arr, + _ => return Err(anyhow::anyhow!("aflj output is not an array for {:?}", self.file_path)), + }; + + let functions_to_be_processed: Vec = array + .into_iter() + .map(|func_json| { + serde_json::from_value::(func_json) + .map(FunctionToBeProcessed::from) + }) + .collect::, _>>() + .with_context(|| format!("Failed to parse aflj function details for {:?}", self.file_path))?; + + debug!("Done getting function list for {:?}", self.get_file_name()); + Ok(functions_to_be_processed) } // Helper Functions From cd7fb7ed06783a3fbb82374f4af0655da9912ab8 Mon Sep 17 00:00:00 2001 From: valbucci Date: Fri, 17 Oct 2025 10:02:51 +0100 Subject: [PATCH 30/42] Function index for func-cfg and bytes extract modes --- src/extract.rs | 40 +++++++++++++++++++++++++++++++++------- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index ab16875..0a8d681 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -873,6 +873,16 @@ impl FunctionToBeProcessed { } // 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 => "bin", + ExtractionJobType::FunctionBytesMasked => "masked.bin", + ExtractionJobType::FunctionCFG => "json", + _ => "", + } + } + fn get_hex_address(&self) -> String { format!("{:x}", self.addr) .trim_start_matches("0x") @@ -890,8 +900,9 @@ impl FunctionToBeProcessed { }; func_filename = sanitize_filename(&func_filename); - if ["symbol", "address"].contains(&template) { + 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 @@ -1181,8 +1192,12 @@ impl FileToBeProcessed { Ok(()) } - pub fn write_function_index(&self, functions: &Vec, output_dirpath: &PathBuf, ext: &str) -> Result<()> { - let file = File::create(output_dirpath.join(".func-index.csv"))?; + 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"])?; @@ -1459,8 +1474,12 @@ impl FileToBeProcessed { .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; } + let function_list = self.get_function_list(r2p)?.to_vec(); + // Write index for the CFGs + self.write_function_index(&function_list, &output_dirpath, ExtractionJobType::FunctionCFG)?; + // Extract the CFGs for each function - for function in self.get_function_list(r2p)? { + for function in function_list { debug!( "Extracting CFG for function {:?} @ {:?}", function.name, function.addr @@ -1661,15 +1680,22 @@ impl FileToBeProcessed { ) -> Result<()> { info!("Starting function bytes extraction"); - let function_details = self.get_function_list(r2p)?; - let functions_count = function_details.len(); + let functions = self.get_function_list(r2p)?; + let functions_count = functions.len(); 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 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)?; + } + let mut success_count: u32 = 0; - for function in function_details { + for function in functions { debug!( "Function Name: {} Address: {} Size: {}", function.name, function.addr, function.size From d90a25ebef34aa1825410414b835f45083beb474 Mon Sep 17 00:00:00 2001 From: valbucci Date: Fri, 17 Oct 2025 10:42:25 +0100 Subject: [PATCH 31/42] Filter functions based on minimum number of blocks --- src/extract.rs | 99 ++++++++++++++++++++++++++++++++++++++++---------- src/main.rs | 6 +++ 2 files changed, 86 insertions(+), 19 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 0a8d681..05e6958 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -17,6 +17,7 @@ use serde_json; use glob::glob; use md5; +use once_cell::sync::OnceCell; use serde::ser::{SerializeSeq, Serializer}; use serde_json::{json, value::RawValue, Deserializer, Value}; use sha1; @@ -31,7 +32,6 @@ use std::io::{BufReader, BufWriter, Read}; use std::path::PathBuf; use std::string::String; use std::sync::LazyLock; -use once_cell::sync::OnceCell; use walkdir::WalkDir; #[derive(PartialEq, Debug)] @@ -102,6 +102,7 @@ pub struct FileToBeProcessed { pub retry_aborted: bool, pub func_filename_template: String, pub function_list: OnceCell>, + pub min_basic_blocks: Option, } #[derive(Debug, Clone)] @@ -305,6 +306,7 @@ impl bool, bool, String, + Option, )> for FileToBeProcessed { fn from( @@ -316,6 +318,7 @@ impl bool, bool, String, + Option, ), ) -> FileToBeProcessed { FileToBeProcessed { @@ -327,6 +330,7 @@ impl retry_aborted: orig.5, func_filename_template: orig.6, function_list: OnceCell::new(), + min_basic_blocks: orig.7, } } } @@ -627,6 +631,7 @@ impl ExtractionJob { timeout: &Option, with_annotations: &bool, retry_aborted: &bool, + min_basic_blocks: &Option, ) -> Result { let mut job_types = vec![]; let mut extraction_job_types = vec![]; @@ -668,6 +673,7 @@ impl ExtractionJob { retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), function_list: OnceCell::new(), + min_basic_blocks: *min_basic_blocks, }; Ok(ExtractionJob { @@ -701,6 +707,7 @@ impl ExtractionJob { retry_aborted: *retry_aborted, func_filename_template: func_filename_template.to_string(), function_list: OnceCell::new(), + min_basic_blocks: *min_basic_blocks, }) .collect(); @@ -1192,7 +1199,12 @@ impl FileToBeProcessed { Ok(()) } - pub fn write_function_index(&self, functions: &Vec, output_dirpath: &PathBuf, job_type: ExtractionJobType) -> Result<()> { + 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)); @@ -1200,11 +1212,19 @@ impl FileToBeProcessed { 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"])?; - + 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); + let output_path = + function.get_output_filepath(output_dirpath, &self.func_filename_template, ext); writer.write_record(&[ &function.name, &function.addr.to_string(), @@ -1413,7 +1433,9 @@ impl FileToBeProcessed { pub fn extract_function_info(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { info!("Starting function metdata extraction"); - let json = r2p.cmdj("aflj").with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; + let json = r2p + .cmdj("aflj") + .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; self.write_to_json(&json, output_path) .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; Ok(()) @@ -1476,7 +1498,11 @@ impl FileToBeProcessed { let function_list = self.get_function_list(r2p)?.to_vec(); // Write index for the CFGs - self.write_function_index(&function_list, &output_dirpath, ExtractionJobType::FunctionCFG)?; + self.write_function_index( + &function_list, + &output_dirpath, + ExtractionJobType::FunctionCFG, + )?; // Extract the CFGs for each function for function in function_list { @@ -1688,10 +1714,18 @@ impl FileToBeProcessed { } // Write index for the raw bytes - self.write_function_index(&functions.to_vec(), &output_dirpath, ExtractionJobType::FunctionBytes)?; + 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)?; + self.write_function_index( + &functions.to_vec(), + &output_dirpath, + ExtractionJobType::FunctionBytesMasked, + )?; } let mut success_count: u32 = 0; @@ -1821,30 +1855,57 @@ impl FileToBeProcessed { }) } - fn setup_function_list( - &self, - r2p: &mut R2Pipe, - ) -> Result> { - info!("Setting up function list for {:?} ...", self.get_file_name()); + fn setup_function_list(&self, r2p: &mut R2Pipe) -> Result> { + info!( + "Setting up function list for {:?} ...", + self.get_file_name() + ); let json = r2p .cmdj("aflj") .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; - + let array = match json { serde_json::Value::Array(arr) => arr, - _ => return Err(anyhow::anyhow!("aflj output is not an array for {:?}", self.file_path)), + _ => { + return Err(anyhow::anyhow!( + "aflj output is not an array for {:?}", + self.file_path + )) + } }; - + + if let Some(min_blocks) = self.min_basic_blocks { + info!( + "Filtering functions with less than {} basic blocks", + min_blocks + ); + } else { + info!("No basic blocks filter applied"); + } + 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(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::, _>>() - .with_context(|| format!("Failed to parse aflj function details for {:?}", self.file_path))?; - + .with_context(|| { + format!( + "Failed to parse aflj function details for {:?}", + self.file_path + ) + })?; debug!("Done getting function list for {:?}", self.get_file_name()); + Ok(functions_to_be_processed) } diff --git a/src/main.rs b/src/main.rs index f22c110..808d9d9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -360,6 +360,10 @@ enum Commands { /// 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 /// @@ -1143,6 +1147,7 @@ fn main() { timeout, with_annotations, retry_aborted, + min_basic_blocks, } => { info!("Creating extraction job with {} modes", modes.len()); if !output_dir.exists() { @@ -1178,6 +1183,7 @@ fn main() { timeout, with_annotations, retry_aborted, + min_basic_blocks, ) .unwrap_or_else(|e| { error!("Failed to create extraction job: {}", e); From f995d31c4235f88e4c96eb72cb0ff89220d6a4d8 Mon Sep 17 00:00:00 2001 From: valbucci Date: Fri, 17 Oct 2025 23:09:09 +0100 Subject: [PATCH 32/42] Tidier bytes-masked extraction and keep_raw_bytes option --- src/extract.rs | 97 +++++++++++++++++++++++++++++++------------------- src/main.rs | 5 +++ 2 files changed, 66 insertions(+), 36 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 05e6958..62af1a5 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -99,6 +99,7 @@ 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>, @@ -305,6 +306,7 @@ impl R2PipeConfig, bool, bool, + bool, String, Option, )> for FileToBeProcessed @@ -317,6 +319,7 @@ impl R2PipeConfig, bool, bool, + bool, String, Option, ), @@ -327,10 +330,11 @@ impl job_types: orig.2, r2p_config: orig.3, with_annotations: orig.4, - retry_aborted: orig.5, - func_filename_template: orig.6, + keep_raw_bytes: orig.5, + retry_aborted: orig.6, + func_filename_template: orig.7, function_list: OnceCell::new(), - min_basic_blocks: orig.7, + min_basic_blocks: orig.8, } } } @@ -630,6 +634,7 @@ impl ExtractionJob { func_filename_template: &str, timeout: &Option, with_annotations: &bool, + keep_raw_bytes: &bool, retry_aborted: &bool, min_basic_blocks: &Option, ) -> Result { @@ -640,13 +645,16 @@ impl ExtractionJob { let job_type = Self::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 - ); - } + // 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); + } + 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 { @@ -670,6 +678,7 @@ impl ExtractionJob { 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(), @@ -704,6 +713,7 @@ impl ExtractionJob { 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(), @@ -809,27 +819,38 @@ impl FunctionToBeProcessed { 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))?; - let bytes_filepath = self.get_output_filepath(output_dirpath, filename_template, "bin"); - let masked_bytes_filepath = - self.get_output_filepath(output_dirpath, filename_template, "masked.bin"); - - 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 - ) - })?; + + // 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")?; - info!( + + // 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 ); @@ -883,7 +904,7 @@ impl FunctionToBeProcessed { 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 => "bin", + ExtractionJobType::FunctionBytes => "raw.bin", ExtractionJobType::FunctionBytesMasked => "masked.bin", ExtractionJobType::FunctionCFG => "json", _ => "", @@ -965,21 +986,22 @@ impl FunctionToBeProcessed { // Ensure function_bytes and bytes_mask have the same length if function_bytes.len() != function_mask.len() { // TODO: Iteratively identify and fix missing bytes in the mask - return Err(anyhow::anyhow!( + error!( "Function bytes length ({}) and mask length ({}) do not match.\n\ Output from `{}`: '{}'", function_bytes.len(), function_mask.len(), cmd_str, function_bytes_and_mask - )); - } - - 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 = 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); } - masked_bytes = Some(masked_bytes_tmp); } else { cmd_str = format!("p8f @ {}", self.addr); debug!("Getting function bytes for function: `{}`", cmd_str); @@ -1713,12 +1735,14 @@ impl FileToBeProcessed { .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; } - // Write index for the raw bytes - self.write_function_index( - &functions.to_vec(), - &output_dirpath, - ExtractionJobType::FunctionBytes, - )?; + 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( @@ -1739,6 +1763,7 @@ impl FileToBeProcessed { &output_dirpath, &self.func_filename_template, apply_mask, + self.keep_raw_bytes, ) { Ok(()) => { debug!( diff --git a/src/main.rs b/src/main.rs index 808d9d9..ef43464 100644 --- a/src/main.rs +++ b/src/main.rs @@ -357,6 +357,9 @@ enum Commands { #[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, @@ -1146,6 +1149,7 @@ fn main() { func_filename, timeout, with_annotations, + keep_raw_bytes, retry_aborted, min_basic_blocks, } => { @@ -1182,6 +1186,7 @@ fn main() { func_filename, timeout, with_annotations, + keep_raw_bytes, retry_aborted, min_basic_blocks, ) From 3ab5104655a14877f006ce302cdb575e14f83842 Mon Sep 17 00:00:00 2001 From: valbucci Date: Sat, 18 Oct 2025 00:44:54 +0100 Subject: [PATCH 33/42] Fixup mismatched function bytes mask --- src/extract.rs | 157 +++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 145 insertions(+), 12 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 62af1a5..849d9f2 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -650,11 +650,19 @@ impl ExtractionJob { // 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); + warn!( + "Annotations are only supported for decompilation extraction (mode: {})", + mode + ); } - if !extraction_job_types.contains(&ExtractionJobType::FunctionBytesMasked) && *keep_raw_bytes { + 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); + warn!( + "Keep raw bytes is only supported for masked bytes extraction (mode: {})", + mode + ); } let r2_handle_config = R2PipeConfig { @@ -829,8 +837,10 @@ impl FunctionToBeProcessed { // 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); + 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(|| { @@ -847,8 +857,11 @@ impl FunctionToBeProcessed { .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); + 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: {:?}", @@ -950,10 +963,125 @@ impl FunctionToBeProcessed { 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 function_mask; + let mut function_mask; let function_bytes_and_mask; let mut masked_bytes: Option> = None; @@ -985,13 +1113,18 @@ impl FunctionToBeProcessed { // Ensure function_bytes and bytes_mask have the same length if function_bytes.len() != function_mask.len() { - // TODO: Iteratively identify and fix missing bytes in the mask + 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!( - "Function bytes length ({}) and mask length ({}) do not match.\n\ - Output from `{}`: '{}'", + "After fix-up, lengths still differ: bytes={} mask={}. Original p8fm was '{}'", function_bytes.len(), function_mask.len(), - cmd_str, function_bytes_and_mask ); masked_bytes = None; From dc966149e5992bc7a7c6c30e92a15f6ec155eb7c Mon Sep 17 00:00:00 2001 From: valbucci Date: Sat, 18 Oct 2025 09:51:30 +0100 Subject: [PATCH 34/42] Changed minbound from i64 to u64 --- src/extract.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/extract.rs b/src/extract.rs index 849d9f2..0e5e1d6 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -176,7 +176,7 @@ pub struct AFLJFuncDetails { #[serde(default)] pub signature: String, #[serde(alias = "minaddr")] - pub minbound: i64, + pub minbound: u64, #[serde(alias = "maxaddr")] pub maxbound: u64, #[serde(default)] From 924c046f52fe559613d6d4c6dc6180a29f4c10be Mon Sep 17 00:00:00 2001 From: valbucci Date: Sat, 18 Oct 2025 22:30:00 +0100 Subject: [PATCH 35/42] Removed wasteful aflj parsing --- src/extract.rs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 0e5e1d6..b6f04f1 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1587,12 +1587,15 @@ impl FileToBeProcessed { } pub fn extract_function_info(&self, r2p: &mut R2Pipe, output_path: &PathBuf) -> Result<()> { - info!("Starting function metdata extraction"); - let json = r2p - .cmdj("aflj") + info!("Starting function info extraction"); + let json_raw = r2p + .cmd("aflj") .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; - self.write_to_json(&json, output_path) - .with_context(|| format!("Unable to convert {:?} to JSON object!", json))?; + + 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); Ok(()) } From 4ca6aa7ce8387d2f7249cdee5c2c4bfc5ac10c31 Mon Sep 17 00:00:00 2001 From: valbucci Date: Sat, 18 Oct 2025 18:36:31 +0100 Subject: [PATCH 36/42] Warning instead of error when function list is empty --- src/extract.rs | 34 ++++++++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index b6f04f1..e2237b6 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1655,13 +1655,22 @@ impl FileToBeProcessed { } let function_list = self.get_function_list(r2p)?.to_vec(); + let functions_count = function_list.len(); // Write index for the CFGs self.write_function_index( &function_list, &output_dirpath, ExtractionJobType::FunctionCFG, )?; + if functions_count == 0 { + warn!( + "No functions to be processed in {:?}. Skipping CFG extraction.", + self.file_path + ); + return Ok(()); + } + let mut success_count: u32 = 0; // Extract the CFGs for each function for function in function_list { debug!( @@ -1675,6 +1684,7 @@ impl FileToBeProcessed { "Successfully extracted CFG for function {:?} @ {:?}", function.name, function.addr ); + success_count += 1; } Err(e) => { let error_path = function.get_output_filepath( @@ -1695,6 +1705,18 @@ impl FileToBeProcessed { } } } + + if success_count > 0 { + info!( + "CFGs extracted for {}/{} functions in {:?}", + success_count, functions_count, self.file_path + ); + } else { + return Err(anyhow::anyhow!( + "Failed to extract CFGs for all functions in {:?}", + self.file_path + )); + } } Ok(()) } @@ -1864,6 +1886,7 @@ impl FileToBeProcessed { ) -> Result<()> { info!("Starting function bytes extraction"); + let file_name = self.get_file_name()?; let functions = self.get_function_list(r2p)?; let functions_count = functions.len(); if !output_dirpath.is_dir() { @@ -1888,6 +1911,14 @@ impl FileToBeProcessed { )?; } + if functions.len() == 0 { + warn!( + "No functions to be processed in {:?}. Skipping bytes extraction.", + file_name + ); + return Ok(()); + } + let mut success_count: u32 = 0; for function in functions { debug!( @@ -1928,7 +1959,6 @@ impl FileToBeProcessed { } } - let file_name = self.get_file_name()?; if success_count > 0 { info!( "Bytes extracted for {}/{} functions in {:?}", @@ -1937,7 +1967,7 @@ impl FileToBeProcessed { Ok(()) } else { Err(anyhow::anyhow!( - "Failed to extract bytes for any function in {:?}", + "Failed to extract bytes for all functions in {:?}", file_name )) } From 6cc92ebf44e5d41f61f088cd8901277811def489 Mon Sep 17 00:00:00 2001 From: valbucci Date: Sat, 18 Oct 2025 19:01:00 +0100 Subject: [PATCH 37/42] Fixed bug related to error log file names --- src/extract.rs | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/src/extract.rs b/src/extract.rs index e2237b6..98cf8a3 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1285,6 +1285,28 @@ impl FileToBeProcessed { Ok(output_filename) } + /// 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) + } + + /// 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()); @@ -1428,7 +1450,13 @@ impl FileToBeProcessed { continue; } }; - let error_path = output_path.with_extension("error.log"); + 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!( From de76fb82240bd0f2b7f38150b04dc417d94dd69b Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 20 Oct 2025 12:52:40 +0100 Subject: [PATCH 38/42] Resume function-level data extraction --- src/extract.rs | 347 ++++++++++++++++++++++++++++++++++--------------- 1 file changed, 240 insertions(+), 107 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 98cf8a3..199a4c0 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -1655,6 +1655,176 @@ impl FileToBeProcessed { Ok(()) } + /// 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(()); + } + + // Phase 1: Scan for existing files to enable resume functionality + struct FileStatus { + required_files_exist: Vec, // One per job_type + error_exists: bool, + } + + // 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; + } + } + } + + // 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, @@ -1682,69 +1852,43 @@ impl FileToBeProcessed { .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(); - let functions_count = function_list.len(); - // Write index for the CFGs self.write_function_index( &function_list, &output_dirpath, ExtractionJobType::FunctionCFG, )?; - if functions_count == 0 { - warn!( - "No functions to be processed in {:?}. Skipping CFG extraction.", - self.file_path - ); - return Ok(()); - } - let mut success_count: u32 = 0; - // Extract the CFGs for each function - for function in function_list { - debug!( - "Extracting CFG for function {:?} @ {:?}", - function.name, function.addr - ); - match function.write_cfg_to_json(r2p, &output_dirpath, &self.func_filename_template) - { - Ok(()) => { + // 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 {:?} @ {:?}", + "Successfully extracted CFG for function {:?} @ {:#x}", function.name, function.addr ); - success_count += 1; - } - Err(e) => { - let error_path = function.get_output_filepath( - &output_dirpath, - &self.func_filename_template, - "error.log", - ); - error!( - "Failed to extract CFG for function {:?} @ {:?}: {}", - 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 { - info!("Error stored at {:?}", error_path); - } - continue; } - } - } + result + }, + )?; - if success_count > 0 { - info!( - "CFGs extracted for {}/{} functions in {:?}", - success_count, functions_count, self.file_path - ); - } else { - return Err(anyhow::anyhow!( - "Failed to extract CFGs for all functions in {:?}", - self.file_path - )); - } + info!( + "Finished extracting CFGs for each function of {:?}", + self.file_path + ); } Ok(()) } @@ -1916,7 +2060,6 @@ impl FileToBeProcessed { let file_name = self.get_file_name()?; let functions = self.get_function_list(r2p)?; - let functions_count = functions.len(); if !output_dirpath.is_dir() { std::fs::create_dir_all(&output_dirpath) .with_context(|| format!("Failed to create directory {:?}", output_dirpath))?; @@ -1939,66 +2082,56 @@ impl FileToBeProcessed { )?; } - if functions.len() == 0 { - warn!( - "No functions to be processed in {:?}. Skipping bytes extraction.", - file_name - ); - return Ok(()); - } + // Determine which job types need to be extracted + let write_raw = !apply_mask || self.keep_raw_bytes; + let write_masked = apply_mask; - let mut success_count: u32 = 0; - for function in functions { - debug!( - "Function Name: {} Address: {} Size: {}", - function.name, function.addr, function.size - ); - match function.write_to_bin( - r2p, - &output_dirpath, - &self.func_filename_template, - apply_mask, - self.keep_raw_bytes, - ) { - Ok(()) => { + 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 ); - success_count += 1; - } - Err(e) => { - let error_path = function.get_output_filepath( - output_dirpath, - &self.func_filename_template, - "error.log", - ); - error!( - "Failed to extract bytes for function {:?} @ {:?}: {}", - 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 { - info!("Error stored in {:?}", error_path); - } - continue; } - } - } + result + }, + )?; - if success_count > 0 { - info!( - "Bytes extracted for {}/{} functions in {:?}", - success_count, functions_count, file_name - ); - Ok(()) - } else { - Err(anyhow::anyhow!( - "Failed to extract bytes for all functions in {:?}", - file_name - )) - } + info!( + "Finished extracting bytes for each function of {:?}", + file_name + ); + Ok(()) } fn get_checksums(&self) -> Result { @@ -2077,7 +2210,7 @@ impl FileToBeProcessed { fn setup_function_list(&self, r2p: &mut R2Pipe) -> Result> { info!( "Setting up function list for {:?} ...", - self.get_file_name() + self.get_file_name()? ); let json = r2p .cmdj("aflj") @@ -2123,7 +2256,7 @@ impl FileToBeProcessed { self.file_path ) })?; - debug!("Done getting function list for {:?}", self.get_file_name()); + debug!("Done getting function list for {:?}", self.get_file_name()?); Ok(functions_to_be_processed) } From 658976151426fe2ae6331e840061809e9ce330a3 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 20 Oct 2025 12:59:00 +0100 Subject: [PATCH 39/42] Process FuncInfo extraction first --- src/extract.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 199a4c0..ac265e1 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -643,8 +643,15 @@ impl ExtractionJob { for mode in modes { let job_type = Self::extraction_job_matcher(mode)?; - job_types.push((job_type, mode.clone())); - extraction_job_types.push(job_type); // Store just the job type + + // 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 { + job_types.push((job_type, mode.clone())); + extraction_job_types.push(job_type); + } } // Warn the user if mode-specific flags are turned on for no reason From 231ca6ce8fb9ef102f5a569296eb5392fe421135 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 20 Oct 2025 13:10:30 +0100 Subject: [PATCH 40/42] Reuse data from finfo to populate function list cache --- src/extract.rs | 34 +++++++++++++++++++++++++++------- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index ac265e1..3b14f65 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -643,7 +643,7 @@ impl ExtractionJob { 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())); @@ -1424,7 +1424,7 @@ impl FileToBeProcessed { pub fn get_function_list(&self, r2p: &mut R2Pipe) -> Result<&[FunctionToBeProcessed]> { self.function_list - .get_or_try_init(|| self.setup_function_list(r2p)) + .get_or_try_init(|| self.setup_function_list(r2p, None)) .map(|v| v.as_slice()) } @@ -1628,9 +1628,21 @@ impl FileToBeProcessed { .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; info!("Writing function info to {:?}", output_path); - std::fs::write(output_path, json_raw) + std::fs::write(output_path, &json_raw) .with_context(|| format!("Failed to write JSON to {:?}", output_path))?; info!("Function info written to {:?}", output_path); + + // 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); + } + Ok(()) } @@ -2214,14 +2226,22 @@ impl FileToBeProcessed { }) } - fn setup_function_list(&self, r2p: &mut R2Pipe) -> Result> { + fn setup_function_list( + &self, + r2p: &mut R2Pipe, + aflj_json: Option, + ) -> Result> { info!( "Setting up function list for {:?} ...", self.get_file_name()? ); - let json = r2p - .cmdj("aflj") - .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?; + + let json = match aflj_json { + Some(value) => value, + None => r2p + .cmdj("aflj") + .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?, + }; let array = match json { serde_json::Value::Array(arr) => arr, From 7fcd62721963b4d622eede3ab30b04f5b4dc5903 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 20 Oct 2025 13:38:48 +0100 Subject: [PATCH 41/42] Only cache FuncInfo list if other modes need it --- src/extract.rs | 56 +++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 44 insertions(+), 12 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 3b14f65..696ec69 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -104,6 +104,7 @@ pub struct FileToBeProcessed { pub func_filename_template: String, pub function_list: OnceCell>, pub min_basic_blocks: Option, + needs_func_list: bool, } #[derive(Debug, Clone)] @@ -309,6 +310,7 @@ impl bool, String, Option, + bool, )> for FileToBeProcessed { fn from( @@ -322,6 +324,7 @@ impl bool, String, Option, + bool, ), ) -> FileToBeProcessed { FileToBeProcessed { @@ -335,6 +338,7 @@ impl func_filename_template: orig.7, function_list: OnceCell::new(), min_basic_blocks: orig.8, + needs_func_list: orig.9, } } } @@ -681,8 +685,12 @@ impl ExtractionJob { timeout: *timeout, }; - let p_type = Self::get_path_type(input_path); + // 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 @@ -698,6 +706,7 @@ impl ExtractionJob { func_filename_template: func_filename_template.to_string(), function_list: OnceCell::new(), min_basic_blocks: *min_basic_blocks, + needs_func_list, }; Ok(ExtractionJob { @@ -733,6 +742,7 @@ impl ExtractionJob { func_filename_template: func_filename_template.to_string(), function_list: OnceCell::new(), min_basic_blocks: *min_basic_blocks, + needs_func_list, }) .collect(); @@ -792,6 +802,23 @@ impl ExtractionJob { } } + /// 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(); @@ -1422,6 +1449,9 @@ impl FileToBeProcessed { 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)) @@ -1632,15 +1662,18 @@ impl FileToBeProcessed { .with_context(|| format!("Failed to write JSON to {:?}", output_path))?; info!("Function info written to {:?}", output_path); - // 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); + // 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); + } } Ok(()) @@ -2040,7 +2073,6 @@ impl FileToBeProcessed { .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))?; @@ -2283,7 +2315,7 @@ impl FileToBeProcessed { self.file_path ) })?; - debug!("Done getting function list for {:?}", self.get_file_name()?); + debug!("Done setting up function list for {:?}", self.get_file_name()?); Ok(functions_to_be_processed) } From 39d547cf80d6c0b9207b8e4d253e5d93230cdcd2 Mon Sep 17 00:00:00 2001 From: valbucci Date: Mon, 20 Oct 2025 13:58:29 +0100 Subject: [PATCH 42/42] Populate function list cache from existing finfo JSON file --- src/extract.rs | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/src/extract.rs b/src/extract.rs index 696ec69..d46bf53 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -2268,11 +2268,34 @@ impl FileToBeProcessed { self.get_file_name()? ); + // Determine the source of function list data let json = match aflj_json { + // If JSON was provided, use it directly Some(value) => value, - None => r2p - .cmdj("aflj") - .with_context(|| format!("Failed executing aflj on {:?}", self.file_path))?, + + // 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))? + } + } }; let array = match json { @@ -2315,7 +2338,10 @@ impl FileToBeProcessed { self.file_path ) })?; - debug!("Done setting up function list for {:?}", self.get_file_name()?); + debug!( + "Done setting up function list for {:?}", + self.get_file_name()? + ); Ok(functions_to_be_processed) }