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
9pub const HEADER_LENGTH: u64 = 64;
12
13#[derive(Debug, Default)]
14pub struct DosHeader {
15 pub e_magic: HeaderField<u16>, e_cblp: HeaderField<u16>, e_cp: HeaderField<u16>, e_crlc: HeaderField<u16>, e_cparhdr: HeaderField<u16>, e_minalloc: HeaderField<u16>, e_maxalloc: HeaderField<u16>, e_ss: HeaderField<u16>, e_sp: HeaderField<u16>, e_csum: HeaderField<u16>, e_ip: HeaderField<u16>, e_cs: HeaderField<u16>, e_lfarlc: HeaderField<u16>, e_ovno: HeaderField<u16>, e_res: HeaderField<[u16; 4]>, e_oemid: HeaderField<u16>, e_oeminfo: HeaderField<u16>, e_res2: HeaderField<[u16; 10]>, pub e_lfanew: HeaderField<u32>, }
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}