pes 0.2.0

use std::fs;
use std::path::PathBuf;

use super::images;
use super::pe::Pe;

#[derive(PartialEq, Debug)]
enum Status {
    FileBuffer,
    ImageBuffer,
}

pub struct PeFile {
    image: Vec<u8>,
    status: Status,
}

// impl Pe
impl Pe for PeFile {
    fn image(&self) -> &[u8] {
        &self.image
    }
}

// private
impl PeFile {
    fn valide_header(&self) -> Result<(), &str> {
        if self.dos_header().e_magic != images::IMAGE_DOS_SIGNATURE {
            return Err("Bad Dos Signature!");
        }

        if self.nt_headers().Signature != images::IMAGE_NT_HEADERS_SIGNATURE {
            return Err("Bad Pe Signature!");
        }

        Ok(())
    }
}

// pub
impl PeFile {
    pub fn from_bytes(bytes: Vec<u8>) -> PeFile {
        PeFile {
            image: bytes,
            status: Status::FileBuffer,
        }
    }

    pub fn read(path: PathBuf) -> Result<PeFile, std::io::Error> {
        let data: Vec<u8> = fs::read(path)?;
        Ok(PeFile {
            image: data,
            status: Status::FileBuffer,
        })
    }

    /// 将 PeFile.image 从 file view 转为 image view
    pub fn into_image_view(self) -> Self {
        assert!(self.status == Status::FileBuffer);

        // 1. 申请 image buffer 空间
        let (size_of_images, size_of_headers) = {
            (
                self.opt_header().SizeOfImage as usize,
                self.opt_header().SizeOfHeaders as usize,
            )
        };

        let mut image_view = vec![0u8; size_of_images];

        // 2. Copy headers
        unsafe {
            let dest_headers = image_view.get_unchecked_mut(..size_of_headers);
            let src_headers = self.image().get_unchecked(..size_of_headers);
            dest_headers.copy_from_slice(src_headers);
        }

        // 3. Copy sections
        for sec in self.section_headers() {
            let sec: &images::IMAGE_SECTION_HEADER = sec;
            let dest = image_view.get_mut(
                sec.VirtualAddress as usize..(sec.VirtualAddress + sec.SizeOfRawData) as usize,
            );
            let src = self.image().get(
                sec.PointerToRawData as usize..(sec.PointerToRawData + sec.SizeOfRawData) as usize,
            );
            // Skip invalid sections
            if let (Some(dest), Some(src)) = (dest, src) {
                dest.copy_from_slice(src);
            }
        }

        PeFile {
            image: image_view,
            status: Status::ImageBuffer,
        }
    }

    /// 将 image view 转为 file view
    pub fn into_file_view(self) -> Self {
        assert!(self.status == Status::ImageBuffer);
        // 1. 申请空间
        // 根据最后一节的 PointerToRawData, 和 SizeOfRawData 来计算
        let secs = self.section_headers();
        let last = secs.last().unwrap(); // 没有找到最后一节就 panic
        let size_of_file_image = last.PointerToRawData as usize + last.SizeOfRawData as usize;
        let mut file_view = vec![0u8; size_of_file_image];
        // 2. Copy headers
        let size_of_headers = self.opt_header().SizeOfHeaders as usize;
        unsafe {
            let dest = file_view.get_unchecked_mut(..size_of_headers);
            let src = self.image().get_unchecked(..size_of_headers);
            dest.copy_from_slice(src);
        }
        // 3. Copy sections
        for sec in self.section_headers() {
            let sec: &images::IMAGE_SECTION_HEADER = sec;
            let dest = file_view.get_mut(
                sec.PointerToRawData as usize..(sec.PointerToRawData + sec.SizeOfRawData) as usize,
            );
            let src = self.image().get(
                sec.VirtualAddress as usize..(sec.VirtualAddress + sec.SizeOfRawData) as usize,
            );

            if let (Some(dest), Some(src)) = (dest, src) {
                dest.copy_from_slice(src);
            }
        }

        PeFile {
            image: file_view,
            status: Status::FileBuffer,
        }
    }

    /// 保存到磁盘
    pub fn save_to_disk(self, file_path: PathBuf) -> Result<(), std::io::Error> {
        assert!(self.status == Status::FileBuffer);
        fs::write(file_path, self.image())?;
        Ok(())
    }
}

// print pe info methods
impl PeFile {
    /// 打印  位 pe 文件信息
    fn dos_file_info(&self) {
        assert!(self.image.len() > 0);
        println!("- Pe INFO");
        println!("-- DOS Header");
        println!("---- DOS Signature: {:>#06X}", self.dos_header().e_magic);
        println!("---- lfanew: {:>#010X}", self.dos_header().e_lfanew);

        println!("-- Pe Signature");
        println!("---- Pe Signature: {:#010X}", self.nt_headers().Signature);

        println!("-- File Header");
        println!("---- Machine: {:#06X}", self.file_header().Machine);
        println!(
            "---- NumberOfSections: {:#06X}",
            self.file_header().NumberOfSections
        );
        println!(
            "---- SizeOfOptionalHeader: {:#06X}",
            self.file_header().SizeOfOptionalHeader
        );
        println!(
            "---- Characteristics: {:#06X}",
            self.file_header().Characteristics
        );
    }

