mod fields;
mod utils;

use crate::exec::{
    Architecture, ExecutableFile, ExecutableType, Imports, OperatingSystem, Ordering, Section,
    Sections,
};
use crate::utils::{bytes_offset_match, u32_from_offset, EntropyCalc};
use crate::TypeMagic;

use std::fmt::{Display, Formatter};

use anyhow::{bail, Context, Result};
use chrono::{DateTime, Utc};
use flagset::FlagSet;
use tracing::debug;

// Microsoft documentation:
// https://learn.microsoft.com/en-us/windows/win32/debug/pe-format
// Wikipedia entry:
// https://en.wikipedia.org/wiki/Portable_Executable

const EXE_MAGIC: [u8; 2] = [b'M', b'Z'];
const PE_MAGIC: [u8; 4] = [0x50, 0x45, 0x00, 0x00];

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SubType {
    /// MS-DOS, IBM DOS, FreeDOS executable
    DosExe,

    /// .Net executable, PE32
    DotNet,

    /// Portable Executable, MS Windows
    PE32,

    /// 64-bit Portable Executable, MS Windows
    PE32_64,

    /// New Executable, 16-bit, Windows 3.1
    NE,

    /// Linear Executable, IBM OS/2
    LE,
}

impl Display for SubType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            SubType::DosExe => write!(f, "MZ DOS executable"),
            SubType::DotNet => write!(f, ".Net executable"),
            SubType::PE32 => write!(f, "Portable Executable (PE32)"),
            SubType::PE32_64 => write!(f, "64-bit Portable Executable (PE32+)"),
            SubType::NE => write!(f, "New Executable"),
            SubType::LE => write!(f, "Linear Executable"),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum SubSystem {
    /// Unknown subsystem
    Unknown,

    /// Graphical Windows
    WindowsGUI,

    /// Command-line Windows
    WindowsCLI,

    /// Command-line OS/2
    OS2CLI,

    /// Posix subsystem
    POSIX,

    /// Native code
    NativeCode,

    /// Native driver
    NativeDriver,

    /// Windows CE (embedded)
    WindowsCE,

    /// EFI (Extensible Firmware Interface) application
    EFI,

    /// EFI driver with boot services
    EFIBoot,

    /// EFI ROM image
    EFIRom,

    /// EFI driver with runtime services
    EFIRuntimeDriver,

    /// Xbox application
    Xbox,

    /// Windows Boot application
    WindowsBoot,
}

impl Display for SubSystem {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            SubSystem::Unknown => write!(f, "Unknown subsystem"),
            SubSystem::WindowsGUI => write!(f, "Windows GUI"),
            SubSystem::WindowsCLI => write!(f, "Windows CLI"),
            SubSystem::OS2CLI => write!(f, "OS/2 CLI"),
            SubSystem::POSIX => write!(f, "Posix"),
            SubSystem::NativeCode => write!(f, "Native Code"),
            SubSystem::NativeDriver => write!(f, "Native Driver"),
            SubSystem::WindowsCE => write!(f, "Windows CE"),
            SubSystem::EFI => write!(f, "EFI"),
            SubSystem::EFIBoot => write!(f, "EFI boot"),
            SubSystem::EFIRom => write!(f, "EFI ROM"),
            SubSystem::EFIRuntimeDriver => write!(f, "EFI driver with runtime"),
            SubSystem::Xbox => write!(f, "Xbox"),
            SubSystem::WindowsBoot => write!(f, "Windows Boot"),
        }
    }
}

impl From<u16> for SubSystem {
    fn from(value: u16) -> Self {
        match value {
            1 => SubSystem::NativeCode,
            2 => SubSystem::WindowsGUI,
            3 => SubSystem::WindowsCLI,
            5 => SubSystem::OS2CLI,
            7 => SubSystem::POSIX,
            8 => SubSystem::NativeDriver,
            9 => SubSystem::WindowsCE,
            10 => SubSystem::EFI,
            11 => SubSystem::EFIBoot,
            12 => SubSystem::EFIRuntimeDriver,
            13 => SubSystem::EFIRom,
            14 => SubSystem::Xbox,
            16 => SubSystem::WindowsBoot,
            _ => SubSystem::Unknown,
        }
    }
}

