use std::ffi::CString;
use std::fs;
use iced_x86::code_asm::{eax, r10, r8, r9, rax, rbp, rcx, rdi, rdx, rsi, rsp};
use nix::sys::ptrace;
use nix::sys::wait::{waitpid, WaitPidFlag};
use nix::unistd::Pid;
use crate::helper::{assemble, Process};


const GREEN: &str = "\x1b[32m";
const RESET: &str = "\x1b[0m";

pub fn inject1(proc: &Process, seg_rw: (u64, u64)) -> Result<(), Box<dyn std::error::Error>> // Simple injection
{
    let regs = ptrace::getregs(proc.get_pid())?;
    let injected_data = "You are injected. \r\n";
    proc.write_memory_vm(seg_rw.0 as usize, &injected_data.as_bytes())?;
    println!(
        "{GREEN}[trace]{RESET} write \"{}\" to {:#016x}",
        injected_data, seg_rw.0
    );

    let injected_inst = assemble(regs.rip as u64, |asm| {
        asm.mov(rax, 1u64)?; // Syscall 1 (write)
        asm.mov(rdi, 1u64)?; // Fd 1 (STDOUT)
        asm.mov(rsi, seg_rw.0)?; // Buffer pointer (Here is rw segment in target)
        asm.mov(rdx, injected_data.as_bytes().len() as u64)?; // Buffer length
        asm.syscall()?; // Syscall interrupt
        asm.int3()?; // (Important!!!) Use int3 interrupt to retrieve control flow
        Ok(())
    })?;

    proc.write_memory_ptrace(regs.rip as usize, &injected_inst)?;
    println!(
        "{GREEN}[trace]{RESET} write instructions to {:#016x}",
        regs.rip
    );

    // Continue target
    ptrace::cont(proc.get_pid(), None)?;
    println!("{GREEN}[trace]{RESET} continue from {:#016x}", regs.rip);
    proc.wait();

    let regs = ptrace::getregs(proc.get_pid())?;
    println!("{GREEN}[trace]{RESET} int3 at {:#016x}", regs.rip);

    Ok(())
}

pub fn inject2(proc: &Process, seg_rw: (u64, u64)) -> Result<(), Box<dyn std::error::Error>> // ld injection
{
    let regs = ptrace::getregs(proc.get_pid())?;
    // Get the absolute path to our shared library
    let lib_path = fs::canonicalize("./target/debug/libproject_hbj_attacker.so")?
        .to_string_lossy()
        .into_owned();

    // Start Inject
    let c_lib_path = CString::new(lib_path).unwrap();
    proc.write_memory_vm(seg_rw.0 as usize, c_lib_path.as_bytes_with_nul())?;
    println!(
        "{GREEN}[trace]{RESET} write {} to {:#016x}",
        &c_lib_path.to_string_lossy(),
        seg_rw.0
    );

    let Some(target_dlopen_addr) = proc.find_remote_proc("/usr/lib/libc.so.6", "dlopen") else {
        return Err(Box::new(std::io::Error::new(
            std::io::ErrorKind::Other,
            "first rw segment not found",
        )));
    };

    let injected_inst = assemble(regs.rip as u64, |asm| {
        asm.mov(rdi, seg_rw.0)?; // Param 1: Filename
        asm.mov(rsi, 2u64)?; // Param 2: Flag, 2(RTLD_NOW)
        asm.call(target_dlopen_addr)?; // Syscall interrupt
        asm.int3()?; // (Important!!!) Use int3 interrupt to retrieve control flow
        Ok(())
    })?;

    proc.write_memory_ptrace(regs.rip as usize, &injected_inst)?;
    println!(
        "{GREEN}[trace]{RESET} write instructions to {:#016x}",
        regs.rip
    );

    // Continue target
    ptrace::cont(proc.get_pid(), None)?;
    println!("{GREEN}[trace]{RESET} continue from {:#016x}", regs.rip);
    proc.wait();

    let regs = ptrace::getregs(proc.get_pid())?;
    println!("{GREEN}[trace]{RESET} int3 at {:#016x}", regs.rip);

    Ok(())
}

