pub mod fields;

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

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

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

// Wikipedia article:
// https://en.wikipedia.org/wiki/Executable_and_Linkable_Format

const MAGIC: [u8; 4] = [0x7f, 0x45, 0x4c, 0x46];

/// The struct to partially represent the ELF (Executable and Linkable File) format
///
/// This is the file type used for Linux, *BSD (FreeBSD, OpenBSD, NetBSD, etc), Haiku, Solaris,
/// and possibly some others.
///
/// Effort is made to fail gracefully, since malware may not obey all the rules, and some information
/// is better than none because some part of the data wasn't parsed correctly.
#[derive(Clone, Debug)]
pub struct Elf<'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
    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 path for the ELF loader (or interpreter)
    pub interpreter: Option<String>,

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

impl<'a> Elf<'a> {
    pub fn from(contents: &'a [u8]) -> Result<Self> {
        if !bytes_offset_match(contents, 0, &MAGIC) {
            bail!("Not an ELF file");
        }

        let is_64bit = contents[4] == 2;
        let ordering = {
            if contents[5] == 2 {
                Ordering::BigEndian
            } else {
                Ordering::LittleEndian
            }
        };
        let mut os = match contents[7] {
            1 => OperatingSystem::HPUX,
            2 => OperatingSystem::NetBSD,

            // Assume SystemV is Linux, some compilers don't set this correctly
            // https://unix.stackexchange.com/questions/132036/why-does-readelf-show-system-v-as-my-os-instead-of-linux
            0 | 3 => OperatingSystem::Linux,
            6 => OperatingSystem::Solaris,
            7 => OperatingSystem::AIX,
            8 => OperatingSystem::Irix,
            9 => OperatingSystem::FreeBSD,
            0xC => OperatingSystem::OpenBSD,
            other => OperatingSystem::Other(other as u16),
        };

        let elf_type = match u16_from_offset(contents, 0x10, ordering) {
            1 | 2 => ExecutableType::Program,
            3 => ExecutableType::Library,
            4 => ExecutableType::Core,
            other => ExecutableType::Unknown(other),
        };

        let arch = match u16_from_offset(contents, 0x12, ordering) {
            // https://refspecs.linuxfoundation.org/elf/gabi4+/ch4.eheader.html
            0 => Architecture::Unknown,
            2 => Architecture::Sparc,
            3 => Architecture::X86,
            4 => Architecture::M68k,
            5 => Architecture::M88k,
            8 => {
                if is_64bit {
                    Architecture::MIPS64
                } else {
                    Architecture::MIPS
                }
            }
            0x0A => {
                if is_64bit {
                    Architecture::MIPSEL64
                } else {
                    Architecture::MIPSEL
                }
            }
            0x14 => {
                if ordering == Ordering::BigEndian {
                    Architecture::PowerPC
                } else {
                    Architecture::PowerPCLE
                }
            }
            0x15 => {
                if ordering == Ordering::BigEndian {
                    Architecture::PowerPC64
                } else {
                    Architecture::PowerPC64LE
                }
            }
            0x16 => {
                if is_64bit {
                    Architecture::S390x
                } else {
                    Architecture::S390
                }
            }

            0x28 => Architecture::ARM,
            0x29 => Architecture::Alpha,
            0x2b => Architecture::Sparc64,

            0x32 => Architecture::Itanium,
            0x3E => Architecture::X86_64,

            0xB7 => Architecture::ARM64,

            0xF3 => {
                if is_64bit {
                    Architecture::RISCV64
                } else {
                    Architecture::RISCV
                }
            }

            0x39d => Architecture::Hobbit,

            other => Architecture::Other(other),
        };

        // Offset for section table
        let e_shoff = {
            if is_64bit {
                u64_from_offset(contents, 0x28, ordering) as usize
            } else {
                u32_from_offset(contents, 0x20, ordering) as usize
            }
        };

        let e_phentsize = {
            if is_64bit {
                u16_from_offset(contents, 0x36, ordering) as usize
            } else {
                u16_from_offset(contents, 0x2A, ordering) as usize
            }
        };

        let e_phnum = {
            if is_64bit {
                u16_from_offset(contents, 0x38, ordering) as usize
            } else {
                u16_from_offset(contents, 0x2C, ordering) as usize
            }
        };

        let mut interpreter = None;
        for p_header_index in 0..e_phnum {
            let start_index = {
                if is_64bit {
                    0x40 + p_header_index * e_phentsize
                } else {
                    0x34 + p_header_index * e_phentsize
                }
            };
            let p_type = FlagSet::<fields::ProgramHeaderFlags>::new_truncated(u32_from_offset(
                contents,
                start_index,
                ordering,
            ));

            if p_type.contains(fields::ProgramHeaderFlags::Interpreter) {
                let header = &contents[start_index..start_index + e_phentsize];
                let p_offset = {
                    if is_64bit {
                        u64_from_offset(header, 0x08, ordering) as usize
                    } else {
                        u32_from_offset(header, 0x04, ordering) as usize
                    }
                };
                let p_filesz = {
                    if is_64bit {
                        u64_from_offset(header, 0x20, ordering) as usize
                    } else {
                        u32_from_offset(header, 0x10, ordering) as usize
                    }
                };

                if p_offset > 0 && p_filesz > 0 {
                    let interpreter_path =
                        String::from_utf8(Vec::from(&contents[p_offset..p_offset + p_filesz]))
                            .unwrap();
                    if interpreter_path.contains("/system/runtime_loader") {
                        os = OperatingSystem::Haiku;
                    }
                    interpreter = Some(interpreter_path);
                    break;
                }
            }
        }

        // Size of a section header table entry
        let e_shentsize = {
            if is_64bit {
                u16_from_offset(contents, 0x3A, ordering)
            } else {
                u16_from_offset(contents, 0x2E, ordering)
            }
        };

        // Number of entries in the section header table
        let e_shnum = {
            if is_64bit {
                u16_from_offset(contents, 0x3C, ordering)
            } else {
                u16_from_offset(contents, 0x30, ordering)
            }
        };

        // Index of the section header table with section names
        let e_shstrndx = {
            if is_64bit {
                u16_from_offset(contents, 0x3E, ordering)
            } else {
                u16_from_offset(contents, 0x32, ordering)
            }
        };

        // Find the offset for section name strings
        let section_names_offset = {
            if is_64bit {
                u64_from_offset(
                    contents,
                    e_shoff + (e_shstrndx * e_shentsize) as usize + 0x18,
                    ordering,
                ) as usize
            } else {
                u32_from_offset(
                    contents,
                    e_shoff + (e_shstrndx * e_shentsize) as usize + 0x10,
                    ordering,
                ) as usize
            }
        };

        let mut section_offset = e_shoff;
        let mut sections = Sections::default();
        for _section_index in 0..e_shnum {
            let section_name_offset =
                u32_from_offset(contents, section_offset, ordering) as usize + section_names_offset;
            let section_name = string_from_offset(contents, section_name_offset);

            let section_type = FlagSet::<fields::SectionHeaderTypes>::new_truncated(
                u32_from_offset(contents, section_offset + 0x4, ordering),
            );

            if section_type.contains(fields::SectionHeaderTypes::DynamicSymbolsTable) {
                // Parse imports here
            }

            let this_section_offset = {
                if is_64bit {
                    u64_from_offset(contents, section_offset + 0x18, ordering) as usize
                } else {
                    u32_from_offset(contents, section_offset + 0x10, ordering) as usize
                }
            };
            let section_size = {
                if is_64bit {
                    u64_from_offset(contents, section_offset + 0x20, ordering) as usize
                } else {
                    u32_from_offset(contents, section_offset + 0x14, ordering) as usize
                }
            };
            let section_flags = {
                if is_64bit {
                    u64_from_offset(contents, section_offset + 0x08, ordering) as usize
                } else {
                    u32_from_offset(contents, section_offset + 0x08, ordering) as usize
                }
            };

            sections.0.push(Section {
                name: section_name,
                is_executable: (section_flags & 4) != 0,
                size: section_size,
                offset: this_section_offset,
                virtual_address: 0,
                virtual_size: 0,
                entropy: contents[this_section_offset..this_section_offset + section_size]
                    .to_vec()
                    .entropy(),
                data: Some(&contents[this_section_offset..this_section_offset + section_size]),
            });

            // Advance to the next section
            section_offset += e_shentsize as usize;
        }

        Ok(Self {
            is64bit: is_64bit,
            arch,
            has_overlay: Some(section_offset < contents.len()),
            ordering,
            executable_type: elf_type,
            os,
            sections: Some(sections),
            imports: None,
            interpreter,
            contents,
        })
    }
}

