Check for macho in constructor

crates/goblin-pyo3/src/lib.rs

@@ -24,16 +24,22 @@ struct Object {
     ptr: *mut u8,
     #[pyo3(get)]
     path: String,
-    inner: Option<goblin::Object<'static>>,
+    inner: Option<goblin::mach::MachO<'static>>,
 }
 
 // SAFETY: We only use `ptr` in `drop` to reconstruct a `Vec`
 unsafe impl Send for Object {}
 
+impl Object {
+    fn macho(&self) -> &goblin::mach::MachO {
+        self.inner.as_ref().unwrap()
+    }
+}
+
 #[pymethods]
 impl Object {
     #[new]
-    fn new(path: String) -> Self {
+    fn new(path: String) -> PyResult<Self> {
         let path = fs::canonicalize(path).unwrap();
         let mut file = File::open(&path).unwrap();
         let size = file.metadata().map(|m| m.len() as usize).ok();
@@ -49,91 +55,69 @@ impl Object {
         let obj = vec.leak();
         let object = goblin::Object::parse(obj).unwrap();
 
-        Self {
+        let macho = match object {
+            goblin::Object::Mach(Mach::Binary(macho)) => macho,
+            _ => return Err(PyErr::new::<PyTypeError, _>("not a macho file")),
+        };
+
+        Ok(Self {
             len,
             ptr,
             path: path.display().to_string(),
-            inner: Some(object),
+            inner: Some(macho),
-        }
+        })
     }
 
     #[getter]
     fn header(&self) -> Header {
-        match self.inner.as_ref().unwrap() {
+        self.macho().header.into()
-            goblin::Object::Mach(Mach::Binary(macho)) => Header::from(macho.header),
-            _ => unimplemented!(),
-        }
     }
 
     #[getter]
     fn name(&self) -> Option<&str> {
-        match self.inner.as_ref().unwrap() {
+        self.macho().name
-            goblin::Object::Mach(Mach::Binary(macho)) => macho.name,
-            _ => unimplemented!(),
-        }
     }
 
     fn symbols(&self) -> Symbols {
-        match self.inner.as_ref().unwrap() {
+        Symbols::from(self.macho().symbols())
-            goblin::Object::Mach(Mach::Binary(macho)) => Symbols::from(macho.symbols()),
-            _ => unimplemented!(),
-        }
     }
 
     fn sections(&self) -> Sections {
-        match self.inner.as_ref().unwrap() {
+        let macho = self.macho();
-            goblin::Object::Mach(Mach::Binary(macho)) => {
+        let mut sections = vec![];
-                let mut sections = vec![];
+        for sect_iter in macho.segments.sections() {
-                for sect_iter in macho.segments.sections() {
+            sections.extend(sect_iter.map(|section| {
-                    sections.extend(sect_iter.map(|section| {
+                let (sect, _data) = section.unwrap();
-                        let (sect, _data) = section.unwrap();
+                Section::from(sect)
-                        Section::from(sect)
+            }));
-                    }));
-                }
-                Sections { sections }
-            }
-            _ => unimplemented!(),
         }
+        Sections { sections }
     }
 
     #[getter]
     fn libs(&self) -> Vec<&str> {
-        match self.inner.as_ref().unwrap() {
+        self.macho().libs.clone()
-            goblin::Object::Mach(Mach::Binary(macho)) => macho.libs.clone(),
-            _ => unimplemented!(),
-        }
     }
 
     #[getter]
     fn rpaths(&self) -> Vec<&str> {
-        match self.inner.as_ref().unwrap() {
+        self.macho().rpaths.clone()
-            goblin::Object::Mach(Mach::Binary(macho)) => macho.rpaths.clone(),
-            _ => unimplemented!(),
-        }
     }
 
     fn exports(&self) -> Result<Vec<Export>, PyErr> {
-        match self.inner.as_ref().unwrap() {
+        let macho = self.macho();
-            goblin::Object::Mach(Mach::Binary(macho)) => {
+        let exports = macho
-                let exports = macho
+            .exports()
-                    .exports()
+            .map_err(|_| PyErr::new::<PyTypeError, _>("failed"))?;
-                    .map_err(|_| PyErr::new::<PyTypeError, _>("failed"))?;
+        Ok(exports.into_iter().map(|exp| exp.into()).collect())
-                Ok(exports.into_iter().map(|exp| exp.into()).collect())
-            }
-            _ => unimplemented!(),
-        }
     }
 
     fn imports(&self) -> Result<Vec<Import>, PyErr> {
-        match self.inner.as_ref().unwrap() {
+        let macho = self.macho();
-            goblin::Object::Mach(Mach::Binary(macho)) => {
+        let imports = macho
-                let imports = macho
+            .imports()
-                    .imports()
+            .map_err(|_| PyErr::new::<PyTypeError, _>("failed"))?;
-                    .map_err(|_| PyErr::new::<PyTypeError, _>("failed"))?;
+        Ok(imports.into_iter().map(|exp| exp.into()).collect())
-                Ok(imports.into_iter().map(|exp| exp.into()).collect())
-            }
-            _ => unimplemented!(),
-        }
     }
 }
 

crates/goblin-pyo3/test.py

@@ -1,6 +1,7 @@
-import goblin
+import oelf
 
-g = goblin.Object("mylib.dylib")
+g = oelf.Object("test.py")
+g = oelf.Object("mylib.dylib")
 print(g.header)
 print(g.name)