    fn section_info(&self) {
        println!("-- Section Headers");
        let secs: &[images::IMAGE_SECTION_HEADER] = self.section_headers();
        for (i, sec) in secs.iter().enumerate() {
            let sec: &images::IMAGE_SECTION_HEADER = sec; // 这句只是为了 rls 提示, 用完注释掉
            println!("---- Section {}", (i + 1));
            println!("-------- Name: {}", sec.name().unwrap());
            println!("-------- VirtualSize: {:#010X}", sec.VirtualSize);
            println!("-------- VirtualAddress: {:#010X}", sec.VirtualAddress);
            println!("-------- SizeOfRawData: {:#010X}", sec.SizeOfRawData);
            println!("-------- PointerToRawData: {:#010X}", sec.PointerToRawData);
            println!("-------- Characteristics: {:#010X}", sec.Characteristics);
        }
    }

    pub fn pe_info(&self) {
        self.dos_file_info();

        println!("-- OptionalHeader");
        println!("---- Magic: {:#06X}", self.opt_header().Magic);
        println!(
            "---- AddressOfEntryPoint: {:#010X}",
            self.opt_header().AddressOfEntryPoint
        );
        println!("---- ImageBase: {:#010X}", self.opt_header().ImageBase);
        println!(
            "---- SectionAlignment: {:#010X}",
            self.opt_header().SectionAlignment
        );
        println!(
            "---- FileAlignment: {:#010X}",
            self.opt_header().FileAlignment
        );
        println!("---- SizeOfImage: {:#010X}", self.opt_header().SizeOfImage);
        println!(
            "---- SizeOfHeaders: {:#010X}",
            self.opt_header().SizeOfHeaders
        );
        println!(
            "---- NumberOfRvaAndSizes: {:#010X}",
            self.opt_header().NumberOfRvaAndSizes
        );

        self.section_info();
    }
}

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

    #[test]
    fn len_test() {
        let v = vec![1, 2];
        assert_eq!(v.len(), 2);
    }

    #[test]
    fn pe_info() {
        let file = "resources/notepad.exe";
        let file = PathBuf::from(file);
        if let Ok(pe_file) = PeFile::read(file) {
            pe_file.pe_info();
        }
    }

    #[test]
    fn into_image_view_works() {
        let file_path = "resources/notepad.exe";
        let file_path = PathBuf::from(file_path);
        let pe_file = PeFile::read(file_path).unwrap();
        assert!(pe_file.status == Status::FileBuffer);
        let pe_file = pe_file.into_image_view();
        assert!(pe_file.status == Status::ImageBuffer);
    }

    #[test]
    fn into_file_view_works() {
        let file_path = "resources/notepad.exe";
        let file_path = PathBuf::from(file_path);
        let pe_file = PeFile::read(file_path).unwrap();
        assert!(pe_file.status == Status::FileBuffer);
        let pe_file = pe_file.into_image_view();
        assert!(pe_file.status == Status::ImageBuffer);
        let pe_file = pe_file.into_file_view();
        assert!(pe_file.status == Status::FileBuffer);
    }

    /// 测试 读取 pe 文件 -> 转为 image buffer -> file buffer -> 存盘
    #[test]
    #[ignore]
    fn file_image_new_file() {
        let file_path = "resources/notepad.exe";
        let file_path = PathBuf::from(file_path);
        let pe_file = PeFile::read(file_path).unwrap();
        assert!(pe_file.status == Status::FileBuffer);
        let pe_file = pe_file.into_image_view();
        assert_eq!(pe_file.status, Status::ImageBuffer);
        let pe_file = pe_file.into_file_view();
        assert_eq!(pe_file.status, Status::FileBuffer);
        let file_path = "resources/rust_new_notepad.exe";
        let file_path = PathBuf::from(file_path);
        let _ = pe_file.save_to_disk(file_path);
    }

    #[test]
    fn rva_to_foa_works() {
        let file_path = "resources/notepad.exe";
        let file_path = PathBuf::from(file_path);
        let pe_file: PeFile = PeFile::read(file_path).unwrap();
        let actual_foa = pe_file
            .rva_to_foa(0x1000u32)
            .expect("Some thing goes wrong!");
        assert_eq!(actual_foa, 0x400u32);
    }

    #[test]
    fn foa_to_rva_works() {
        let fp = "resources/notepad.exe";
        let fp = PathBuf::from(fp);
        let pe_file: PeFile = PeFile::read(fp).unwrap();
        let actual_rva = pe_file.foa_to_rva(0x7c00u32).expect("...");
        assert_eq!(actual_rva, 0x9000u32);
    }
}