#[derive(Clone, Debug)]
pub struct EXE<'a> {
    /// If the program is 64-bit
    pub is64bit: bool,

    /// If the binary has extra data after the last section, could be used to hide something
    pub has_overlay: Option<bool>,

    /// Instruction set architecture for this binary
    pub arch: Architecture,

    /// EXE sub-type, mostly if it's for DOS, Windows, OS/2
    pub sub_type: SubType,

    /// Operating System for this binary, likely Windows
    pub os: OperatingSystem,

    /// COFF (Common Object File Format) header of the program
    pub coff_header: Option<fields::COFFHeader>,

    /// Optional Header for this program, not optional if for Windows
    pub optional_header: Option<fields::OptionalHeader>,

    /// Executable subtype: Program or Library?
    pub exec_type: ExecutableType,

    /// Windows Subsystem used by this program
    pub subsystem: Option<SubSystem>,

    /// Sections of this binary
    pub sections: Option<Sections<'a>>,

    /// External libraries used by this application or library
    pub imports: Option<Imports>,

    /// The array containing the raw bytes used to parse this program
    pub contents: &'a [u8],
}

impl<'a> EXE<'a> {
    pub fn from(contents: &'a [u8]) -> Result<Self> {
        if !((contents[0] == EXE_MAGIC[0] && contents[1] == EXE_MAGIC[1])
            || (contents[0] == EXE_MAGIC[1] && contents[1] == EXE_MAGIC[0]))
        {
            bail!("Not a MZ, MS-DOS, or PE32 file");
        }

        let pe_magic_offset = u32_from_offset(contents, 0x3C, Ordering::LittleEndian) as usize;
        let coff_header_offset = pe_magic_offset + 4;
        if pe_magic_offset > contents.len()
            || pe_magic_offset + PE_MAGIC.len() > contents.len()
            || !bytes_offset_match(contents, pe_magic_offset, &PE_MAGIC)
        {
            return Ok(Self {
                is64bit: false,
                has_overlay: None,
                arch: Architecture::X86,
                sub_type: SubType::DosExe,
                os: OperatingSystem::DOS,
                subsystem: None,
                sections: None,
                coff_header: None,
                optional_header: None,
                exec_type: ExecutableType::Program,
                imports: None,
                contents,
            });
        }

        let mut sections = Sections::default();
        sections.0.push(Section {
            name: "DOS Header".into(),
            is_executable: false,
            size: 0x3F,
            offset: 0,
            virtual_address: 0,
            virtual_size: 0,
            entropy: contents[..0x3F].to_vec().entropy(),
            data: Some(&contents[..0x3F]),
        });
        sections.0.push(Section {
            name: "DOS Stub".into(),
            is_executable: true,
            size: (pe_magic_offset - 0x3F),
            offset: 0x3F,
            virtual_address: 0,
            virtual_size: 0,
            entropy: contents[0x3F..pe_magic_offset].to_vec().entropy(),
            data: Some(&contents[0x3F..pe_magic_offset]),
        });
        sections.0.push(Section {
            name: "COFF Header".into(),
            is_executable: false,
            size: 20,
            offset: coff_header_offset,
            virtual_address: 0,
            virtual_size: 0,
            entropy: contents[coff_header_offset..coff_header_offset + 20]
                .to_vec()
                .entropy(),
            data: Some(&contents[coff_header_offset..coff_header_offset + 20]),
        });

        let coff = fields::COFFHeader::from(&contents[coff_header_offset..]);
        let optional_header = fields::OptionalHeader::from(&contents[coff_header_offset + 20..]);

        let is64bit = optional_header.is_64bit();
        let sub_type = {
            let mut sub_type = SubType::PE32;
            if is64bit {
                sub_type = SubType::PE32_64;
            }

            if let Some(data_directories) = optional_header.data_directories() {
                if data_directories.clr_runtime_header.virtual_address > 0 {
                    sub_type = SubType::DotNet;
                }
            }

            sub_type
        };

        let exec_type = {
            if coff
                .characteristics
                .contains(fields::CoffCharacteristics::FileDLL)
            {
                ExecutableType::Library
            } else if coff
                .characteristics
                .contains(fields::CoffCharacteristics::ExecutableImage)
            {
                ExecutableType::Program
            } else {
                ExecutableType::Unknown(0)
            }
        };

        let mut subsystem = None;
        let mut has_overlay = false;
        if coff
            .characteristics
            .contains(fields::CoffCharacteristics::ExecutableImage)
        {
            let mut offset = u32_from_offset(contents, 0x3C, Ordering::LittleEndian) + 24;
            subsystem = Some(optional_header.subsystem().unwrap().into());
            sections.0.push(Section {
                name: "Optional Header".into(),
                is_executable: false,
                size: 240,
                offset: offset as usize,
                virtual_address: 0,
                virtual_size: 0,
                entropy: contents[offset as usize..offset as usize + 240]
                    .to_vec()
                    .entropy(),
                data: Some(&contents[offset as usize..offset as usize + 240]),
            });
            if is64bit {
                offset += 240;
            } else {
                offset += 224;
            }
            let mut sections_offset = 0;
            for sec_num in 0..coff.num_sections {
                let section_name =
                    String::from_utf8(contents[offset as usize..(offset + 8) as usize].to_vec())?;
                let section_virtual_size =
                    u32_from_offset(contents, (offset + 8) as usize, Ordering::LittleEndian);
                let section_virtual_address =
                    u32_from_offset(contents, (offset + 12) as usize, Ordering::LittleEndian);
                let section_size =
                    u32_from_offset(contents, (offset + 16) as usize, Ordering::LittleEndian);
                let section_offset =
                    u32_from_offset(contents, (offset + 20) as usize, Ordering::LittleEndian);
                let section_characteristics = FlagSet::<fields::SectionFlags>::new(
                    u32_from_offset(contents, (offset + 36) as usize, Ordering::LittleEndian),
                )
                .unwrap();
                let section_executable = section_characteristics.contains(
                    fields::SectionFlags::Executable | fields::SectionFlags::ContainsCode,
                );

                if section_offset + section_size < contents.len() as u32 {
                    sections.0.push(Section {
                        name: section_name,
                        is_executable: section_executable,
                        size: section_size as usize,
                        offset: section_offset as usize,
                        virtual_address: section_virtual_address,
                        virtual_size: section_virtual_size,
                        entropy: contents
                            [section_offset as usize..(section_offset + section_size) as usize]
                            .to_vec()
                            .entropy(),
                        data: Some(
                            &contents
                                [section_offset as usize..(section_offset + section_size) as usize],
                        ),
                    });
                } else {
                    debug!("Skipping section {sec_num} '{section_name}': offset {section_offset} + size {section_size} >= file length {}", contents.len());
                }

                if section_size + section_offset > sections_offset {
                    sections_offset = section_size + section_offset;
                }
                offset += 40; // Section entries are 40 bytes
            }
            has_overlay =
                sections_offset as usize > 0 && contents.len() - sections_offset as usize > 0;
        }

        let mut imports = None;
        if let Some(data_directories) = optional_header.data_directories() {
            let import_section = sections
                .0
                .iter()
                .find(|x| x.virtual_address == data_directories.import_table.virtual_address);

            if let Some(import_section) = import_section {
                if import_section.size > 8 && import_section.offset > 0 {
                    let num_imports = match optional_header {
                        fields::OptionalHeader::OptionalHeader32(h32) => h32.num_rva_and_sizes,
                        fields::OptionalHeader::OptionalHeader64(h64) => h64.num_rva_and_sizes,
                    }
                    .unwrap();
                    imports = Some(data_directories.imports(num_imports as usize, import_section));
                }
            }
        }

        Ok(Self {
            is64bit,
            has_overlay: Some(has_overlay),
            arch: coff.architecture(),
            sub_type,
            os: OperatingSystem::Windows,
            subsystem,
            sections: Some(sections),
            exec_type,
            coff_header: Some(*coff),
            optional_header: Some(optional_header),
            imports,
            contents,
        })
    }
}

