// Panopticon - A libre program analysis library for machine code
// Copyright (C) 2014-2018  The Panopticon Developers
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 2.1 of the License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

//! Loader for 32 and 64-bit ELF, PE, and Mach-o files.

use std::fs::File;
use std::path::Path;

use memmap::{MmapOptions, Mmap};
use goblin::{self, Hint, pe, elf, mach};
use goblin::elf::program_header;
use goblin::pe::section_table::SectionTable;

use {Str, Region, Result};

/// CPU the binary file is intended for.
#[derive(Clone,Copy,Debug)]
pub enum Machine {
    /// 8-bit AVR
    Avr,
    /// AMD64
    Amd64,
    /// Intel x86
    Ia32,
}

/// Named pointer inside a region.
#[derive(Debug)]
pub struct Pointer {
    /// Region name
    pub segment: Str,
    /// Pointed-to name
    pub name: Str,
    /// Offset from the start of the region
    pub offset: u64,
}

/// Binary file meta information. Starting point for disassembler routines.
#[derive(Debug)]
pub struct Content {
    /// CPU ISA
    pub machine: Machine,
    /// Public functions
    pub entry_points: Vec<Pointer>,
    /// Regions
    pub segments: Vec<Region>,
    /// Imported functions
    pub symbolic: Vec<Pointer>,
}

impl Content {
    /// Load an ELF or PE file from disk and creates a `Project` from it. Returns the `Project` instance and
    /// the CPU its intended for.
    pub fn load(path: &Path) -> Result<Self> {
        Self::load_all(path).and_then(|x| match x.into_iter().next() {
            Some(x) => Ok(x),
            None => Err(format!("Not a supported file format").into()),
        })
    }

    /// Load all programs inside a ELF or PE file.
    pub fn load_all(path: &Path) -> Result<Vec<Self>> {
        let fd = File::open(path)?;
        let map = unsafe { Mmap::map(&fd)? };
        let mut magic = [0u8; 16];

        magic.copy_from_slice(&map[0..16]);

        match goblin::peek_bytes(&magic)? {
            Hint::Unknown(magic) => Err(format!("Tried to load an unknown file. Magic: {}", magic).into()),
            Hint::Elf(_) => Self::load_elf(&map, &fd, path).map(|x| vec![x]),
            Hint::PE => Self::load_pe(&map, &fd, path).map(|x| vec![x]),
            Hint::Mach(_) => Self::load_mach(&map, &fd, path).map(|x| vec![x]),
            Hint::MachFat(_) => {
                unimplemented!()
            }
            Hint::Archive => {
                unimplemented!()
            }
        }
    }

    fn load_mach(map: &Mmap, fd: &File, path: &Path) -> Result<Content> {
        let binary = mach::MachO::parse(&map[..], 0)?;
        debug!("mach: {:#?}", &binary);
        let mut base = 0x0;
        let cputype = binary.header.cputype;
        let (machine, addr_bits) = match cputype {
            mach::cputype::CPU_TYPE_X86 => {
                (Machine::Ia32, 32)
            }
            mach::cputype::CPU_TYPE_X86_64 => {
                (Machine::Amd64, 64)
            }
            machine => {
                return Err(
                    format!(
                        "Unsupported machine ({:#x})",
                        machine,
                        )
                    .into()
                    )
            }
        };
        let mut regs = Vec::default();
        let mut syms = Vec::default();
        let mut entries = Vec::default();

        for segment in &*binary.segments {
            let offset = segment.fileoff as usize;
            let filesize = segment.filesize as usize;
            if offset + filesize > map.len() {
                return Err(
                    format!(
                        "Failed to read segment: range {:?} greater than len {}",
                        offset..offset + filesize,
                        map.len()
                        )
                    .into()
                    );
            }
            let start = segment.vmaddr;
            let name = segment.name()?;

            debug!(
                "Load mach segment {:?}: {} bytes segment to {:#x}",
                name,
                segment.vmsize,
                start
                );

            let reg = if filesize > 0 {
                Self::load_section(name.to_string().into(), fd, path, filesize, addr_bits, offset, start)?
            } else {
                Region::undefined(name.to_string(), addr_bits)
            };
            regs.push(reg);

            if name == "__TEXT" {
                base = segment.vmaddr;
                debug!("Setting vm address base to {:#x}", base);
            }
        }

        let entry = binary.entry;

        if entry != 0 {
            match Self::resolve_reference(entry as u64, "(entry)", &regs) {
                Some(e) => { entries.push(e); }
                None => { /* do nothing */ }
            }
        }

        for export in binary.exports()? {
            if export.offset != 0 {
                debug!("adding: {:?}", &export);

                match Self::resolve_reference(export.offset as u64 + base, &export.name, &regs) {
                    Some(e) => { entries.push(e); }
                    None => { /* do nothing */ }
                }
            }
        }

        for import in binary.imports()? {
            debug!("Import {}: {:#x}", import.name, import.offset);

            match Self::resolve_reference(import.offset as u64, import.name, &regs) {
                Some(e) => { syms.push(e); }
                None => { /* do nothing */ }
            }
        }

        let c = Content{
            machine: machine,
            entry_points: entries,
            symbolic: syms,
            segments: regs,
        };
        Ok(c)
    }

