malwaredb-types 0.3.4

// SPDX-License-Identifier: Apache-2.0

use crate::exec::{Architecture, Import, Imports, Section};
use crate::utils::{string_from_offset, u16_from_offset, u32_from_offset, u64_from_offset};
use crate::Ordering;

use std::fmt::{Debug, Display, Formatter};
use std::mem::size_of;
use std::time::{Duration, SystemTime};

use anyhow::Result;
use chrono::{DateTime, Utc};
use flagset::{flags, FlagSet};
use tracing::{info, instrument};

/// COFF header within a PE-32 file
#[repr(C, packed)]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct COFFHeader {
    /// Identified target machine
    pub machine: u16,

    /// The number of sections in the section table, after the headers.
    pub num_sections: u16,

    /// Seconds since midnight 1 Jan 1970
    pub time_date_stamp: u32,

    /// Offset to the COFF symbol table
    pub pointer_to_symbol_table: u32,

    /// Number of entries in the symbol table
    pub num_symbols: u32,

    /// Size of the Optional Header, which is required for executable files
    pub size_of_optional_header: u16,

    /// Flags which indicate the attributes of the file.
    pub characteristics: FlagSet<CoffCharacteristics>,
}

impl AsRef<[u8; size_of::<Self>()]> for COFFHeader {
    #[allow(clippy::transmute_ptr_to_ptr)]
    fn as_ref(&self) -> &[u8; size_of::<Self>()] {
        unsafe { std::mem::transmute::<_, &[u8; size_of::<Self>()]>(self) }
    }
}

flags! {
    /// COFF Characteristics flags
    pub enum CoffCharacteristics: u16 {
        /// File does not contain base relocations
        RelocsStripped = 0x0001,

        /// Image can run, linker error if not set
        ExecutableImage = 0x0002,

        /// COFF line numbers removed, deprecated
        LineNumsStripped = 0x0004,

        /// COFF symbol table entries have been removed, deprecated
        LocalSymbolsStripped = 0x0008,

        /// Aggressively trim working set, obsolete
        AggressiveTrimWorkingSet = 0x0010,

        /// Application can handle > 2GB addresses
        LargeAddressAware = 0x0020,

        /// Reserved for future use, should be zero
        Reserved = 0x0040,

        /// Little Endian, deprecated
        BytesReversedLo = 0x0080,

        /// Machine is based on 32-bit word architecture
        ThirtyTwoBitMachine = 0x0100,

        /// Debugging information removed from image
        DebugStripped = 0x0200,

        /// If the image is on removable media, load it fully and copy into swap
        RemovableRunFromSwap = 0x0400,

        /// If the image is on network media, load if fully and copy into swap
        NetRunFromSwap = 0x0800,

        /// Image is a system file, not a user program
        FileSystem = 0x1000,

        /// File is a Dynamic Link Library (DLL)
        FileDLL = 0x2000,

        /// Run only on uni-processor machines
        UPSystemOnly = 0x4000,

        /// Big Endian, deprecated
        BytesReversedHi = 0x8000,
    }
}

impl COFFHeader {
    /// Casts a byte buffer into a [`COFFHeader`].
    #[allow(clippy::transmute_ptr_to_ptr)]
    pub fn cast(raw: &[u8; size_of::<Self>()]) -> &Self {
        unsafe { std::mem::transmute(raw) }
    }

    pub fn from(contents: &[u8]) -> Result<&Self> {
        let temp: &[u8; 20] = contents[0..20].try_into()?;
        Ok(COFFHeader::cast(temp))
    }

    pub fn architecture(&self) -> Architecture {
        match self.machine {
            0x184 => Architecture::Alpha,
            0x284 => Architecture::Alpha64,
            0x1c0 => Architecture::ARM,
            0x1c4 | 0x1c2 => Architecture::ARMThumb,
            0xAA64 | 0xA641 | 0xA64E => Architecture::ARM64, // 0xA641 and 0xA64E are for ARM/Intel interop and hybrids
            0x1a2 | 0x1a3 => Architecture::HitachiSH3,       // 0x1a3 is SH3 DSP
            0x1a6 => Architecture::HitachiSH4,
            0x1a8 => Architecture::HitachiSH5,
            0x6232 => Architecture::LoongArch32,
            0x6264 => Architecture::LoongArch64,
            0x169 | 0x266 | 0x366 | 0x466 | 0x166 => Architecture::MIPSEL,
            0x1f0 | 0x1f1 => Architecture::PowerPCLE, // 0x1f1 is PPCLE with FPU
            0x5032 => Architecture::RISCV,
            0x5064 => Architecture::RISCV64,
            0x5128 => Architecture::RISCV128,
            0x14c => Architecture::X86,
            0x8664 => Architecture::X86_64,
            0 => Architecture::Unknown,
            other => Architecture::Other(u32::from(other)),
        }
    }