impl<'a> ExecutableFile for EXE<'a> {
    fn type_name(&self) -> String {
        self.sub_type.to_string()
    }

    fn architecture(&self) -> &Architecture {
        &self.arch
    }

    fn pointer_size(&self) -> usize {
        if self.sub_type == SubType::PE32_64 {
            64
        } else {
            32
        }
    }

    fn operating_system(&self) -> &OperatingSystem {
        &self.os
    }
    fn compiled_timestamp(&self) -> Option<DateTime<Utc>> {
        self.coff_header.map(|c| c.compiled_date())
    }

    fn num_sections(&self) -> u32 {
        if let Some(sections) = &self.sections {
            sections.0.len() as u32
        } else {
            0
        }
    }

    fn sections(&self) -> Result<&Sections> {
        self.sections
            .as_deref()
            .ok_or_else(|| anyhow::Error::msg("No sections"))
            .context("")
    }
}

impl<'a> TypeMagic for EXE<'a> {
    const MAGIC: &'static [&'static [u8]] = &[&[b'M', b'Z'], &[b'Z', b'M']];
}

impl<'a> Display for EXE<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "EXE file:")?;
        writeln!(f, "\tSubtype: {}", self.sub_type)?;
        writeln!(f, "\tExecutable type: {}", self.exec_type)?;
        writeln!(f, "\tOS: {}", self.os)?;
        if let Some(coff) = self.coff_header {
            writeln!(
                f,
                "\tCompiled: {:?}",
                coff.compiled_date().format("%Y-%m-%d %H:%M:%S").to_string()
            )?;
        }
        if let Some(subsys) = &self.subsystem {
            writeln!(f, "\tSubsystem: {subsys}")?;
        }
        if let Some(sections) = &self.sections {
            writeln!(f, "\t{} sections:", sections.0.len())?;
            for section in sections.0.iter() {
                writeln!(f, "\t\t{section}")?;
            }
        }
        if let Some(imports) = &self.imports {
            if imports.imports.len() == imports.expected_imports as usize {
                writeln!(f, "\t{} imports:", imports.imports.len())?;
            } else {
                writeln!(
                    f,
                    "\t{} out of {} imports:",
                    imports.imports.len(),
                    imports.expected_imports
                )?;
            }
            for import in imports.imports.iter() {
                writeln!(f, "\t\t{import}")?;
            }
            writeln!(f, "\t\tImport hash: {}", hex::encode(imports.hash()))?;
            writeln!(f, "\t\tFuzzy Import hash: {}", imports.fuzzy_hash())?;
        }
        if self.has_overlay == Some(true) {
            writeln!(f, "\tHas extra bytes at the end (overlay).")?;
        }
        writeln!(f, "\tSize: {}", self.contents.len())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn arm64_gnu() {
        const BYTES: &[u8] =
            include_bytes!("../../../testdata/exe/pe64_win32_console_arm64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        // To-do: add check for imports
        // Currently not found, but are present
        assert_eq!(exe.arch, Architecture::ARM64);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::WindowsCLI));
    }

    #[test]
    fn x86_64_gnu_console() {
        const BYTES: &[u8] =
            include_bytes!("../../../testdata/exe/pe64_win32_console_x86_64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert_eq!(
            "610b6b6aa37c8e01c9855a05dcf00565",
            hex::encode(exe.imports.unwrap().hash())
        );
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::WindowsCLI));
    }

    #[test]
    fn x86_64_gnu_gui() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/exe/pe64_win32_gui_x86_64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert_eq!(
            "610b6b6aa37c8e01c9855a05dcf00565",
            hex::encode(exe.imports.unwrap().hash())
        );
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::WindowsGUI));
    }

    #[test]
    fn x86_64_gnu_posix() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/exe/pe64_posix_x86_64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert_eq!(
            "610b6b6aa37c8e01c9855a05dcf00565",
            hex::encode(exe.imports.unwrap().hash())
        );
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::POSIX));
    }

    #[test]
    fn x86_64_gnu_xbox() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/exe/pe64_xbox_x86_64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert_eq!(
            "610b6b6aa37c8e01c9855a05dcf00565",
            hex::encode(exe.imports.unwrap().hash())
        );
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::Xbox));
    }

    #[test]
    fn x86_64_gnu_lib_console() {
        const BYTES: &[u8] =
            include_bytes!("../../../testdata/exe/pe64_lib_console_x86_64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert_eq!(
            "466e0075eba65b201b4cc8d4d3f85cbb",
            hex::encode(exe.imports.unwrap().hash())
        );
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Library);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::WindowsCLI));
    }

    #[test]
    fn x86_64_gnu_lib_gui() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/exe/pe64_lib_gui_x86_64_gnu.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert_eq!(
            "466e0075eba65b201b4cc8d4d3f85cbb",
            hex::encode(exe.imports.unwrap().hash())
        );
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Library);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::WindowsGUI));
    }

    #[test]
    fn dot_net_mono() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/exe/pe32_dotnet.exe");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        // To-do: add check for imports
        // Currently not found, but are present
        assert_eq!(exe.arch, Architecture::X86);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::DotNet);
        assert_eq!(exe.subsystem, Some(SubSystem::WindowsCLI));
    }

    #[test]
    fn efi() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/exe/efi/hello.efi");

        let exe = EXE::from(BYTES).unwrap();
        eprintln!("{exe}");
        assert!(exe.imports.is_none());
        assert_eq!(exe.arch, Architecture::X86_64);
        assert_eq!(exe.exec_type, ExecutableType::Program);
        assert_eq!(exe.sub_type, SubType::PE32_64);
        assert_eq!(exe.subsystem, Some(SubSystem::EFI));
    }
}