    /// Parses an ELF 32/64-bit binary from `bytes` and creates a `Project` from it. Returns the `Project` instance and
    /// the CPU its intended for.
    fn load_elf(map: &Mmap, fd: &File, path: &Path) -> Result<Content> {
        use std::collections::HashSet;

        let binary = elf::Elf::parse(&map[..])?;
        let mut regs = vec![];
        let mut entries: Vec<Pointer> = vec![];
        let mut syms: Vec<Pointer> = vec![];

        debug!("elf: {:#?}", &binary);

        let (machine, addr_bits) = match binary.header.e_machine {
            elf::header::EM_X86_64 => {
                (Machine::Amd64, 64)
            }
            elf::header::EM_386 => {
                (Machine::Ia32, 32)
            }
            elf::header::EM_AVR => {
                (Machine::Avr, 22)
            }
            machine => return Err(format!("Unsupported machine: {}", machine).into()),
        };


        for (idx, ph) in binary.program_headers.iter().enumerate() {
            if ph.p_type == program_header::PT_LOAD {
                debug!("Load ELF {} bytes segment to {:#x}", ph.p_filesz, ph.p_vaddr);

                let reg = Self::load_section(format!("sec{}", idx).into(), fd, path, ph.p_filesz as usize, addr_bits, ph.p_offset as usize, ph.p_vaddr)?;
                regs.push(reg);
            }
        }
        let mut seen_syms = HashSet::<u64>::new();

        // add dynamic symbol information (non-strippable)
        for sym in &binary.dynsyms {
            let name = &binary.dynstrtab[sym.st_name];

            Self::add_elf_symbol(&sym, name, &regs, &mut syms, &mut entries);
            seen_syms.insert(sym.st_value);

            let name = &binary.dynstrtab[sym.st_name];
            if !Self::resolve_elf_import_address(&binary.pltrelocs, name, &regs, &binary, &mut syms) {
                if sym.is_function() {
                    if !Self::resolve_elf_import_address(&binary.dynrelas, name, &regs, &binary, &mut syms) {
                        Self::resolve_elf_import_address(&binary.dynrels, name, &regs, &binary, &mut syms);
                    }
                }
            }
        }

        // add strippable symbol information
        for sym in &binary.syms {
            let name = &binary.strtab[sym.st_name];
            if !seen_syms.contains(&sym.st_value) {
                Self::add_elf_symbol(&sym, &name, &regs, &mut syms, &mut entries);
            }
            seen_syms.insert(sym.st_value);
        }

        // binary entry point
        match Self::resolve_reference(binary.entry, "(entry)", &regs) {
            Some(e) => { entries.push(e); }
            None => { /* do nothing */ }
        }

        let c = Content{
            machine: machine,
            entry_points: entries,
            symbolic: syms,
            segments: regs,
        };
        Ok(c)
    }

    fn add_elf_symbol(sym: &elf::Sym, name: &str, regs: &[Region], syms: &mut Vec<Pointer>, entries: &mut Vec<Pointer>) {
        let name = name.to_string();
        let addr = sym.st_value;

        debug!("Symbol: {} @ 0x{:x}: {:?}", name, addr, sym);

        if sym.is_function() {
            match Self::resolve_reference(addr, &name, &regs) {
                Some(e) =>
                    if sym.is_import() { syms.push(e); }
                    else { entries.push(e); },
                None => { /* do nothing */ }
            }
        }
    }

    fn resolve_elf_import_address(relocs: &elf::RelocSection, name: &str, regs: &[Region], binary: &elf::Elf, syms: &mut Vec<Pointer>) -> bool {
        for reloc in relocs.iter() {
            if let Some(pltsym) = &binary.dynsyms.get(reloc.r_sym) {
                let pltname = &binary.dynstrtab[pltsym.st_name];
                if pltname == name {
                    debug!("Import match {}: {:#x} {:?}", name, reloc.r_offset, pltsym);

                    match Self::resolve_reference(reloc.r_offset as u64, name.into(), &regs) {
                        Some(e) => { syms.push(e); }
                        None => { /* do nothing */ }
                    }
                    return true;
                }
            }
        }
        false
    }