    /// Compilation date as UTC
    pub fn compiled_date(&self) -> DateTime<Utc> {
        let d = SystemTime::UNIX_EPOCH + Duration::from_secs(u64::from(self.time_date_stamp));
        DateTime::<Utc>::from(d)
    }
}

/// PE-32 Optional header
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum OptionalHeader {
    OptionalHeader32(OptionalHeader32),
    OptionalHeader64(OptionalHeader64),
}

impl OptionalHeader {
    /// Optional Header from a byte sequence
    pub fn from(contents: &[u8]) -> Option<Self> {
        match u16_from_offset(contents, 0, Ordering::LittleEndian) {
            Some(0x020B) => Some(OptionalHeader::OptionalHeader64(OptionalHeader64::from(
                contents,
            ))),
            Some(_) => Some(OptionalHeader::OptionalHeader32(OptionalHeader32::from(
                contents,
            ))),
            None => None,
        }
    }

    /// PE-32 subsystem as an integer
    pub fn subsystem(&self) -> Option<u16> {
        match self {
            OptionalHeader::OptionalHeader32(h32) => h32.subsystem,
            OptionalHeader::OptionalHeader64(h64) => h64.subsystem,
        }
    }

    /// If the program is 64-bit
    pub fn is_64bit(&self) -> bool {
        match self {
            OptionalHeader::OptionalHeader32(_) => false,
            OptionalHeader::OptionalHeader64(_) => true,
        }
    }

    /// Data directories
    pub fn data_directories(&self) -> Option<DataDirectories> {
        match self {
            OptionalHeader::OptionalHeader32(h32) => h32.data_directories,
            OptionalHeader::OptionalHeader64(h64) => h64.data_directories,
        }
    }
}

/// Import Directory data for a PE-32 binary
#[repr(C)]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct ImportDirectoryEntry {
    /// Contains the RVA of the import table
    pub import_lookup_table_rva: u32,

    /// The date and time stamp of the DLL, set after the image is loaded in to memory
    pub date_time_stamp: u32,

    /// The index of the first forwarder reference
    pub forwarder_chain: u32,

    /// The address of an ASCII string which contains the name of the DLL
    pub name_rva: u32,

    /// Address of the import address table
    pub import_address_table_rva: u32,
}

/// Entry in the Import Directory for a PE-32 binary
impl ImportDirectoryEntry {
    /// Casts a byte buffer into a [`ImportDirectoryEntry`].
    #[allow(clippy::transmute_ptr_to_ptr)]
    pub fn cast(raw: &[u8; size_of::<Self>()]) -> &Self {
        unsafe { std::mem::transmute(raw) }
    }

    pub fn from(contents: &[u8]) -> &Self {
        let temp: &[u8; 20] = contents[0..20].try_into().unwrap();
        ImportDirectoryEntry::cast(temp)
    }
}

/// Data Directory for a PE-32 binary
#[repr(C)]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct DataDirectory {
    /// Virtual address, location in memory
    pub virtual_address: u32,

    /// Data size
    pub size: u32,

    /// Name of the data directory
    pub name: &'static str,
}

impl DataDirectory {
    #[inline]
    pub fn new(
        virtual_address: Option<u32>,
        size: Option<u32>,
        name: &'static str,
    ) -> Option<Self> {
        Some(Self {
            virtual_address: virtual_address?,
            size: size?,
            name,
        })
    }
}

impl Display for DataDirectory {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}: Address: 0x{:02X}, Size: 0x{:02X}",
            self.name, self.virtual_address, self.size
        )
    }
}

