pe-assembler 0.1.1

use gaia_binary::{LittleEndian, ReadBytesExt};
use gaia_types::GaiaError;
use serde::{Deserialize, Serialize};
use std::io::Read;

/// DOS header structure
///
/// Contains basic information about a DOS executable, the first structure in a PE file.
/// Although modern Windows programs don't run in DOS mode, the PE format still retains this structure for compatibility.
#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
pub struct DosHeader {
    /// DOS signature, usually 0x5A4D ("MZ")
    pub e_magic: u16,
    /// Number of bytes in the last page of the file
    pub e_cblp: u16,
    /// Total number of pages in the file
    pub e_cp: u16,
    /// Number of relocation items
    pub e_crlc: u16,
    /// Size of the header in paragraphs
    pub e_cparhdr: u16,
    /// Minimum number of paragraphs required
    pub e_min_allocate: u16,
    /// Maximum number of paragraphs required
    pub e_max_allocate: u16,
    /// Initial SS register value
    pub e_ss: u16,
    /// Initial SP register value
    pub e_sp: u16,
    /// Checksum
    pub e_check_sum: u16,
    /// Initial IP register value
    pub e_ip: u16,
    /// Initial CS register value
    pub e_cs: u16,
    /// File offset of the relocation table
    pub e_lfarlc: u16,
    /// Overlay number
    pub e_ovno: u16,
    /// Reserved fields, usually 0
    pub e_res: [u16; 4],
    /// OEM identifier
    pub e_oem_id: u16,
    /// OEM information
    pub e_oem_info: u16,
    /// Reserved fields, usually 0
    pub e_res2: [u16; 10],
    /// File offset of the PE header, pointing to the real PE structure
    pub e_lfanew: u32,
}

impl Default for DosHeader {
    fn default() -> Self {
        Self {
            e_magic: 0x5A4D, // "MZ"
            e_cblp: 0x90,
            e_cp: 0x03,
            e_crlc: 0x00,
            e_cparhdr: 0x04,
            e_min_allocate: 0x00,
            e_max_allocate: 0xFFFF,
            e_ss: 0x00,
            e_sp: 0xB8,
            e_check_sum: 0x00,
            e_ip: 0x00,
            e_cs: 0x00,
            e_lfarlc: 0x40,
            e_ovno: 0x00,
            e_res: [0; 4],
            e_oem_id: 0x00,
            e_oem_info: 0x00,
            e_res2: [0; 10],
            e_lfanew: 0x80, // PE header offset
        }
    }
}

impl DosHeader {
    /// Read DOS header from binary reader
    ///
    /// # Arguments
    /// * `reader` - Binary reader
    ///
    /// # Returns
    /// Returns DOS header structure or error
    pub fn read<R: Read>(mut reader: R) -> Result<Self, GaiaError> {
        let e_magic = reader.read_u16::<LittleEndian>()?;
        let e_cblp = reader.read_u16::<LittleEndian>()?;
        let e_cp = reader.read_u16::<LittleEndian>()?;
        let e_crlc = reader.read_u16::<LittleEndian>()?;
        let e_cparhdr = reader.read_u16::<LittleEndian>()?;
        let e_min_allocate = reader.read_u16::<LittleEndian>()?;
        let e_max_allocate = reader.read_u16::<LittleEndian>()?;
        let e_ss = reader.read_u16::<LittleEndian>()?;
        let e_sp = reader.read_u16::<LittleEndian>()?;
        let e_check_sum = reader.read_u16::<LittleEndian>()?;
        let e_ip = reader.read_u16::<LittleEndian>()?;
        let e_cs = reader.read_u16::<LittleEndian>()?;
        let e_lfarlc = reader.read_u16::<LittleEndian>()?;
        let e_ovno = reader.read_u16::<LittleEndian>()?;

        let mut e_res = [0u16; 4];
        for i in 0..4 {
            e_res[i] = reader.read_u16::<LittleEndian>()?;
        }

        let e_oem_id = reader.read_u16::<LittleEndian>()?;
        let e_oem_info = reader.read_u16::<LittleEndian>()?;

        let mut e_res2 = [0u16; 10];
        for i in 0..10 {
            e_res2[i] = reader.read_u16::<LittleEndian>()?;
        }

        let e_lfanew = reader.read_u32::<LittleEndian>()?;

        Ok(DosHeader {
            e_magic,
            e_cblp,
            e_cp,
            e_crlc,
            e_cparhdr,
            e_min_allocate,
            e_max_allocate,
            e_ss,
            e_sp,
            e_check_sum,
            e_ip,
            e_cs,
            e_lfarlc,
            e_ovno,
            e_res,
            e_oem_id,
            e_oem_info,
            e_res2,
            e_lfanew,
        })
    }
}

impl DosHeader {
    /// Create a standard DOS header with specified PE header offset
    pub fn new(header_offset: u32) -> Self {
        Self { e_lfanew: header_offset, ..Default::default() }
    }
}