// SPDX-License-Identifier: Apache-2.0

/// Flags for some fields within an ELF header
pub mod fields;

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

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

use anyhow::{Result, anyhow, bail};
use chrono::{DateTime, Utc};
use flagset::FlagSet;
use tracing::{error, info, instrument};
use uuid::Uuid;

// 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> {
    /// ELF parsed from a sequence of bytes
    ///
    /// # Errors
    ///
    /// Returns an error if the parsing fails.
    #[allow(
        clippy::too_many_lines,
        clippy::cast_possible_truncation,
        clippy::similar_names
    )] // TODO: Revisit these
    #[instrument(name = "ELF parser", skip(contents))]
    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(u16::from(other)),
        };

        let elf_type = match u16_from_offset(contents, 0x10, ordering)
            .ok_or(anyhow!("ELF buffer too small for elf type"))?
        {
            1 | 2 => ExecutableType::Program,
            3 => ExecutableType::Library,
            4 => ExecutableType::Core,
            other => ExecutableType::Unknown(other),
        };

        let arch = match u16_from_offset(contents, 0x12, ordering)
            .ok_or(anyhow!("ELF buffer too small for architecture"))?
        {
            // 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 => {
                if is_64bit {
                    Architecture::Alpha64
                } else {
                    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(u32::from(other)),
        };

        // Offset for section table
        let e_shoff = {
            if is_64bit {
                u64_from_offset(contents, 0x28, ordering)
                    .ok_or(anyhow!("ELF contents too short for section offset"))?
                    as usize
            } else {
                u32_from_offset(contents, 0x20, ordering)
                    .ok_or(anyhow!("ELF contents too short for section offset"))?
                    as usize
            }
        };

        let e_phentsize = {
            if is_64bit {
                u16_from_offset(contents, 0x36, ordering)
            } else {
                u16_from_offset(contents, 0x2A, ordering)
            }
        }
        .ok_or(anyhow!("ELF contents too short program entry size"))?
            as usize;

        let e_phnum = {
            if is_64bit {
                u16_from_offset(contents, 0x38, ordering)
            } else {
                u16_from_offset(contents, 0x2C, ordering)
            }
        }
        .ok_or(anyhow!("ELF contents too short program entries"))? 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).unwrap_or_default(),
            );

            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).unwrap_or_default() as usize
                    } else {
                        u32_from_offset(header, 0x04, ordering).unwrap_or_default() as usize
                    }
                };
                let p_filesz = {
                    if is_64bit {
                        u64_from_offset(header, 0x20, ordering).unwrap_or_default() as usize
                    } else {
                        u32_from_offset(header, 0x10, ordering).unwrap_or_default() as usize
                    }
                };

                if p_offset > 0 && p_filesz > 0 {
                    let interpreter_path =
                        String::from_utf8(Vec::from(&contents[p_offset..p_offset + p_filesz]))
                            .map_err(|e| {
                                error!(
                                    "Interpreter error {e}, bytes: {:?}",
                                    &contents[p_offset..p_offset + p_filesz]
                                );
                            })
                            .unwrap_or_default();
                    if !interpreter_path.is_empty() {
                        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)
            }
        }
        .ok_or(anyhow!("ELF contents too short for section entry size"))?;

        // 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)
            }
        }
        .ok_or(anyhow!("ELF contents too short for section count"))?;

        // 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)
            }
        }
        .ok_or(anyhow!("ELF contents too short for section header table with section names"))?;

        // 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,
                )
                .ok_or(anyhow!("ELF contents too short for section name"))? as usize
            } else {
                u32_from_offset(
                    contents,
                    e_shoff + (e_shstrndx * e_shentsize) as usize + 0x10,
                    ordering,
                )
                .ok_or(anyhow!("ELF contents too short for section name"))? 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)
                .unwrap_or_default() as usize
                + section_names_offset;
            let section_name = if section_name_offset < contents.len() {
                string_from_offset(contents, section_name_offset).unwrap_or_default()
            } else {
                info!(
                    "ELF: section name offset {section_name_offset} greater than buffer length {}.",
                    contents.len()
                );
                String::new()
            };

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

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

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

            if this_section_offset + section_size <= contents.len() {
                let section_bytes =
                    &contents[this_section_offset..this_section_offset + section_size];

                sections.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: section_bytes.to_vec().entropy(),
                    data: Some(section_bytes),
                });
            } else {
                error!(
                    "Section {section_index}: {section_name} offset {this_section_offset} + size {section_size} (end {}) is beyond the ELF buffer {}!",
                    this_section_offset + section_size,
                    contents.len()
                );
            }

            // 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 ExecutableFile for Elf<'_> {
    fn architecture(&self) -> Option<Architecture> {
        Some(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
    }

    #[allow(clippy::cast_possible_truncation)]
    fn num_sections(&self) -> u32 {
        self.sections.as_ref().unwrap_or(&Sections::default()).len() as u32
    }

    fn sections(&self) -> Option<&Sections<'_>> {
        self.sections.as_ref()
    }

    fn import_hash(&self) -> Option<Uuid> {
        self.imports.as_ref().map(Imports::hash)
    }

    fn fuzzy_imports(&self) -> Option<String> {
        self.imports.as_ref().map(Imports::fuzzy_hash)
    }
}