/// Entry in a Data Directory for a PE-32 binary
#[repr(C)]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct DataDirectories {
    /// Exports
    pub export_table: DataDirectory,

    /// Imports
    pub import_table: DataDirectory,

    /// Resources
    pub resource_table: DataDirectory,

    /// Exceptions
    pub exception_table: DataDirectory,

    /// Certificates, the virtual address here is a file pointer
    pub certificate_table: DataDirectory,

    /// Entries for all base relocations in the image
    pub base_relocation_table: DataDirectory,

    /// Debugging data
    pub debug: DataDirectory,

    /// Reserved, must be zero
    pub architecture: DataDirectory,

    /// Location of global pointer register
    pub global_ptr: DataDirectory,

    /// Thread local storage
    pub tls_table: DataDirectory,

    /// Load configuration table
    pub load_config_table: DataDirectory,

    /// Bound import table
    pub bound_import: DataDirectory,

    /// Import address table
    pub iat: DataDirectory,

    /// Delay import descriptor
    pub delay_import_descriptor: DataDirectory,

    /// Dot net runtime header
    pub clr_runtime_header: DataDirectory,

    /// Reserved, must be zero
    pub reserved: DataDirectory,

    /// Whether or not the directories are from a 64-bit PE file
    is_64bit: bool,
}

/// To be loaded from the Optional Header, same for both 32- and 64-bit
impl DataDirectories {
    #[allow(clippy::too_many_lines)]
    pub fn from(contents: &[u8], offset: usize, is_64bit: bool) -> Option<Self> {
        if contents.len() < offset + 128 {
            return None;
        }
        // The integer from bytes conversions below should succeed since the lines above
        // ensure the buffer is large enough.
        Some(Self {
            export_table: DataDirectory::new(
                u32_from_offset(contents, offset, Ordering::LittleEndian),
                u32_from_offset(contents, offset + 4, Ordering::LittleEndian),
                "Export Directory",
            )?,
            import_table: DataDirectory::new(
                u32_from_offset(contents, offset + 8, Ordering::LittleEndian),
                u32_from_offset(contents, offset + 12, Ordering::LittleEndian),
                "Import Directory",
            )?,
            resource_table: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 16, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 20, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Resource Directory",
            },
            exception_table: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 24, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 28, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Exception Directory",
            },
            certificate_table: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 32, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 36, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Certificate Directory",
            },
            base_relocation_table: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 40, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 44, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Base Relocation Table",
            },
            debug: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 48, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 52, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Debug Directory",
            },
            architecture: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 56, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 60, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Architecture Specific Data",
            },
            global_ptr: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 64, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 68, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "RVA of Global Pointer Register",
            },
            tls_table: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 72, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 76, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Thread Local Storage Directory",
            },
            load_config_table: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 80, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 84, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Load Configuration Directory",
            },
            bound_import: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 88, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 92, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Bound Import Directory",
            },
            iat: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 96, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 100, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Import Address Table",
            },
            delay_import_descriptor: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 104, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 108, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Delay Load Import Descriptors",
            },
            clr_runtime_header: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 112, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 116, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: ".Net (Common Language Runtime) Header",
            },
            reserved: DataDirectory {
                virtual_address: u32_from_offset(contents, offset + 120, Ordering::LittleEndian)
                    .unwrap_or_default(),
                size: u32_from_offset(contents, offset + 124, Ordering::LittleEndian)
                    .unwrap_or_default(),
                name: "Reserved (should be zero)",
            },
            is_64bit,
        })
    }

    /// Get imports for a PE-32 binary
    #[allow(clippy::cast_possible_truncation)]
    #[instrument]
    pub fn imports(&self, num_imports: usize, section: &Section) -> Imports {
        const IMPORT_ENTRY_SIZE: usize = 20usize;

        let mut imports = Vec::new();

        if self.import_table.size == 0
            || self.import_table.virtual_address == 0
            || section.data.is_none()
        {
            return Imports {
                imports,
                expected_imports: num_imports as u32,
            };
        }

        if let Some(section_data) = section.data {
            for import_num in 0..num_imports {
                let offset = import_num * IMPORT_ENTRY_SIZE;

                if offset > section_data.len() || offset + IMPORT_ENTRY_SIZE > section_data.len() {
                    // If these offsets are too large for the section data, then no other iterations
                    // will be possible either.
                    info!("PE32: Offset {offset} or offset + import entry size {IMPORT_ENTRY_SIZE} is beyond section data size {}.", section_data.len());
                    break;
                }

                let import_directory =
                    ImportDirectoryEntry::from(&section_data[offset..offset + IMPORT_ENTRY_SIZE]);

                if import_directory.name_rva == 0 {
                    continue;
                }

                if import_directory.name_rva < section.virtual_address {
                    continue;
                }

                if import_directory.import_lookup_table_rva < section.virtual_address {
                    continue;
                }

                let dll_name_offset =
                    import_directory.name_rva as usize - section.virtual_address as usize;
                let dll_name = if dll_name_offset < section_data.len() {
                    string_from_offset(section_data, dll_name_offset).unwrap_or_default()
                } else {
                    info!("PE2: failed to get dll string name, offset {dll_name_offset} beyond buffer size {}.", section_data.len());
                    String::new()
                };

                let import_data_offset =
                    (import_directory.import_lookup_table_rva - section.virtual_address) as usize;
                if 4 + import_data_offset > section_data.len() {
                    info!(
                        "PE32: Import data offset {import_data_offset} beyond section data buffer size {}.",
                        section_data.len()
                    );
                    continue;
                }

                let mut the_imports = &section_data[import_data_offset..];
                loop {
                    if the_imports.is_empty() {
                        info!("PE32: Import data buffer is empty at {import_data_offset}.");
                        break;
                    }
                    let va = if self.is_64bit {
                        u64_from_offset(the_imports, 0, Ordering::LittleEndian)
                    } else {
                        u32_from_offset(the_imports, 0, Ordering::LittleEndian).map(u64::from)
                    }
                    .unwrap_or_default();

                    if va == 0 {
                        break;
                    }
                    if (va & 0x8000_0000_0000_0000) > 0 {
                        info!("PE32: binary uses ordinals, which isn't yet implemented.");
                        // TODO: Implement ordinal parsing
                    } else {
                        let function_name = if section.virtual_address as usize + 2 < va as usize {
                            let function_name_offset =
                                va as usize - section.virtual_address as usize + 2;
                            if function_name_offset < section_data.len() {
                                string_from_offset(section_data, function_name_offset)
                                    .unwrap_or_default()
                            } else {
                                info!("PE32: function name offset {function_name_offset} beyond section data size {}.", section_data.len());
                                String::new()
                            }
                        } else {
                            info!(
                                "PE32: function name: virtual address {}+2 is greater than va {va}.",
                                section.virtual_address
                            );
                            String::new()
                        };

                        imports.push(Import {
                            library: dll_name.clone(),
                            function: function_name,
                        });
                    }

                    the_imports = &the_imports[4..];
                    if self.is_64bit {
                        the_imports = &the_imports[4..];
                    }
                }
            }
        } else {
            info!("PE32: Import section data buffer is empty.");
        }

        Imports {
            imports,
            expected_imports: num_imports as u32,
        }
    }
}

