pub mod fat;

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

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

use anyhow::{bail, Context, Result};
use chrono::{DateTime, Utc};

// Apple's Mach-O loader:
// https://opensource.apple.com/source/xnu/xnu-2050.18.24/EXTERNAL_HEADERS/mach-o/loader.h

const MAGIC32: [u8; 4] = [0xfe, 0xed, 0xfa, 0xce];
const CIGAM32: [u8; 4] = [0xce, 0xfa, 0xed, 0xfe];
const MAGIC64: [u8; 4] = [0xfe, 0xed, 0xfa, 0xcf];
const CIGAM64: [u8; 4] = [0xcf, 0xfa, 0xed, 0xfe];

/// A struct presenting Mach Objects (Mach-O) files
///
/// These are executables or libraries used on macOS, iOS, iPadOS, tvOS, watchOS, etc. They began
/// as the file format for NeXTSTEP.
///
/// Because of the different architectures, these files could be bi-endian.
#[derive(Clone, Debug)]
pub struct Macho<'a> {
    /// If the program is 64-bit
    pub is64bit: bool,

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

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

    /// Byte ordering for this binary
    pub ordering: Ordering,

    /// Executable subtype: Program, Library, or Core file?
    pub executable_type: ExecutableType,

    /// Operating System for this binary, going to be Mac OS or some derivative, could be NeXTSTEP
    pub os: OperatingSystem,

    /// 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> Macho<'a> {
    pub fn from(contents: &'a [u8]) -> Result<Self> {
        let (is_64bit, ordering) = {
            if bytes_offset_match(contents, 0, &MAGIC32) {
                (false, Ordering::BigEndian)
            } else if bytes_offset_match(contents, 0, &CIGAM32) {
                (false, Ordering::LittleEndian)
            } else if bytes_offset_match(contents, 0, &MAGIC64) {
                (true, Ordering::BigEndian)
            } else if bytes_offset_match(contents, 0, &CIGAM64) {
                (true, Ordering::LittleEndian)
            } else {
                bail!("Not a Mach-O file");
            }
        };

        // The CPU type has an upper end flag to indicate if 64-bit, in addition
        // to the different machine number. Check it, and clear it.
        let mut arch = u32_from_offset(contents, 4, ordering);
        let _should_be_64bit = (arch & 0x01000000) == 0x01000000;
        arch &= 0x00FFFFFF;
        let arch = match arch {
            0x06 => Architecture::M68k,
            0x07 => {
                if is_64bit {
                    Architecture::X86_64
                } else {
                    Architecture::X86
                }
            }
            0x8 => {
                if is_64bit {
                    Architecture::MIPS64
                } else {
                    Architecture::MIPS
                }
            }
            0x0C => {
                if is_64bit {
                    Architecture::ARM64
                } else {
                    Architecture::ARM
                }
            }
            0x0D => Architecture::M88k,
            0x0E => {
                if is_64bit {
                    Architecture::Sparc64
                } else {
                    Architecture::Sparc
                }
            }
            0x10 => Architecture::Alpha,
            0x12 => {
                if is_64bit {
                    Architecture::PowerPC64
                } else {
                    Architecture::PowerPC
                }
            }
            other => Architecture::Other(other as u16),
        };

        let exec_type = u16_from_offset(contents, 4, ordering);
        let exec_type = {
            if (exec_type & 0x02) != 0 || (exec_type & 0x05) != 0 {
                ExecutableType::Program
            } else if (exec_type & 0x06) != 0 || (exec_type & 0x09) != 0 {
                ExecutableType::Library
            } else if (exec_type & 0x04) != 0 {
                ExecutableType::Core
            } else {
                ExecutableType::Unknown(exec_type)
            }
        };

        Ok(Self {
            is64bit: is_64bit,
            arch,
            has_overlay: None,
            ordering,
            executable_type: exec_type,
            os: OperatingSystem::MacOS,
            sections: None,
            imports: None,
            contents,
        })
    }
}

impl<'a> ExecutableFile for Macho<'a> {
    fn type_name(&self) -> String {
        "Mach-O".into()
    }

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

    fn pointer_size(&self) -> usize {
        if self.is64bit {
            64
        } else {
            32
        }
    }

    fn operating_system(&self) -> &OperatingSystem {
        &self.os
    }
    fn compiled_timestamp(&self) -> Option<DateTime<Utc>> {
        None
    }

    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 Macho<'a> {
    const MAGIC: &'static [&'static [u8]] = &[&MAGIC32, &MAGIC64, &CIGAM32, &CIGAM64];
}

impl<'a> Display for Macho<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Mach-O file:")?;
        writeln!(f, "\tOS: {}", self.os)?;
        writeln!(f, "\tArchitecture: {}", self.arch)?;
        writeln!(f, "\tOrdering: {}", self.ordering)?;
        writeln!(f, "\tType: {}", self.executable_type)?;
        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 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() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/macho/macho_arm64");

        let macho = Macho::from(BYTES).unwrap();
        assert!(macho.is64bit);
        assert_eq!(macho.executable_type, ExecutableType::Program);
        assert_eq!(macho.arch, Architecture::ARM64);
    }

    #[test]
    fn ppc() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/macho/macho_ppc");

        let macho = Macho::from(BYTES).unwrap();
        assert!(!macho.is64bit);
        //assert_eq!(macho.executable_type, ExecutableType::Program);
        assert_eq!(macho.arch, Architecture::PowerPC);
    }

    #[test]
    fn ppc64() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/macho/macho_ppc64");

        let macho = Macho::from(BYTES).unwrap();
        assert!(macho.is64bit);
        //assert_eq!(macho.executable_type, ExecutableType::Program);
        assert_eq!(macho.arch, Architecture::PowerPC64);
    }

    #[test]
    fn x86_64() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/macho/macho_x86_64");

        let macho = Macho::from(BYTES).unwrap();
        assert!(macho.is64bit);
        assert_eq!(macho.executable_type, ExecutableType::Program);
        assert_eq!(macho.arch, Architecture::X86_64);
    }
}
