use std::fs;
use std::path::Path;
use std::error::Error;
use std::fmt;
use rayon::prelude::*;
use rustc_hash::FxHashSet;
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub struct Segment {
pub addr: u64,
pub bytes: Vec<u8>,
}
impl Segment {
pub fn new(addr: u64, bytes: Vec<u8>) -> Segment {
Segment { addr, bytes }
}
pub fn contains(&self, addr: u64) -> bool {
(self.addr <= addr) && (addr < (self.addr + self.bytes.len() as u64))
}
pub fn get_matching_offsets(&self, vals: &[u8]) -> Vec<usize> {
self.bytes
.par_iter()
.enumerate()
.filter(|&(_, b)| vals.contains(b))
.map(|(i, _)| i)
.collect()
}
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Format {
Unknown,
ELF,
PE,
Raw,
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum Arch {
Unknown,
X8086,
X86,
X64,
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Binary {
pub name: String,
pub format: Format,
pub arch: Arch,
pub entry: u64,
pub segments: FxHashSet<Segment>,
}
impl Binary {
pub fn from_bytes(name: &str, bytes: &[u8]) -> Result<Binary, Box<dyn Error>> {
Binary::priv_from_buf(name, bytes)
}
pub fn from_path_str(path: &str) -> Result<Binary, Box<dyn Error>> {
let name = Path::new(path).file_name().ok_or("No filename.")?;
let name_str = name.to_str().ok_or("Failed filename decode.")?;
let bytes = fs::read(path)?;
Binary::priv_from_buf(name_str, &bytes)
}
fn priv_new() -> Binary {
Binary {
name: String::from("None"),
format: Format::Unknown,
arch: Arch::Unknown,
entry: 0,
segments: FxHashSet::default(),
}
}
fn priv_from_buf(name: &str, bytes: &[u8]) -> Result<Binary, Box<dyn Error>> {
match goblin::Object::parse(&bytes)? {
goblin::Object::Elf(elf) => {
Binary::from_elf(name, &bytes, &elf)
}
goblin::Object::PE(pe) => {
Binary::from_pe(name, &bytes, &pe)
}
goblin::Object::Unknown(_) => {
Ok(Binary::from_raw(name, bytes))
}
_ => {
Err("Unsupported file format!".into())
}
}
}
fn from_elf(name: &str, bytes: &[u8], elf: &goblin::elf::Elf) -> Result<Binary, Box<dyn Error>> {
let mut bin = Binary::priv_new();
bin.name = name.to_string();
bin.entry = elf.entry;
bin.format = Format::ELF;
bin.arch = match elf.header.e_machine {
goblin::elf::header::EM_X86_64 => {
Arch::X64
}, goblin::elf::header::EM_386 => {
Arch::X86
},
_ => {
return Err("Unsupported architecture!".into());
}
};
for prog_hdr in elf.program_headers.iter().filter(|&p| (p.p_flags & goblin::elf::program_header::PF_X) != 0) {
let start_offset = prog_hdr.p_offset as usize;
let end_offset = start_offset + prog_hdr.p_filesz as usize;
bin.segments.insert(
Segment::new(
prog_hdr.p_vaddr,
bytes[start_offset..end_offset].to_vec(),
)
);
}
bin.remove_sub_segs();
Ok(bin)
}
fn from_pe(name: &str, bytes: &[u8], pe: &goblin::pe::PE) -> Result<Binary, Box<dyn Error>> {
let mut bin = Binary::priv_new();
bin.name = name.to_string();
bin.entry = pe.entry as u64;
bin.format = Format::PE;
bin.arch = match pe.header.coff_header.machine {
goblin::pe::header::COFF_MACHINE_X86_64 => {
Arch::X64
}, goblin::pe::header::COFF_MACHINE_X86 => {
Arch::X86
},
_ => {
return Err("Unsupported architecture!".into());
}
};
for sec_tab in pe.sections.iter()
.filter(|&p| (p.characteristics & goblin::pe::section_table::IMAGE_SCN_MEM_EXECUTE) != 0) {
let start_offset = sec_tab.pointer_to_raw_data as usize;
let end_offset = start_offset + sec_tab.size_of_raw_data as usize;
bin.segments.insert(
Segment::new(
sec_tab.virtual_address as u64 + pe.image_base as u64,
bytes[start_offset..end_offset].to_vec(),
)
);
}
bin.remove_sub_segs();
Ok(bin)
}
fn from_raw(name: &str, bytes: &[u8]) -> Binary {
let mut bin = Binary::priv_new();
bin.name = name.to_string();
bin.entry = 0;
bin.format = Format::Raw;
bin.segments.insert(
Segment::new(
0,
bytes[..].to_vec()
)
);
bin
}
fn remove_sub_segs(&mut self) {
let mut sub_segs = Vec::new();
for seg in &self.segments {
let mut local_sub_segs = self.segments.iter()
.cloned()
.filter(|s| (s.addr == seg.addr) && (s.bytes.len() < seg.bytes.len()))
.collect();
sub_segs.append(&mut local_sub_segs);
}
for s in sub_segs {
self.segments.remove(&s);
}
}
}
impl fmt::Display for Binary {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "\'{}\': {:?}-{:?}, entry 0x{:016x}, {}/{} executable bytes/segments",
self.name,
self.format,
self.arch,
self.entry,
self.segments.iter().fold(0, | bytes, seg | bytes + seg.bytes.len()),
self.segments.len(),
)
}
}