/// Data Directory iterator convenience object
pub struct DataDirectoryIterator {
    item: DataDirectories,
    index: usize,
}

impl IntoIterator for DataDirectories {
    type Item = DataDirectory;
    type IntoIter = DataDirectoryIterator;

    fn into_iter(self) -> Self::IntoIter {
        DataDirectoryIterator {
            item: self,
            index: 0,
        }
    }
}

impl Iterator for DataDirectoryIterator {
    type Item = DataDirectory;

    fn next(&mut self) -> Option<Self::Item> {
        let result = match self.index {
            0 => self.item.export_table,
            1 => self.item.import_table,
            2 => self.item.resource_table,
            3 => self.item.exception_table,
            4 => self.item.certificate_table,
            5 => self.item.base_relocation_table,
            6 => self.item.debug,
            7 => self.item.architecture,
            8 => self.item.global_ptr,
            9 => self.item.tls_table,
            10 => self.item.load_config_table,
            11 => self.item.bound_import,
            12 => self.item.iat,
            13 => self.item.delay_import_descriptor,
            14 => self.item.clr_runtime_header,
            15 => self.item.reserved,
            _ => {
                return None;
            }
        };
        self.index += 1;
        Some(result)
    }

    fn count(self) -> usize
    where
        Self: Sized,
    {
        16
    }

