From 621e90f1a02952188499931810958b25f6529a0a Mon Sep 17 00:00:00 2001 From: lapla-cogito Date: Tue, 20 Aug 2024 08:32:59 +0900 Subject: [PATCH] add custom error --- src/error.rs | 25 +++++++++++++++++++++++++ src/main.rs | 21 ++++++++++++++------- src/obfus.rs | 53 +++++++++++++++++++++++++++++----------------------- 3 files changed, 69 insertions(+), 30 deletions(-) create mode 100644 src/error.rs diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..c46c4d3 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,25 @@ +#[derive(thiserror::Error, Debug)] +pub enum Error { + #[error("invalid option: {0}")] + InvalidOption(&'static str), + + #[error("failed to open file: {0}")] + OpenFile(std::io::Error), + + #[error("failed to create file: {0}")] + CreateFile(std::io::Error), + + #[error("invalid ELF magic")] + InvalidMagic, + + #[error("failed to mmap: {0}")] + Mmap(std::io::Error), + + #[error("failed in I/O operation: {0}")] + Io(std::io::Error), + + #[error("")] + ExitWithCode(std::process::ExitCode), +} + +pub type Result = std::result::Result; diff --git a/src/main.rs b/src/main.rs index cfc0b1d..47e57e0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +mod error; mod obfus; mod util; @@ -47,13 +48,15 @@ struct Args { recursive: String, } -fn main() { +fn main() -> crate::error::Result<()> { use clap::Parser as _; let args = Args::parse(); if args.recursive.is_empty() { if args.input.is_empty() { - panic!("input file name is required"); + return Err(crate::error::Error::InvalidOption( + "input file name is required", + )); } let output_path = if !args.output.is_empty() { @@ -62,13 +65,15 @@ fn main() { "obfuscated" }; - exec_obfus(&args.input, output_path, &args).unwrap_or(()); + exec_obfus(&args.input, output_path, &args).unwrap_or(()) } else { if !args.input.is_empty() { - panic!("both input file name and recursive option are not allowed"); + return Err(crate::error::Error::InvalidOption( + "both input file name and recursive option are not allowed", + )); } if !args.output.is_empty() { - println!("output file name will be ignored"); + eprintln!("output file name will be ignored"); } let entries = util::RecursiveDir::new(&args.recursive) @@ -86,9 +91,11 @@ fn main() { exec_obfus(entry.to_str().unwrap(), &output_path, &args).unwrap_or(()); } } + + Ok(()) } -fn exec_obfus(input_path: &str, output_path: &str, args: &Args) -> std::io::Result<()> { +fn exec_obfus(input_path: &str, output_path: &str, args: &Args) -> crate::error::Result<()> { let loader = obfus::Obfuscator::open(input_path, output_path); let mut obfuscator = loader.unwrap(); @@ -119,7 +126,7 @@ fn exec_obfus(input_path: &str, output_path: &str, args: &Args) -> std::io::Resu Ok(()) } false => { - panic!("not a valid ELF file: {}", args.input); + return Err(crate::error::Error::InvalidMagic); } } } diff --git a/src/obfus.rs b/src/obfus.rs index e54c772..b820782 100644 --- a/src/obfus.rs +++ b/src/obfus.rs @@ -32,11 +32,11 @@ pub struct Obfuscator { } impl Obfuscator { - pub fn open(input_path: &str, output_path: &str) -> std::io::Result { + pub fn open(input_path: &str, output_path: &str) -> crate::error::Result { let file = match std::fs::OpenOptions::new().read(true).open(input_path) { Ok(file) => file, Err(e) => { - panic!("failed to open file: {}", e); + return Err(crate::error::Error::OpenFile(e)); } }; @@ -49,18 +49,23 @@ impl Obfuscator { { Ok(file) => file, Err(e) => { - panic!("failed to create file: {}", e); + return Err(crate::error::Error::CreateFile(e)); } }; let mut input_contents = Vec::new(); - file.try_clone()? + file.try_clone() + .map_err(crate::error::Error::Io)? .take(usize::MAX as u64) - .read_to_end(&mut input_contents)?; - output_file.write_all(&input_contents)?; + .read_to_end(&mut input_contents) + .map_err(crate::error::Error::Io)?; + output_file + .write_all(&input_contents) + .map_err(crate::error::Error::Io)?; - let input = unsafe { memmap2::Mmap::map(&file)? }; - let output = unsafe { memmap2::MmapMut::map_mut(&output_file)? }; + let input = unsafe { memmap2::Mmap::map(&file).map_err(crate::error::Error::Mmap)? }; + let output = + unsafe { memmap2::MmapMut::map_mut(&output_file).map_err(crate::error::Error::Mmap)? }; let elf_hdr: ElfHeader = unsafe { std::ptr::read(input.as_ptr() as *const ElfHeader) }; @@ -130,14 +135,14 @@ impl Obfuscator { } fn is_stripped(&self) -> bool { - self.get_section(".symtab").0 == 0 + self.get_section(".symtab").unwrap().0 == 0 } // (section_addr, section_size, entry_size, vaddr) - fn get_section(&self, section: &str) -> (usize, usize, usize, usize) { + fn get_section(&self, section: &str) -> crate::error::Result<(usize, usize, usize, usize)> { let searched_idx = self.sec_hdr.find(section).unwrap_or(usize::MAX); if searched_idx == usize::MAX { - panic!("section not found"); + return Err(crate::error::Error::InvalidOption("section not found")); } for i in 0..self.sec_hdr_num { @@ -147,25 +152,25 @@ impl Obfuscator { let string_offset = u32::from_le_bytes(sec_hdr[0..4].try_into().unwrap()); if string_offset == searched_idx as u32 { if self.is_64bit() { - return ( + return Ok(( u64::from_le_bytes(sec_hdr[24..32].try_into().unwrap()) as usize, u64::from_le_bytes(sec_hdr[32..40].try_into().unwrap()) as usize, u64::from_le_bytes(sec_hdr[56..64].try_into().unwrap()) as usize, u64::from_le_bytes(sec_hdr[16..24].try_into().unwrap()) as usize, - ); + )); } else { - return ( + return Ok(( u32::from_le_bytes(sec_hdr[16..20].try_into().unwrap()) as usize, u32::from_le_bytes(sec_hdr[20..24].try_into().unwrap()) as usize, u32::from_le_bytes(sec_hdr[36..40].try_into().unwrap()) as usize, u32::from_le_bytes(sec_hdr[12..16].try_into().unwrap()) as usize, - ); + )); } } } // section not found - (usize::MAX, usize::MAX, usize::MAX, usize::MAX) + Err(crate::error::Error::InvalidOption("section not found")) } pub fn change_class(&mut self) { @@ -185,15 +190,17 @@ impl Obfuscator { } } - pub fn nullify_section(&mut self, section: &str) { - let (section_addr, section_size, _, _) = self.get_section(section); + pub fn nullify_section(&mut self, section: &str) -> crate::error::Result<()> { + let (section_addr, section_size, _, _) = self.get_section(section).unwrap(); if section_addr == usize::MAX { - panic!("section not found"); + return Err(crate::error::Error::InvalidOption("section not found")); } for i in section_addr..section_addr + section_size { self.output[i] = 0; } + + Ok(()) } pub fn got_overwrite(&self, function: &str, new_func_addr: &str) { @@ -203,16 +210,16 @@ impl Obfuscator { println!("cannot overwrite GOT with stripped binary") } - // let mut addr; - if self.is_64bit() { - let (section_addr, section_size, entry_size, vaddr) = self.get_section(".rela.plt"); + let (section_addr, section_size, entry_size, vaddr) = + self.get_section(".rela.plt").unwrap(); for i in 0..section_size / entry_size { let entry = &self.input[section_addr..section_addr + section_size] [i * entry_size..(i + 1) * entry_size]; } } else { - let (section_addr, section_size, entry_size, vaddr) = self.get_section(".rel.plt"); + let (section_addr, section_size, entry_size, vaddr) = + self.get_section(".rel.plt").unwrap(); } } }