impl<'a> ExecutableFile for Elf<'a> {
    fn type_name(&self) -> String {
        "ELF".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 Elf<'a> {
    const MAGIC: &'static [&'static [u8]] = &[&MAGIC];
}

impl<'a> Display for Elf<'a> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "ELF 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(interp) = &self.interpreter {
            writeln!(f, "\tInterpreter: {interp}")?;
        }
        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::*;

    // assert_eq!(elf.executable_type, ExecutableType::Program); should pass for all
    // But for some targets, the file is called a Shared Object

    #[test]
    fn arm32() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_arm");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(!elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.executable_type, ExecutableType::Program);
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::ARM);
    }

    #[test]
    fn arm64() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_arm64");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::ARM64);
    }

    #[test]
    fn hobbit() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_aclock_hobbit_beos");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(!elf.is64bit);
        assert_eq!(elf.ordering, Ordering::BigEndian);
        assert_eq!(elf.arch, Architecture::Hobbit);
    }

    #[test]
    fn mips32() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_mips");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(!elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.executable_type, ExecutableType::Program);
        assert_eq!(elf.ordering, Ordering::BigEndian);
        assert_eq!(elf.arch, Architecture::MIPS);
    }

    #[test]
    fn mips64() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_mips64");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.executable_type, ExecutableType::Program);
        assert_eq!(elf.ordering, Ordering::BigEndian);
        assert_eq!(elf.arch, Architecture::MIPS64);
    }

    #[test]
    fn ppc64le() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_ppc64le");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.executable_type, ExecutableType::Program);
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::PowerPC64LE);
    }

    #[test]
    fn ppc64le_lib() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_ppc64le.so");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_none());
        assert_eq!(elf.executable_type, ExecutableType::Library);
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::PowerPC64LE);
    }

    #[test]
    fn s390x() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_s390x");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.ordering, Ordering::BigEndian);
        assert_eq!(elf.arch, Architecture::S390x);
    }

    #[test]
    fn riscv() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_linux_riscv64");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Linux);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::RISCV64);
    }

    #[test]
    fn x86() {
        const BYTES: &[u8] = include_bytes!("../../../testdata/elf/elf_haiku_x86");

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(!elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::Haiku);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::X86);
    }

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

        let elf = Elf::from(BYTES).unwrap();
        eprintln!("{elf}");
        assert!(elf.is64bit);
        assert_eq!(elf.os, OperatingSystem::FreeBSD);
        assert!(elf.interpreter.is_some());
        assert_eq!(elf.executable_type, ExecutableType::Program);
        assert_eq!(elf.ordering, Ordering::LittleEndian);
        assert_eq!(elf.arch, Architecture::X86_64);
    }
}