impl SpecimenFile for Elf<'_> {
    const MAGIC: &'static [&'static [u8]] = &[&MAGIC];

    fn type_name(&self) -> &'static str {
        "ELF"
    }
}

impl Display for Elf<'_> {
    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.len())?;
            for section in sections {
                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())?;
        writeln!(f, "\tEntropy: {:.4}", self.contents.entropy())
    }
}

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

    use rstest::rstest;

    #[rstest]
    #[case::arm32(include_bytes!("../../../testdata/elf/elf_linux_arm"), false, OperatingSystem::Linux, Architecture::ARM, Ordering::LittleEndian, ExecutableType::Program)]
    #[case::arm64(include_bytes!("../../../testdata/elf/elf_linux_arm64"), true, OperatingSystem::Linux, Architecture::ARM64, Ordering::LittleEndian, ExecutableType::Library /* Not really a library, but compiler said it is*/)]
    #[case::mips32(include_bytes!("../../../testdata/elf/elf_linux_mips"), false, OperatingSystem::Linux, Architecture::MIPS, Ordering::BigEndian, ExecutableType::Program)]
    #[case::mips64(include_bytes!("../../../testdata/elf/elf_linux_mips64"), true, OperatingSystem::Linux, Architecture::MIPS64, Ordering::BigEndian, ExecutableType::Program)]
    #[case::ppc64le(include_bytes!("../../../testdata/elf/elf_linux_ppc64le"), true, OperatingSystem::Linux, Architecture::PowerPC64LE, Ordering::LittleEndian, ExecutableType::Program)]
    #[case::ppc64le_lib(include_bytes!("../../../testdata/elf/elf_linux_ppc64le.so"), true, OperatingSystem::Linux, Architecture::PowerPC64LE, Ordering::LittleEndian, ExecutableType::Library)]
    #[case::riscv(include_bytes!("../../../testdata/elf/elf_linux_riscv64"), true, OperatingSystem::Linux, Architecture::RISCV64, Ordering::LittleEndian, ExecutableType::Library /* Not really a library, but compiler said it is*/)]
    #[case::s390x(include_bytes!("../../../testdata/elf/elf_linux_s390x"), true, OperatingSystem::Linux, Architecture::S390x, Ordering::BigEndian, ExecutableType::Library /* Not really a library, but compiler said it is*/)]
    #[case::x86_haiku(include_bytes!("../../../testdata/elf/elf_haiku_x86"), false, OperatingSystem::Haiku, Architecture::X86, Ordering::LittleEndian, ExecutableType::Library /* Not really a library, but compiler said it is*/)]
    #[case::x86_64_freebsd(include_bytes!("../../../testdata/elf/elf_freebsd_x86_64"), true, OperatingSystem::FreeBSD, Architecture::X86_64, Ordering::LittleEndian, ExecutableType::Program)]
    #[test]
    fn binaries(
        #[case] bytes: &[u8],
        #[case] is_64bit: bool,
        #[case] os: OperatingSystem,
        #[case] arch: Architecture,
        #[case] ordering: Ordering,
        #[case] elf_type: ExecutableType,
    ) {
        let elf = Elf::from(bytes).unwrap();
        eprintln!("{elf}");
        assert_eq!(elf.is64bit, is_64bit);
        assert_eq!(elf.os, os);
        if elf_type == ExecutableType::Program
            && arch != Architecture::ARM64
            && arch != Architecture::RISCV64
            && arch != Architecture::S390x
        {
            assert!(elf.interpreter.is_some());
        }
        assert_eq!(elf.executable_type, elf_type);
        assert_eq!(elf.ordering, ordering);
        assert_eq!(elf.arch, arch);
    }

    #[test]
    fn hobbit() {
        // This test is separate, as the parser reports this program is for Linux when it's for BeOS.
        // Interpreter not found.
        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);
    }
}
