lib.rs

extern crate byteorder;

use std::io::{Read, Cursor, Seek, SeekFrom};
use std::fs::File;

use byteorder::{LittleEndian, ReadBytesExt};


pub type PeResult<T> = Result<T, PeError>;


#[derive(Debug)]
pub struct PeError {
    pub kind: PeErrorKind,
    pub details: String,
}

impl PeError {
    fn new<T: AsRef<str>>(kind: PeErrorKind, details: T) -> PeError {
        let description = match kind {
            PeErrorKind::WrongMagicNumbers => "Wrong magic numbers",
            _ => "BMP Error",
        };

        PeError { kind: kind, details: format!("{}: {}", description, details.as_ref()) }
    }
}

impl From<std::io::Error> for PeError {
    fn from(err: std::io::Error) -> PeError {
        PeError::new(PeErrorKind::PeIoError(err), "Io Error")
    }
}

impl From<byteorder::Error> for PeError {
    fn from(err: byteorder::Error) -> PeError {
        PeError::new(PeErrorKind::PeByteorderError(err), "Io Error")
    }
}


/// The different kinds of possible PE errors.
#[derive(Debug)]
pub enum PeErrorKind {
    WrongMagicNumbers,
    EndOfFileReached,
    PeIoError(std::io::Error),
    PeByteorderError(byteorder::Error),
}

#[derive(Debug)]
pub struct PeFile {
    pub dos_header: DosHeader,
}

#[derive(Debug, Default)]
pub struct DosHeader {
    pub magic: [u8; 2],
    pub cblp: u16,
    pub cp: u16,
    pub crlc: u16,
    pub cparhdr: u16,
    pub minalloc: u16,
    pub maxalloc: u16,
    pub ss: u16,
    pub sp: u16,
    pub csum: u16,
    pub ip: u16,
    pub cs: u16,
    pub lfarc: u16,
    pub ovno: u16,
    pub reserved1: [u8; 8], 
    pub oemid: u16,
    pub oeminfo: u16,
    pub reserved2: [u8; 20], 
    pub lfanew: i32, // file address of PE header
}

#[derive(Debug)]
pub struct NtHeader {
    pub magic: [u8; 2],
}

// TODO : is it really needed to read the whole file first?
pub fn open(path: &str) -> PeResult<PeFile> {
    let mut bytes = Vec::new();
    let mut f = try!(std::fs::File::open(path));
    try!(f.read_to_end(&mut bytes));
    let mut pe_data = Cursor::new(bytes);

    parse_pe(&mut pe_data)
}


fn parse_pe(pe_data: &mut Cursor<Vec<u8>>) -> PeResult<PeFile> {

    let dos_header = try!(parse_dos_header(pe_data));

    // advance the cursor to the PE header
    try!(pe_data.seek(SeekFrom::Start(dos_header.lfanew as u64)));

    let nt_header = try!(parse_nt_header(pe_data));
    let pefile = PeFile { dos_header: dos_header };
    Ok(pefile)

}

fn parse_dos_header(pe_data: &mut Cursor<Vec<u8>>) -> PeResult<DosHeader> {

    let mut dos_header = DosHeader {.. Default::default()};
    
    try!(pe_data.read(&mut dos_header.magic));
    dos_header.cblp = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.cp = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.crlc = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.cparhdr = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.minalloc = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.maxalloc = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.ss = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.sp = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.csum = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.ip = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.cs = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.lfarc = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.ovno = try!(pe_data.read_u16::<LittleEndian>());
    try!(pe_data.read(&mut dos_header.reserved1));
    dos_header.oemid = try!(pe_data.read_u16::<LittleEndian>());
    dos_header.oeminfo = try!(pe_data.read_u16::<LittleEndian>());
    try!(pe_data.read(&mut dos_header.reserved2));
    dos_header.lfanew = try!(pe_data.read_i32::<LittleEndian>());
    
    if dos_header.magic != b"MZ"[..] {
        return Err(PeError::new(PeErrorKind::WrongMagicNumbers, 
                format!("Expected 'MZ', but was {:?}", dos_header.magic)));
    }

    Ok(dos_header)

}


fn parse_nt_header(pe_data: &mut Cursor<Vec<u8>>) -> PeResult<NtHeader> {
    
    let mut magic = [0, 0];
    try!(pe_data.read(&mut magic));
    
    if magic != b"PE"[..] {
        return Err(PeError::new(PeErrorKind::WrongMagicNumbers, 
                format!("Expected 'PE', but was {:?}", magic)));
    }

    let nt_header = NtHeader {
        magic : magic,
    };

    Ok(nt_header)
}