    fn last(self) -> Option<Self::Item>
    where
        Self: Sized,
    {
        // Technically not the last, but getting the CRL runtime
        // header is far more useful than getting the reserved
        // entry, which should just be zeroes.
        Some(self.item.clr_runtime_header)
    }
}

impl Display for DataDirectories {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for directory in self.into_iter() {
            writeln!(f, "{directory}")?;
        }

        Ok(())
    }
}

/// Optional Header for 32-bit PE-32 files
#[repr(C)]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct OptionalHeader32 {
    /// Magic: `0x10B` for 32-bit, `0x107` for ROM, `0x20B` for 64-bit
    pub magic: u16,

    /// Linker major version
    pub major_linker_version: u8,

    /// Linker minor version
    pub minor_linker_version: u8,

    /// Size of the ".text" section, or sum of multiple code sections
    pub size_of_code: u32,

    /// Size of initialized data section, or sum if there are multiple such sections
    pub size_of_initialized_data: u32,

    /// Size of the uninitialized data section, or sum if there are multiple such sections
    pub size_of_uninitialized_data: u32,

    /// Entry of starting code for executables or device drivers, zero for libraries
    pub address_of_entry_point: u32,

    /// The address that is relative to the image base of the beginning of code section
    pub base_of_code: u32,

    /// The address that is relative to the image base of the beginning of data section.
    pub base_of_data: u32,

    /// Address of the first byte of image when in memory
    pub image_base: Option<u32>,

    /// Byte alignment of sections when loaded in memory
    pub section_alignment: Option<u32>,

    /// Byte alignment of sections in the image file
    pub file_alignment: Option<u32>,

    /// The major version of the required operating system
    pub major_operating_system_version: Option<u16>,

    /// The minor version of the required operating system
    pub minor_operating_system_version: Option<u16>,

    /// The major version of the image
    pub major_image_version: Option<u16>,

    /// The minor version of the image
    pub minor_image_version: Option<u16>,

    /// The major version of the subsystem
    pub major_subsystem_version: Option<u16>,

    /// The minor version of the subsystem
    pub minor_subsystem_version: Option<u16>,

    /// Reserved, must be zero
    pub win32_version_value: Option<u32>,

    /// Byte size of the image with headers in memory
    pub size_of_image: Option<u32>,

    /// Combined size of the DOS stub, PE header, and section headers
    pub size_of_headers: Option<u32>,

    /// Image file checksum
    pub checksum: Option<u32>,

    /// The subsystem that is required to run this image
    pub subsystem: Option<u16>,

    /// Additional flags for the DLL
    pub dll_characteristics: Option<FlagSet<DLLCharacteristics>>,

    /// Size of the stack to reserve
    pub size_of_stack_reserve: Option<u32>,

    /// Size of the stack to commit
    pub size_of_stack_commit: Option<u32>,

    /// Size of the local heap space to reserve
    pub size_of_heap_reserve: Option<u32>,

    /// Size of the local heap space to commit
    pub size_of_heap_commit: Option<u32>,

    /// Reserved, must be zero
    pub loader_flags: Option<u32>,

    /// Number of data-directory entries in the remainder of the optional header
    pub num_rva_and_sizes: Option<u32>,

    /// Additional data
    pub data_directories: Option<DataDirectories>,
}

