add custom error

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<T> = std::result::Result<T, Error>;

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);
         }
     }
 }

src/obfus.rs

@@ -32,11 +32,11 @@ pub struct Obfuscator {
 }
 
 impl Obfuscator {
-    pub fn open(input_path: &str, output_path: &str) -> std::io::Result<Obfuscator> {
+    pub fn open(input_path: &str, output_path: &str) -> crate::error::Result<Obfuscator> {
         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();
         }
     }
 }