    /// Parses a PE32/PE32+ file from `bytes` and create a project from it.
    fn load_pe(map: &Mmap, fd: &File, path: &Path) -> Result<Self> {
        let pe = pe::PE::parse(&map[..])?;
        let image_base = pe.image_base as u64;
        let mut regs = vec![];
        let mut entries = vec![];
        let entry = (pe.image_base + pe.entry) as u64;
        let size = fd.metadata()?.len();
        let address_bits = if pe.is_64 { 64 } else { 32 };

        debug!("loading PE: {:#?}", &pe);

        // alloc. segment for each section
        for section in &pe.sections {
            debug!("loading PE section: {:?}", section.name);

            let ret = Self::load_pe_section(section, fd, path, size as usize, image_base, address_bits)?;
            let addr = section.virtual_address as u64 + pe.image_base as u64;

            if entry >= addr && entry < addr + section.virtual_size as u64 {
                let ent = Pointer{
                    segment: ret.name().clone(),
                    name: "main".into(),
                    offset: entry,
                };
                entries.push(ent);
            }

            regs.push(ret);
        }

        debug!("PE file entry at {:#x}", entry);

        // add exported functions as entry points
        for (i, export) in pe.exports.iter().enumerate() {
            debug!("adding PE export: {:?}", &export);

            let nam = export.name.map(|x| x.to_string()).unwrap_or(format!("exp{}", i));
            match Self::resolve_reference(export.rva as u64 + image_base, &nam, &regs) {
                Some(e) => {
                    entries.push(e);
                }
                None => {
                    error!("PE export {:?} at {:#x} not mapped", export.name, export.rva);
                }
            }
        }

        let mut syms = vec![];

        // add imports as symbolic functions
        for import in pe.imports {
            debug!("adding PE import: {:?}", &import);

            match Self::resolve_reference(import.rva as u64 + image_base, &import.name, &regs) {
                Some(e) => {
                    syms.push(e);
                }
                None => {
                    error!("PE import {} at {:#x} not mapped", import.name, import.rva);
                }
            }
        }

        let c = Content{
            machine: Machine::Ia32,
            entry_points: entries,
            symbolic: syms,
            segments: regs,
        };
        Ok(c)
    }

    fn load_pe_section(sec: &SectionTable, fd: &File, path: &Path, fsize: usize, image_base: u64, address_bits: usize) -> Result<Region> {
        let voffset = sec.virtual_address as u64 + image_base;
        let vsize = sec.virtual_size as u64;
        let offset = sec.pointer_to_raw_data as usize;
        let size = sec.size_of_raw_data as usize;
        let name = String::from_utf8(sec.name[..].to_vec())?;

        if size > 0 {
            if offset + size > fsize {
                return Err(format!("PE section out of range: {:#x} + {:#x} >= {:#x}",offset,size,fsize).into());
            }

            // XXX: implemented larger vsize in Region
            debug!("PE section {} mapped from {:?} to {:?}", name, offset..offset + size, voffset..voffset + size as u64);
            Self::load_section(name.into(), fd, path, size, address_bits, offset, voffset)
        } else {
            debug!("PE section {} mapped to {:?}", name, voffset..voffset + vsize);
            Ok(Region::undefined(name, address_bits))
        }
    }

    fn resolve_reference(addr: u64, name: &str, regs: &[Region]) -> Option<Pointer> {
        for r in regs {
            if r.in_range(addr..addr + 1) {
                return Some(Pointer{
                    segment: r.name().clone(),
                    name: name.to_string().into(),
                    offset: addr,
                });
            }
        }

        None
    }

    fn load_section(name: Str, fd: &File, path: &Path, size: usize, address_bits: usize, file_offset: usize, load_offset: u64) -> Result<Region> {
        let mmap = unsafe {
            MmapOptions::new()
                .len(size)
                .offset(file_offset)
                .map(fd)?
        };
        Ok(Region::from_mmap(name, address_bits, mmap, path.to_path_buf(), file_offset as u64, load_offset))
    }
}

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

    #[test]
    fn pe() {
        use std::path::Path;

        let p = format!("{}/tests/data/test.exe", env!("CARGO_MANIFEST_DIR"));
        let c = Content::load(Path::new(&p)).unwrap();

        println!("{:?}", c);
    }

    #[test]
    fn elf() {
        use std::path::Path;

        let p = format!("{}/tests/data/dynamic-32", env!("CARGO_MANIFEST_DIR"));
        let c = Content::load(Path::new(&p)).unwrap();

        println!("{:?}", c);
    }

    #[test]
    fn mach() {
        use std::path::Path;

        let p = format!("{}/tests/data/deadbeef.mach", env!("CARGO_MANIFEST_DIR"));
        let c = Content::load(Path::new(&p)).unwrap();

        println!("{:?}", c);
    }

}