impl OptionalHeader32 {
    pub fn from(contents: &[u8]) -> Self {
        Self {
            magic: u16_from_offset(contents, 0, Ordering::LittleEndian).unwrap_or_default(),
            major_linker_version: contents[2],
            minor_linker_version: contents[3],
            size_of_code: u32_from_offset(contents, 4, Ordering::LittleEndian).unwrap_or_default(),
            size_of_initialized_data: u32_from_offset(contents, 8, Ordering::LittleEndian)
                .unwrap_or_default(),
            size_of_uninitialized_data: u32_from_offset(contents, 12, Ordering::LittleEndian)
                .unwrap_or_default(),
            address_of_entry_point: u32_from_offset(contents, 16, Ordering::LittleEndian)
                .unwrap_or_default(),
            base_of_code: u32_from_offset(contents, 20, Ordering::LittleEndian).unwrap_or_default(),
            base_of_data: u32_from_offset(contents, 24, Ordering::LittleEndian).unwrap_or_default(),
            image_base: u32_from_offset(contents, 28, Ordering::LittleEndian),
            section_alignment: u32_from_offset(contents, 32, Ordering::LittleEndian),
            file_alignment: u32_from_offset(contents, 36, Ordering::LittleEndian),
            major_operating_system_version: u16_from_offset(contents, 40, Ordering::LittleEndian),
            minor_operating_system_version: u16_from_offset(contents, 42, Ordering::LittleEndian),
            major_image_version: u16_from_offset(contents, 44, Ordering::LittleEndian),
            minor_image_version: u16_from_offset(contents, 46, Ordering::LittleEndian),
            major_subsystem_version: u16_from_offset(contents, 48, Ordering::LittleEndian),
            minor_subsystem_version: u16_from_offset(contents, 50, Ordering::LittleEndian),
            win32_version_value: u32_from_offset(contents, 52, Ordering::LittleEndian),
            size_of_image: u32_from_offset(contents, 56, Ordering::LittleEndian),
            size_of_headers: u32_from_offset(contents, 60, Ordering::LittleEndian),
            checksum: u32_from_offset(contents, 64, Ordering::LittleEndian),
            subsystem: u16_from_offset(contents, 68, Ordering::LittleEndian),
            dll_characteristics: {
                if let Some(val) = u16_from_offset(contents, 70, Ordering::LittleEndian) {
                    FlagSet::<DLLCharacteristics>::new(val).ok()
                } else {
                    None
                }
            },
            size_of_stack_reserve: u32_from_offset(contents, 72, Ordering::LittleEndian),
            size_of_stack_commit: u32_from_offset(contents, 76, Ordering::LittleEndian),
            size_of_heap_reserve: u32_from_offset(contents, 80, Ordering::LittleEndian),
            size_of_heap_commit: u32_from_offset(contents, 84, Ordering::LittleEndian),
            loader_flags: u32_from_offset(contents, 88, Ordering::LittleEndian),
            num_rva_and_sizes: u32_from_offset(contents, 92, Ordering::LittleEndian),
            data_directories: DataDirectories::from(contents, 96, false),
        }
    }
}

/// Optional Header for 64-bit PE-32 files
#[repr(C)]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct OptionalHeader64 {
    /// Magic: `0x10B` for 32-bit, `0x107` for ROM, `0x20B` for 64-bit
    pub magic: u16,

    /// Linker major version
    pub major_linker_version: u8,

    /// Linker minor version
    pub minor_linker_version: u8,

    /// Size of the ".text" section, or sum of multiple code sections
    pub size_of_code: u32,

    /// Size of initialized data section, or sum if there are multiple such sections
    pub size_of_initialized_data: u32,

    /// Size of the uninitialized data section, or sum if there are multiple such sections
    pub size_of_uninitialized_data: u32,

    /// Entry of starting code for executables or device drivers, zero for libraries
    pub address_of_entry_point: u32,

    /// The address that is relative to the image base of the beginning of code section
    pub base_of_code: u32,

    /// Address of the first byte of image when in memory
    pub image_base: Option<u64>,

    /// Byte alignment of sections when loaded in memory
    pub section_alignment: Option<u32>,

    /// Byte alignment of sections in the image file
    pub file_alignment: Option<u32>,

    /// The major version of the required operating system
    pub major_operating_system_version: Option<u16>,

    /// The minor version of the required operating system
    pub minor_operating_system_version: Option<u16>,

    /// The major version of the image
    pub major_image_version: Option<u16>,

    /// The minor version of the image
    pub minor_image_version: Option<u16>,

    /// The major version of the subsystem
    pub major_subsystem_version: Option<u16>,

    /// The minor version of the subsystem
    pub minor_subsystem_version: Option<u16>,

    /// Reserved, must be zero
    pub win32_version_value: Option<u32>,

    /// Byte size of the image with headers in memory
    pub size_of_image: Option<u32>,

    /// Combined size of the DOS stub, PE header, and section headers
    pub size_of_headers: Option<u32>,

    /// Image file checksum
    pub checksum: Option<u32>,

    /// The subsystem that is required to run this image
    pub subsystem: Option<u16>,

    /// Additional flags for the DLL
    pub dll_characteristics: Option<FlagSet<DLLCharacteristics>>,

    /// Size of the stack to reserve
    pub size_of_stack_reserve: Option<u64>,

    /// Size of the stack to commit
    pub size_of_stack_commit: Option<u64>,

    /// Size of the local heap space to reserve
    pub size_of_heap_reserve: Option<u64>,

    /// Size of the local heap space to commit
    pub size_of_heap_commit: Option<u64>,

    /// Reserved, must be zero
    pub loader_flags: Option<u32>,

    /// Number of data-directory entries in the remainder of the optional header
    pub num_rva_and_sizes: Option<u32>,

    /// Additional data
    pub data_directories: Option<DataDirectories>,
}

