pe_parser/
pe.rs

1use crate::{coff::CoffFileHeader, optional::{OptionalHeader32, OptionalHeader64, Magic, Optional}, section::{SectionHeader, parse_section_table}, Error};
2use bytemuck::checked::try_from_bytes;
3use num_traits::FromPrimitive;
4use core::fmt;
5use crate::prelude::*;
6
7const IMAGE_DOS_PE_SIGNATURE_OFFSET: usize = 0x3c;
8
9/// Representation of the sections of a Portable Executable
10pub struct PortableExecutable {
11    /// COFF File Header (Object and Image)
12    pub coff: CoffFileHeader,
13    /// PE32 Optional Header (Image Only)
14    pub optional_header_32: Option<OptionalHeader32>,
15    /// PE32+ Optional Header (Image Only)
16    pub optional_header_64: Option<OptionalHeader64>,
17    /// Table containing a list of section headers
18    pub section_table: Vec<SectionHeader>,
19}
20
21/// Parse a Portable Executable from a given byte array
22pub fn parse_portable_executable(binary: &[u8]) -> Result<PortableExecutable, Error> {
23    let mut offset = read_u16(binary, IMAGE_DOS_PE_SIGNATURE_OFFSET)?.into();
24
25    let slice = match binary.get(offset..offset+4) {
26        Some(slice) => slice,
27        None => {
28            return Err(Error::OffsetOutOfRange);
29        }
30    };
31
32    let string = match String::from_utf8(slice.to_vec()) {
33        Ok(string) => string,
34        Err(e) => {
35            return Err(Error::BadString(e));
36        }
37    };
38
39    if string != "PE\0\0" {
40        return Err(Error::MissingPeHeader);
41    }
42
43    offset += 4;
44
45    let mut pe: PortableExecutable = PortableExecutable { 
46        coff: CoffFileHeader::default(),
47        optional_header_32: None, 
48        optional_header_64: None, 
49        section_table: Vec::new()
50    };
51
52    let slice = match binary.get(offset..offset+20) {
53        Some(slice) => slice,
54        None => {
55            return Err(Error::OffsetOutOfRange);
56        }
57    };
58
59    pe.coff = match try_from_bytes::<CoffFileHeader>(slice) {
60        Ok(coff) => *coff,
61        Err(_) => {
62            return Err(Error::MissingCoffHeader);
63        }
64    };
65
66    offset += 20;
67
68    if pe.coff.size_of_optional_header != 0 {
69        let magic = match Magic::from_u16(read_u16(binary, offset)?) {
70            Some(magic) => magic,
71            None => {
72                return Err(Error::MissingMagicNumber);
73            }
74        };
75
76        match magic {
77            Magic::PE32 => {
78                pe.optional_header_32 = Some(OptionalHeader32::parse_optional_header(binary, &mut offset)?);
79            }
80            Magic::PE64 => {
81                pe.optional_header_64 = Some(OptionalHeader64::parse_optional_header(binary, &mut offset)?);
82            }
83        }
84    }
85
86    pe.section_table = parse_section_table(binary, offset, pe.coff.number_of_sections);
87
88    /*for section in pe.section_table.iter() {
89        let name = match section.get_name() {
90            Some(name) => name,
91            None => {
92                return Err(Error::new(ErrorKind::Other, "Failed to get section name"));
93            }
94        };
95
96        match name.trim_end_matches(char::from(0)) {
97            ".edata" => {
98                println!(".edata Section");
99            }
100            ".idata" => {
101                println!(".idata Section");
102            }
103            _ => {}
104        }
105    }*/
106
107    Ok(pe)
108}
109
110impl fmt::Display for PortableExecutable {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        writeln!(f, "{}", self.coff)?;
113
114        match self.optional_header_32 {
115            None => (),
116            Some(header) => {
117                writeln!(f, "{}", header)?;
118            }
119        }
120
121        match self.optional_header_64 {
122            None => (),
123            Some(header) => {
124                writeln!(f, "{}", header)?;
125            }
126        }
127
128        for section in self.section_table.iter() {
129            writeln!(f, "{}", section)?;
130        }
131
132        Ok(())
133    }
134}
135
136fn read_u16(binary: &[u8], offset: usize) -> Result<u16, Error> {
137    if let Some(array) = binary.get(offset..offset+2) {
138        if let Ok(slice) = array.try_into() {
139            Ok(u16::from_le_bytes(slice))
140        } else {
141            unreachable!()
142        }
143    } else {
144        Err(Error::OffsetOutOfRange)
145    }
146}
pe_parser/pe.rs

pe_parser/
pe.rs
