rustbin/pe/
dos.rs

1use crate::{new_header_field, types::{Header, HeaderField}};
2
3use std::{io::Cursor, fmt::Display};
4
5use byteorder::{LittleEndian, ReadBytesExt};
6
7use super::PeError;
8
9//#[allow(unused)]
10
11pub const HEADER_LENGTH: u64 = 64;
12
13#[derive(Debug, Default)]
14pub struct DosHeader {
15    pub e_magic: HeaderField<u16>,    // Magic number
16    e_cblp: HeaderField<u16>,         // Bytes on last page of file
17    e_cp: HeaderField<u16>,           // Pages in file
18    e_crlc: HeaderField<u16>,         // Relocations
19    e_cparhdr: HeaderField<u16>,      // Size of header in paragraphs
20    e_minalloc: HeaderField<u16>,     // Minimum extra paragraphs needed
21    e_maxalloc: HeaderField<u16>,     // Maximum extra paragraphs needed
22    e_ss: HeaderField<u16>,           // Initial (relative) SS value
23    e_sp: HeaderField<u16>,           // Initial SP value
24    e_csum: HeaderField<u16>,         // Checksum
25    e_ip: HeaderField<u16>,           // Initial IP value
26    e_cs: HeaderField<u16>,           // Initial (relative) CS value
27    e_lfarlc: HeaderField<u16>,       // File address of relocation table
28    e_ovno: HeaderField<u16>,         // Overlay number
29    e_res: HeaderField<[u16; 4]>,     // Reserved words
30    e_oemid:  HeaderField<u16>,       // OEM identifier (for e_oeminfo)
31    e_oeminfo: HeaderField<u16>,      // OEM information; e_oemid specific
32    e_res2: HeaderField<[u16; 10]>,   // Reserved words
33    pub e_lfanew: HeaderField<u32>,   // File address of new exe header
34}
35
36impl DosHeader {
37    pub fn new() -> Self {
38        DosHeader {
39            e_magic: Default::default(),     
40            e_cblp: Default::default(),      
41            e_cp: Default::default(),        
42            e_crlc: Default::default(),      
43            e_cparhdr: Default::default(),   
44            e_minalloc: Default::default(),  
45            e_maxalloc: Default::default(),  
46            e_ss: Default::default(),        
47            e_sp: Default::default(),        
48            e_csum: Default::default(),      
49            e_ip: Default::default(),        
50            e_cs: Default::default(),        
51            e_lfarlc: Default::default(),    
52            e_ovno: Default::default(),      
53            e_res: Default::default(),  
54            e_oemid: Default::default(),     
55            e_oeminfo: Default::default(),   
56            e_res2: Default::default(),
57            e_lfanew: Default::default(),
58        }
59    }
60}
61
62impl Header for DosHeader {
63    fn parse_bytes(bytes: Vec<u8>, pos: u64) -> crate::Result<Self> {
64        let bytes_available = (bytes.len() as u64) - pos;
65
66        if bytes_available < HEADER_LENGTH {
67            return Err ( 
68                PeError::BufferTooSmall { target: "DosHeader".into(), expected: HEADER_LENGTH, actual: bytes_available}
69            );
70        }
71
72        let mut cursor = Cursor::new(bytes);
73        let mut offset = pos;
74        let mut dos_header = Self::new();
75
76        dos_header.e_magic = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
77        dos_header.e_cblp = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);    
78        dos_header.e_cp = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
79        dos_header.e_crlc = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
80        dos_header.e_cparhdr = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
81        dos_header.e_minalloc = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
82        dos_header.e_maxalloc = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
83        dos_header.e_ss = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
84        dos_header.e_sp = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
85        dos_header.e_csum = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
86        dos_header.e_ip = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
87        dos_header.e_cs = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
88        dos_header.e_lfarlc = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
89        dos_header.e_ovno = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
90        
91        let mut tmp_buff: [u16; 4] = [0; 4];
92        cursor.read_u16_into::<LittleEndian>(&mut tmp_buff)?;
93        dos_header.e_res = new_header_field!(tmp_buff, offset);
94        dos_header.e_oemid = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
95        dos_header.e_oeminfo = new_header_field!(cursor.read_u16::<LittleEndian>()?, offset);
96        
97        let mut tmp_buff: [u16; 10] = [0; 10];
98        cursor.read_u16_into::<LittleEndian>(&mut tmp_buff)?;
99        dos_header.e_res2 = new_header_field!(tmp_buff, offset);
100
101        dos_header.e_lfanew = new_header_field!(cursor.read_u32::<LittleEndian>()?, offset);
102        Ok(dos_header)
103    }
104
105    fn is_valid(&self) -> bool {
106        self.e_magic.value == 0x5A4D
107    }
108    
109    fn length() -> usize { HEADER_LENGTH as usize}
110}
111
112impl Display for DosHeader {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        write!(f, "{{e_magic: '{}', e_lfanew: {}(0x{:X})}}", std::str::from_utf8(&self.e_magic.value.to_le_bytes()).unwrap_or("ERR"), self.e_lfanew.value, self.e_lfanew.value)
115    }
116}
117
118
119#[cfg(test)]
120mod tests {
121    use crate::types::Header;
122
123    use super::DosHeader;    
124    const RAW_DOS_BYTES: [u8; 64] = [0x4D, 0x5A, 0x90, 0x00, 0x03, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 
125                                    0x00, 0x00, 0xB8, 0x00, 00, 00, 00, 00, 00, 00, 0x40, 00, 00, 00, 00, 00, 00, 00, 
126                                    00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 00, 
127                                    00, 00, 00, 00, 00, 00, 00, 0xF8, 00, 00, 00];
128    #[test]
129    fn parse_valid_header(){
130        let dos_header = DosHeader::parse_bytes(RAW_DOS_BYTES.to_vec(), 0).unwrap();
131        assert!(dos_header.is_valid());
132        assert_eq!(dos_header.e_magic.value, 0x5A4D);
133        assert_eq!(dos_header.e_magic.offset, 0);
134        assert_eq!(dos_header.e_magic.rva, 0);
135        assert_eq!(dos_header.e_lfanew.value, 0x000000F8);
136        assert_eq!(dos_header.e_lfanew.offset, 60);
137        assert_eq!(dos_header.e_lfanew.rva, 60);
138    }
139
140    #[test]
141    fn parse_invalid_header(){
142        let mut buf = RAW_DOS_BYTES.to_vec();
143        buf[0] = 0x4E;
144        let dos_header = DosHeader::parse_bytes(buf, 0).unwrap();
145        assert!(dos_header.is_valid() == false);
146    }
147}