use crate::io::BinaryStream;
use crate::search::{SectionHelper, SearchSection};
use crate::error::{Error, Result};
pub const IMAGE_DOS_SIGNATURE: u16 = 0x5A4D;
pub const IMAGE_NT_SIGNATURE: u32 = 0x00004550;
pub const IMAGE_SCN_CNT_CODE: u32 = 0x00000020;
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x00000040;
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x20000000;
pub const IMAGE_DIRECTORY_ENTRY_EXPORT: usize = 0;
#[derive(Debug, Clone, Default)]
pub struct PeSectionHeader {
pub name: String,
pub virtual_size: u32,
pub virtual_address: u32,
pub size_of_raw_data: u32,
pub pointer_to_raw_data: u32,
pub pointer_to_relocations: u32,
pub pointer_to_linenumbers: u32,
pub number_of_relocations: u16,
pub number_of_linenumbers: u16,
pub characteristics: u32,
}
#[derive(Debug, Clone, Default)]
pub struct DataDirectory {
pub virtual_address: u32,
pub size: u32,
}
pub struct Pe {
pub stream: BinaryStream,
pub is_32bit: bool,
image_base: u64,
pub sections: Vec<PeSectionHeader>,
data_directories: Vec<DataDirectory>,
}
impl Pe {
pub fn new(data: Vec<u8>) -> Result<Self> {
let mut pe = Self {
stream: BinaryStream::new(data),
is_32bit: true,
image_base: 0,
sections: Vec::new(),
data_directories: Vec::new(),
};
pe.load()?;
Ok(pe)
}
fn load(&mut self) -> Result<()> {
self.stream.set_position(0);
let e_magic = self.stream.read_u16()?;
if e_magic != IMAGE_DOS_SIGNATURE {
return Err(Error::InvalidFormat("Invalid DOS signature".into()));
}
self.stream.set_position(0x3C);
let e_lfanew = self.stream.read_u32()?;
self.stream.set_position(e_lfanew as u64);
let nt_sig = self.stream.read_u32()?;
if nt_sig != IMAGE_NT_SIGNATURE {
return Err(Error::InvalidFormat("Invalid NT signature".into()));
}
let _machine = self.stream.read_u16()?;
let number_of_sections = self.stream.read_u16()?;
let _time_date_stamp = self.stream.read_u32()?;
let _pointer_to_symbol_table = self.stream.read_u32()?;
let _number_of_symbols = self.stream.read_u32()?;
let size_of_optional_header = self.stream.read_u16()?;
let _characteristics = self.stream.read_u16()?;
let optional_header_start = self.stream.position();
let optional_magic = self.stream.read_u16()?;
if optional_magic == 0x20B {
self.is_32bit = false;
self.stream.is_32bit = false;
self.stream.set_position(optional_header_start + 24);
self.image_base = self.stream.read_u64()?;
self.stream.set_position(optional_header_start + 108);
let num_rva_and_sizes = self.stream.read_u32()?;
self.data_directories.clear();
for _ in 0..std::cmp::min(num_rva_and_sizes, 16) {
self.data_directories.push(DataDirectory {
virtual_address: self.stream.read_u32()?,
size: self.stream.read_u32()?,
});
}
} else {
self.is_32bit = true;
self.stream.is_32bit = true;
self.stream.set_position(optional_header_start + 28);
self.image_base = self.stream.read_u32()? as u64;
self.stream.set_position(optional_header_start + 92);
let num_rva_and_sizes = self.stream.read_u32()?;
self.data_directories.clear();
for _ in 0..std::cmp::min(num_rva_and_sizes, 16) {
self.data_directories.push(DataDirectory {
virtual_address: self.stream.read_u32()?,
size: self.stream.read_u32()?,
});
}
}
self.stream.image_base = self.image_base;
self.stream.set_position(optional_header_start + size_of_optional_header as u64);
self.sections.clear();
for _ in 0..number_of_sections {
let name_bytes = self.stream.read_bytes(8)?;
let name = String::from_utf8_lossy(
&name_bytes.iter().copied().take_while(|&b| b != 0).collect::<Vec<u8>>()
).to_string();
self.sections.push(PeSectionHeader {
name,
virtual_size: self.stream.read_u32()?,
virtual_address: self.stream.read_u32()?,
size_of_raw_data: self.stream.read_u32()?,
pointer_to_raw_data: self.stream.read_u32()?,
pointer_to_relocations: self.stream.read_u32()?,
pointer_to_linenumbers: self.stream.read_u32()?,
number_of_relocations: self.stream.read_u16()?,
number_of_linenumbers: self.stream.read_u16()?,
characteristics: self.stream.read_u32()?,
});
}
Ok(())
}
pub fn map_vatr(&self, mut addr: u64) -> Result<u64> {
if addr >= self.image_base {
addr -= self.image_base;
}
for section in &self.sections {
let start = section.virtual_address as u64;
let end = start + section.virtual_size as u64;
if addr >= start && addr < end {
return Ok(addr - section.virtual_address as u64 + section.pointer_to_raw_data as u64);
}
}
Err(Error::AddressNotMapped(addr))
}
pub fn map_rtva(&self, addr: u64) -> u64 {
for section in &self.sections {
let start = section.pointer_to_raw_data as u64;
let end = start + section.size_of_raw_data as u64;
if addr >= start && addr < end {
return addr - section.pointer_to_raw_data as u64 + section.virtual_address as u64 + self.image_base;
}
}
0
}
pub fn list_exported_symbols(&mut self) -> Result<Vec<(String, u64)>> {
let mut exports = Vec::new();
if self.data_directories.len() <= IMAGE_DIRECTORY_ENTRY_EXPORT {
return Ok(exports);
}
let export_dir = &self.data_directories[IMAGE_DIRECTORY_ENTRY_EXPORT];
if export_dir.virtual_address == 0 {
return Ok(exports);
}
let export_offset = self.map_vatr(export_dir.virtual_address as u64)?;
self.stream.set_position(export_offset);
let _characteristics = self.stream.read_u32()?;
let _time_date_stamp = self.stream.read_u32()?;
let _major_version = self.stream.read_u16()?;
let _minor_version = self.stream.read_u16()?;
let _name_rva = self.stream.read_u32()?;
let _base = self.stream.read_u32()?;
let _number_of_functions = self.stream.read_u32()?;
let number_of_names = self.stream.read_u32()?;
let address_of_functions = self.stream.read_u32()?;
let address_of_names = self.stream.read_u32()?;
let address_of_name_ordinals = self.stream.read_u32()?;
let names_offset = self.map_vatr(address_of_names as u64)?;
let ordinals_offset = self.map_vatr(address_of_name_ordinals as u64)?;
let functions_offset = self.map_vatr(address_of_functions as u64)?;
for i in 0..number_of_names {
self.stream.set_position(names_offset + i as u64 * 4);
let name_rva = self.stream.read_u32()?;
let name_offset = match self.map_vatr(name_rva as u64) {
Ok(o) => o,
Err(_) => continue,
};
let name = match self.stream.read_string_to_null_at(name_offset) {
Ok(n) => n,
Err(_) => continue,
};
self.stream.set_position(ordinals_offset + i as u64 * 2);
let ordinal = self.stream.read_u16()?;
self.stream.set_position(functions_offset + ordinal as u64 * 4);
let func_rva = self.stream.read_u32()?;
exports.push((name, func_rva as u64));
}
Ok(exports)
}
pub fn symbol_search(&mut self) -> Result<Option<(u64, u64)>> {
if self.data_directories.len() <= IMAGE_DIRECTORY_ENTRY_EXPORT {
return Ok(None);
}
let export_dir = &self.data_directories[IMAGE_DIRECTORY_ENTRY_EXPORT];
if export_dir.virtual_address == 0 {
return Ok(None);
}
let export_offset = self.map_vatr(export_dir.virtual_address as u64)?;
self.stream.set_position(export_offset);
let _characteristics = self.stream.read_u32()?;
let _time_date_stamp = self.stream.read_u32()?;
let _major_version = self.stream.read_u16()?;
let _minor_version = self.stream.read_u16()?;
let _name_rva = self.stream.read_u32()?;
let _base = self.stream.read_u32()?;
let _number_of_functions = self.stream.read_u32()?;
let number_of_names = self.stream.read_u32()?;
let address_of_functions = self.stream.read_u32()?;
let address_of_names = self.stream.read_u32()?;
let address_of_name_ordinals = self.stream.read_u32()?;
let names_offset = self.map_vatr(address_of_names as u64)?;
let ordinals_offset = self.map_vatr(address_of_name_ordinals as u64)?;
let functions_offset = self.map_vatr(address_of_functions as u64)?;
let mut code_reg = 0u64;
let mut metadata_reg = 0u64;
for i in 0..number_of_names {
self.stream.set_position(names_offset + i as u64 * 4);
let name_rva = self.stream.read_u32()?;
let name_offset = self.map_vatr(name_rva as u64)?;
let name = self.stream.read_string_to_null_at(name_offset)?;
self.stream.set_position(ordinals_offset + i as u64 * 2);
let ordinal = self.stream.read_u16()?;
self.stream.set_position(functions_offset + ordinal as u64 * 4);
let func_rva = self.stream.read_u32()?;
if name == "g_CodeRegistration" {
code_reg = func_rva as u64 + self.image_base;
} else if name == "g_MetadataRegistration" {
metadata_reg = func_rva as u64 + self.image_base;
}
}
if code_reg > 0 && metadata_reg > 0 {
Ok(Some((code_reg, metadata_reg)))
} else {
Ok(None)
}
}
pub fn get_section_helper(&self, method_count: usize, type_definitions_count: usize, metadata_usages_count: usize, image_count: usize, version: f64) -> SectionHelper<'_> {
let mut data_list = Vec::new();
let mut exec_list = Vec::new();
let mut all_sections = Vec::new();
for section in &self.sections {
if section.virtual_size == 0 {
continue;
}
let search_section = SearchSection::new(
section.pointer_to_raw_data as u64,
(section.pointer_to_raw_data + section.size_of_raw_data) as u64,
section.virtual_address as u64 + self.image_base,
(section.virtual_address + section.virtual_size) as u64 + self.image_base,
);
all_sections.push(search_section.clone());
let chars = section.characteristics;
let is_code = (chars & IMAGE_SCN_CNT_CODE) != 0 || (chars & IMAGE_SCN_MEM_EXECUTE) != 0;
let is_data = (chars & IMAGE_SCN_CNT_INITIALIZED_DATA) != 0;
if is_code && !is_data {
exec_list.push(search_section);
} else if is_data && !is_code {
data_list.push(search_section);
} else {
}
}
let bss = data_list.clone();
SectionHelper::new(
self.stream.data(),
self.is_32bit,
version,
false,
all_sections,
data_list,
exec_list,
bss,
method_count,
type_definitions_count,
metadata_usages_count,
image_count,
)
}
pub fn check_dump(&self) -> bool {
if self.is_32bit {
self.image_base != 0x10000000
} else {
self.image_base != 0x180000000
}
}
pub fn load_from_memory(&mut self, addr: u64) {
self.image_base = addr;
self.stream.image_base = addr;
for section in &mut self.sections {
section.pointer_to_raw_data = section.virtual_address;
section.size_of_raw_data = section.virtual_size;
}
}
pub fn get_rva(&self, pointer: u64) -> u64 {
pointer - self.image_base
}
pub fn image_base(&self) -> u64 {
self.image_base
}
}