1
Fork 0
oelf/crates/goblin-pyo3/src/lib.rs
2023-12-08 17:44:53 +01:00

154 lines
4.2 KiB
Rust

use std::{fs::File, io::Read};
use goblin::mach::Mach;
use pyo3::{exceptions::PyTypeError, prelude::*};
mod exports;
mod header;
mod imports;
mod sections;
mod symbols;
use exports::Export;
use header::Header;
use imports::Import;
use sections::{Sections, Section};
use symbols::Symbols;
#[pyclass]
struct Object {
len: usize,
ptr: *mut u8,
inner: Option<goblin::Object<'static>>,
}
// SAFETY: We only use `ptr` in `drop` to reconstruct a `Vec`
unsafe impl Send for Object {}
#[pymethods]
impl Object {
#[new]
fn new(path: String) -> Self {
let mut file = File::open(path).unwrap();
let size = file.metadata().map(|m| m.len() as usize).ok();
let mut vec = Vec::with_capacity(size.unwrap_or(0));
file.read_to_end(&mut vec).unwrap();
vec.shrink_to_fit();
let len = vec.len();
let cap = vec.capacity();
let ptr = vec.as_mut_ptr();
assert!(len == cap);
let obj = vec.leak();
let object = goblin::Object::parse(obj).unwrap();
Self {
len,
ptr,
inner: Some(object),
}
}
#[getter]
fn header(&self) -> Header {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => Header::from(macho.header),
_ => unimplemented!(),
}
}
#[getter]
fn name(&self) -> Option<&str> {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => macho.name.clone(),
_ => unimplemented!(),
}
}
fn symbols(&self) -> Symbols {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => Symbols::from(macho.symbols()),
_ => unimplemented!(),
}
}
fn sections(&self) -> Sections {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => {
let mut sections = vec![];
for sect_iter in macho.segments.sections() {
sections.extend(sect_iter.map(|section| {
let (sect, _data) = section.unwrap();
Section::from(sect)
}));
}
Sections { sections }
},
_ => unimplemented!(),
}
}
#[getter]
fn libs(&self) -> Vec<&str> {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => macho.libs.clone(),
_ => unimplemented!(),
}
}
#[getter]
fn rpaths(&self) -> Vec<&str> {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => macho.rpaths.clone(),
_ => unimplemented!(),
}
}
fn exports(&self) -> Result<Vec<Export>, PyErr> {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => {
let exports = macho
.exports()
.map_err(|_| PyErr::new::<PyTypeError, _>("failed"))?;
Ok(exports.into_iter().map(|exp| exp.into()).collect())
}
_ => unimplemented!(),
}
}
fn imports(&self) -> Result<Vec<Import>, PyErr> {
match self.inner.as_ref().unwrap() {
goblin::Object::Mach(Mach::Binary(macho)) => {
let imports = macho
.imports()
.map_err(|_| PyErr::new::<PyTypeError, _>("failed"))?;
Ok(imports.into_iter().map(|exp| exp.into()).collect())
}
_ => unimplemented!(),
}
}
}
impl Drop for Object {
fn drop(&mut self) {
let obj = self.inner.take();
drop(obj);
// SAFETY:
// We took `ptr` and `len` from the vec earlier (and ensured `len` == `cap`),
// then leaked it to get a static reference to it
// which was only held within `self.inner`, which has been dropped above.
unsafe {
let vec = Vec::from_raw_parts(self.ptr, self.len, self.len);
drop(vec);
}
}
}
#[pymodule]
#[pyo3(name = "goblin")]
fn py_goblin(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
m.add_class::<Object>()?;
Ok(())
}