pub fn inject3(proc: &Process, seg_rw: (u64, u64)) -> Result<i32, Box<dyn std::error::Error>> // thread inject
{
    let regs = ptrace::getregs(proc.get_pid())?;
    // Alloc rwx memory

    let injected_inst = assemble(regs.rip as u64, |asm| {
        asm.mov(rax, 9u64)?; // Syscall 9 (mmap)

        asm.mov(rdi, 0u64)?; // Addr
        asm.mov(rsi, 4096u64)?; // Length, we alloc a page (4K)
        asm.mov(
            rdx,
            (libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC) as u64,
        )?; // Set protect to rwx
        asm.mov(r10, (libc::MAP_PRIVATE | libc::MAP_ANONYMOUS) as u64)?; // Private and anonymous
        asm.mov(r8, 01i64)?; // Fd (-1 because we want anonymous)
        asm.mov(r9, 0u64)?; // Offset

        asm.syscall()?; // Syscall interrupt
        asm.int3()?; // (Important!!!) Use int3 interrupt to retrieve control flow
        Ok(())
    })?;

    proc.write_memory_ptrace(regs.rip as usize, &injected_inst)?;
    println!(
        "{GREEN}[trace]{RESET} write instructions to {:#016x}",
        regs.rip
    );

    // Continue target
    ptrace::cont(proc.get_pid(), None)?;
    println!("{GREEN}[trace]{RESET} continue from {:#016x}", regs.rip);
    proc.wait();

    println!(
        "{GREEN}[trace]{RESET} int3 at {:#016x}",
        ptrace::getregs(proc.get_pid())?.rip
    );

    let regs_after_map = ptrace::getregs(proc.get_pid())?;
    let page_addr = regs_after_map.rax as usize;
    println!(
        "{GREEN}[trace]{RESET} allocated page is at {:#016x}",
        page_addr
    );

    let injected_data = "[%d] I am the injected thread, I am running... \r\n";
    proc.write_memory_vm(
        page_addr + 0x200,
        &CString::new(injected_data).unwrap().as_bytes_with_nul(),
    )?;
    proc.write_memory_vm(page_addr + 0x300, &1i64.to_le_bytes())?;
    proc.write_memory_vm(page_addr + 0x308, &0u64.to_le_bytes())?;

    let printf = proc.find_remote_proc("/usr/lib/libc.so.6", "printf").unwrap();

    // Construct inject payload
    let injected_payload = assemble(page_addr as u64, |asm| {
        let mut target_label = asm.create_label();

        asm.mov(rbp, rsp)?;
        asm.mov(rcx, 0u64)?;
        asm.set_label(&mut target_label)?;

        asm.add(rcx, 1i32)?;
        asm.push(rcx)?;

        asm.mov(rdi, (page_addr + 0x200) as u64)?;
        asm.mov(rsi, rcx)?;
        asm.call(printf)?;

        asm.mov(rax, 35u64)?; // Syscall 35 (nano sleep)
        asm.mov(rdi, (page_addr + 0x300) as u64)?; // Req
        asm.mov(rsi, 0u64)?; //Rem
        asm.syscall()?; // Syscall interrupt

        asm.pop(rcx)?;
        asm.jmp(target_label)?; // Jmp back to loop
        Ok(())
    })?;
    proc.write_memory_vm(page_addr, &injected_payload)?;
    println!("{GREEN}[trace]{RESET} write payload to {:#016x}", page_addr);

    // Start Trigger
    // let regs = ptrace::getregs(proc.get_pid())?;
    ptrace::setregs(proc.get_pid(), regs)?;

    let injected_trigger = assemble(regs.rip as u64, |asm| {
        asm.mov(rax, 56u64)?; // Syscall 56 (clone)

        asm.mov(
            rdi,
            (libc::CLONE_VM
                | libc::CLONE_FS
                | libc::CLONE_FILES
                | libc::CLONE_SIGHAND
                | libc::CLONE_THREAD) as u64,
        )?; // Flags
        asm.mov(rsi, (page_addr + 0x1000) as u64)?; // Stack top

        asm.mov(rdx, 0u64)?; // parent_tid = NULL
        asm.mov(r10, 0u64)?; // child_tid = NULL
        asm.mov(r8, 0u64)?; // tls = NULL

        asm.syscall()?; // Syscall interrupt
        asm.test(eax, eax)?; // Syscall returns zero?
        asm.jz(page_addr as u64)?;

        asm.int3()?; // (Important!!!) Use int3 interrupt to retrieve control flow

        Ok(())
    })?;

    proc.write_memory_ptrace(regs.rip as usize, &injected_trigger)?;
    println!("{GREEN}[trace]{RESET} write trigger to {:#016x}", regs.rip);

    ptrace::cont(proc.get_pid(), None)?;
    println!("{GREEN}[trace]{RESET} continue from {:#016x}", regs.rip);
    proc.wait();

    let regs = ptrace::getregs(proc.get_pid())?;
    let pid_new_thread = Pid::from_raw(regs.rax as i32);
    println!("{GREEN}[trace]{RESET} int3 at {:#016x}", regs.rip);
    println!(
        "{GREEN}[trace]{RESET} new thread is {}, which will be suspend.",
        pid_new_thread
    );

    ptrace::attach(pid_new_thread)?;
    waitpid(pid_new_thread, Some(WaitPidFlag::WUNTRACED))?;
    println!("{GREEN}[trace]{RESET} attached new thread.");

    loop {
        let regs = ptrace::getregs(pid_new_thread)?;
        println!(
            "{GREEN}[trace]{RESET} rip in new thread is {:#016x}.",
            regs.rip
        );

        if regs.rip >= page_addr as u64 && regs.rip < (page_addr + 0x1000) as u64 {
            println!("{GREEN}[trace]{RESET} rip in new thread return to inject payload.");
            break;
        }

        ptrace::step(pid_new_thread, None)?;
        waitpid(pid_new_thread, Some(WaitPidFlag::WUNTRACED))?;
    }

    Ok(pid_new_thread.as_raw())
}