impl OptionalHeader64 {
    /// 64-bit Optional Header from a byte sequence
    pub fn from(contents: &[u8]) -> Self {
        Self {
            magic: u16_from_offset(contents, 0, Ordering::LittleEndian).unwrap_or_default(),
            major_linker_version: contents[2],
            minor_linker_version: contents[3],
            size_of_code: u32_from_offset(contents, 4, Ordering::LittleEndian).unwrap_or_default(),
            size_of_initialized_data: u32_from_offset(contents, 8, Ordering::LittleEndian)
                .unwrap_or_default(),
            size_of_uninitialized_data: u32_from_offset(contents, 12, Ordering::LittleEndian)
                .unwrap_or_default(),
            address_of_entry_point: u32_from_offset(contents, 16, Ordering::LittleEndian)
                .unwrap_or_default(),
            base_of_code: u32_from_offset(contents, 20, Ordering::LittleEndian).unwrap_or_default(),
            image_base: u64_from_offset(contents, 24, Ordering::LittleEndian),
            section_alignment: u32_from_offset(contents, 32, Ordering::LittleEndian),
            file_alignment: u32_from_offset(contents, 36, Ordering::LittleEndian),
            major_operating_system_version: u16_from_offset(contents, 40, Ordering::LittleEndian),
            minor_operating_system_version: u16_from_offset(contents, 42, Ordering::LittleEndian),
            major_image_version: u16_from_offset(contents, 44, Ordering::LittleEndian),
            minor_image_version: u16_from_offset(contents, 46, Ordering::LittleEndian),
            major_subsystem_version: u16_from_offset(contents, 48, Ordering::LittleEndian),
            minor_subsystem_version: u16_from_offset(contents, 50, Ordering::LittleEndian),
            win32_version_value: u32_from_offset(contents, 52, Ordering::LittleEndian),
            size_of_image: u32_from_offset(contents, 56, Ordering::LittleEndian),
            size_of_headers: u32_from_offset(contents, 60, Ordering::LittleEndian),
            checksum: u32_from_offset(contents, 64, Ordering::LittleEndian),
            subsystem: u16_from_offset(contents, 68, Ordering::LittleEndian),
            dll_characteristics: {
                if let Some(val) = u16_from_offset(contents, 70, Ordering::LittleEndian) {
                    FlagSet::<DLLCharacteristics>::new(val).ok()
                } else {
                    None
                }
            },
            size_of_stack_reserve: u64_from_offset(contents, 72, Ordering::LittleEndian),
            size_of_stack_commit: u64_from_offset(contents, 80, Ordering::LittleEndian),
            size_of_heap_reserve: u64_from_offset(contents, 88, Ordering::LittleEndian),
            size_of_heap_commit: u64_from_offset(contents, 96, Ordering::LittleEndian),
            loader_flags: u32_from_offset(contents, 104, Ordering::LittleEndian),
            num_rva_and_sizes: u32_from_offset(contents, 108, Ordering::LittleEndian),
            data_directories: DataDirectories::from(contents, 112, true),
        }
    }
}

flags! {
    /// DLL Characteristics flags
    pub enum DLLCharacteristics: u16 {
        /// Reserved, must be zero
        Reserved1 = 0x0001,

        /// Reserved, must be zero
        Reserved2 = 0x0002,

        /// Reserved, must be zero
        Reserved3 = 0x0004,

        /// Reserved, must be zero
        Reserved4 = 0x0008,

        /// Image can handle high entropy 64-bit virtual address space
        HighEntropyVA = 0x0020,

        /// DLL can be relocated at load time
        DynamicBase = 0x0040,

        /// Code Integrity checks are enforced
        ForceIntegrity = 0x0080,

        /// Image is NX (No Execute)-compatible
        NXCompat = 0x0100,

        /// Isolation aware, but don't isolate
        NoIsolation = 0x0200,

        /// Does not use structured exception handling
        NoSEH = 0x0400,

        /// Do not bind the image
        NoBind = 0x0800,

        /// Must execute in an AppContainer
        AppContainer = 0x1000,

        /// Is a WDM (Windows Driver Model) driver
        WDMDriver = 0x2000,

        /// Image supports Control Flow Guard
        GuardCF = 0x4000,

        /// Terminal Server Aware
        TerminalServerAware = 0x8000,
    }

    /// Section Flags
    pub enum SectionFlags: u32 {
        /// Reserved for future use
        Reserved1 = 0x0000_0000,

        /// Reserved for future use
        Reserved2 = 0x0000_0001,

        /// Reserved for future use
        Reserved3 = 0x0000_0002,

        /// Reserved for future use
        Reserved4 = 0x0000_0004,

        /// Obsolete: Section is not padded to next boundary
        NoPadding = 0x0000_0008,

        /// Reserved for future use
        Reserved5 = 0x0000_0010,

        /// Contains executable code
        ContainsCode = 0x0000_0020,

        /// Contains initialized data
        InitializedData = 0x0000_0040,

        /// Contains uninitialized data
        UninitializedData = 0x0000_0080,

        /// Link other, reserved for future use
        LinkOther = 0x0000_0100,

        /// Contains comments or other information
        LinkInfo = 0x0000_0200,

        /// Reserved for future use
        Reserved6 = 0x0000_0400,

        /// Will not become part of the image, for object files
        Remove = 0x0000_0800,

        /// Contains COMDAT data
        Comdat = 0x0000_1000,

        /// Contains data referenced through the global pointer
        GlobalPointerReferenced = 0x0000_8000,

        /// Reserved for future use
        MemPurgeable = 0x0002_0000,

        /// Reserved for future use
        Mem16bit = 0x0002_0000,

        /// Reserved for future use
        MemLocked = 0x0004_0000,

        /// Reserved for future use
        MemPreload = 0x0008_0000,

        /// Object files: align on 1-byte boundary
        Align1Bytes = 0x0010_0000,

        /// Object files: align on 2-byte boundary
        Align2Bytes = 0x0020_0000,

        /// Object files: align on 4-byte boundary
        Align4Bytes = 0x0030_0000,

        /// Object files: align on 8-byte boundary
        Align8Bytes = 0x0040_0000,

        /// Object files: align on 16-byte boundary
        Align16Bytes = 0x0050_0000,

        /// Object files: align on 32-byte boundary
        Align32Bytes = 0x0060_0000,

        /// Object files: align on 64-byte boundary
        Align64Bytes = 0x0070_0000,

        /// Object files: align on 128-byte boundary
        Align128Bytes = 0x0080_0000,

        /// Object files: align on 256-byte boundary
        Align256Bytes = 0x0090_0000,

        /// Object files: align on 512-byte boundary
        Align512Bytes = 0x00A0_0000,

        /// Object files: align on 1024-byte boundary
        Align1024Bytes = 0x00B0_0000,

        /// Object files: align on 2048-byte boundary
        Align2048Bytes = 0x00C0_0000,

        /// Object files: align on 4096-byte boundary
        Align4096Bytes = 0x00D0_0000,

        /// Object files: align on 8192-byte boundary
        Align8192Bytes = 0x00E0_0000,

        /// Contains extended relocations
        ExtendedRel = 0x0100_0000,

        /// Can be discarded as needed
        Discardable = 0x0200_0000,

        /// Cannot be cached
        NotCacheable = 0x0400_0000,

        /// Cannot be paged
        NotPageable = 0x0800_0000,

        /// Can be shared in memory
        Shareable = 0x1000_0000,

        /// Can be executed as code
        Executable = 0x2000_0000,

        /// Can be read
        Readable = 0x4000_0000,

        /// Can be written
        Writeable = 0x8000_0000,
    }
}