#![allow(unused_imports)]

mod helper;

use nix::sys::ptrace;
use nix::sys::wait::waitpid;
use nix::unistd::Pid;
use std::arch::asm;
use std::ffi::CString;

use helper::*;
use iced_x86::code_asm::asm_traits::CodeAsmJmp;
use iced_x86::{Instruction, code_asm::*};
use libc::user_regs_struct;
use libc::{RTLD_NEXT, c_void, dlsym};
use std::fs;
use std::io::BufRead;

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

fn inject1(
    pid: Pid,
    seg_rw: (u64, u64),
    regs: user_regs_struct,
) -> Result<(), Box<dyn std::error::Error>> // Simple injection
{
    let injected_data = "You are injected. \r\n";
    write_memory_vm(pid, seg_rw.0 as usize, &injected_data.as_bytes())?;

    let mut asm = CodeAssembler::new(64)?;

    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

    let injected_inst = asm.assemble(regs.rip as u64)?;
    write_memory_ptrace(pid, regs.rip as usize, &injected_inst)?;

    // Continue target
    ptrace::cont(pid, None)?;
    waitpid(pid, None)?;

    Ok(())
}

fn inject2(
    pid: Pid,
    seg_rw: (u64, u64),
    regs: user_regs_struct,
) -> Result<(), Box<dyn std::error::Error>> // ld injection
{
    // Get the absolute path to our shared library
    let lib_path = fs::canonicalize("./target/debug/libproject_hbj_attacker.so")?
        .to_string_lossy()
        .into_owned();
    let cpid = nix::unistd::getpid().to_string();

    // Read our own process memory maps to find libc base address
    let self_maps = fs::read_to_string(format!("/proc/{}/maps", cpid))?;
    let self_map_lines = self_maps.lines().collect::<Vec<&str>>();
    let mut dlopen_offset: u64 = 0;

    // Find libc base address in our own process
    let Some(libc_base_local) = module_base_address(&self_map_lines, "libc.so") else {
        return Err(Box::new(std::io::Error::new(
            std::io::ErrorKind::Other,
            "libc not found",
        )));
    };

    println!("{GREEN}[local]{RESET} libc base: {:#016x}", libc_base_local);

    // Use dlsym to get the address of dlopen in our own process
    unsafe {
        let dlopen_addr_local = dlsym(RTLD_NEXT, b"dlopen\0".as_ptr() as *const _);
        // Calculate offset of dlopen from libc base in our process
        dlopen_offset = dlopen_addr_local as u64 - libc_base_local;
    }

    println!(
        "{GREEN}[local]{RESET} dlopen offset = {:#016x}",
        dlopen_offset
    );

    // Read target process memory maps to find its libc base address
    let target_maps = fs::read_to_string(format!("/proc/{}/maps", pid))?;
    let target_map_lines = target_maps.lines().collect::<Vec<&str>>();

    // Find libc base address in target process
    let Some(libc_base_target) = module_base_address(&target_map_lines, "libc.so") else {
        return Err(Box::new(std::io::Error::new(
            std::io::ErrorKind::Other,
            "libc not found",
        )));
    };

    println!(
        "{GREEN}[trace]{RESET} libc base = {:#016x}",
        libc_base_target
    );

    // Calculate dlopen address in target process using the same offset
    let target_dlopen_addr = libc_base_target + dlopen_offset;
    println!(
        "{GREEN}[trace]{RESET} dlopen address = {:#016x}",
        target_dlopen_addr
    );

    // Start Inject
    let c_lib_path = CString::new(lib_path).unwrap();
    write_memory_vm(pid, 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 mut asm = CodeAssembler::new(64)?;
    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
    let injected_inst = asm.assemble(regs.rip as u64)?;
    write_memory_ptrace(pid, regs.rip as usize, &injected_inst)?;
    println!(
        "{GREEN}[trace]{RESET} write instructions to {:#016x}",
        regs.rip
    );

    // Continue target
    ptrace::cont(pid, None)?;
    println!("{GREEN}[trace]{RESET} running...");
    waitpid(pid, None)?;
    println!("{GREEN}[trace]{RESET} int3!");
    Ok(())
}

fn inject3(
    pid: Pid,
    seg_rw: (u64, u64),
    regs: user_regs_struct,
) -> Result<(), Box<dyn std::error::Error>> // thread inject
{
    // Alloc rwx memory
    let mut asm = CodeAssembler::new(64)?;

    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_SHARED | 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

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

    // Continue target
    ptrace::cont(pid, None)?;
    println!("{GREEN}[trace]{RESET} running...");
    waitpid(pid, None)?;
    println!("{GREEN}[trace]{RESET} int3!");

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

    let injected_data = "I am injected thread, I am running... \r\n";
    write_memory_vm(pid, page_addr + 0x500, &injected_data.as_bytes())?;
    write_memory_vm(pid, page_addr + 0x600, &1i64.to_le_bytes())?;
    write_memory_vm(pid, page_addr + 0x608, &0u64.to_le_bytes())?;

    asm = CodeAssembler::new(64)?;
    // Construct inject payload
    let mut target_label = asm.create_label();

    asm.set_label(&mut target_label)?;
    asm.mov(rax, 1u64)?; // Syscall 1 (write)
    asm.mov(rdi, 1u64)?; // Fd 1 (STDOUT)
    asm.mov(rsi, (page_addr + 0x500) as u64)?; // Buffer pointer (Here is page_addr + 0x500)
    asm.mov(rdx, injected_data.as_bytes().len() as u64)?; // Buffer length
    asm.syscall()?; // Syscall interrupt

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

    asm.jmp(target_label)?; // Jmp back to loop

    let injected_payload = asm.assemble(page_addr as u64)?;
    write_memory_vm(pid, page_addr, &injected_payload)?;
    println!("{GREEN}[trace]{RESET} write payload to {:#016x}", page_addr);

    // Start Trigger
    let regs = ptrace::getregs(pid)?;

    asm = CodeAssembler::new(64)?;

    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 + 0x800) 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

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

    // Continue target
    ptrace::cont(pid, None)?;
    waitpid(pid, None)?;

    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Find our target program
    let pid = Pid::from_raw(get_pid_by_name("target")?);

    let target = fs::read_link(format!("/proc/{}/exe", pid))?
        .to_string_lossy()
        .into_owned();
    let content = fs::read_to_string(format!("/proc/{}/maps", pid))?;
    let lines: Vec<&str> = content
        .lines()
        .filter(|&line| !line.is_empty() && line.contains(&target))
        .collect();

    for line in &lines {
        println!("{GREEN}[memory map]{RESET} {}", line);
    }

    let Some(seg_rw) = first_rw_segment(&lines) else {
        return Err(Box::new(std::io::Error::new(
            std::io::ErrorKind::Other,
            "first rw segment not found",
        )));
    };

    let Some(seg_x) = first_exec_segment(&lines) else {
        return Err(Box::new(std::io::Error::new(
            std::io::ErrorKind::Other,
            "first exec segment not found",
        )));
    };

    ptrace::attach(pid)?;
    waitpid(pid, None)?;
    ptrace::step(pid, None)?;
    waitpid(pid, None)?;

    // ↓ Old behavior, but maybe the process stay in their libraries forever ?

    // loop{
    //    // Single-stepping, so that RIP returns to the user space of the process itself,
    //    // rather than in some other library
    //    let regs = ptrace::getregs(pid)?;
    //    if is_address_in_range(regs.rip, &lines)
    //    {
    //        println!("{GREEN}[trace]{RESET} Address: {:#016x}", regs.rip);
    //        break;
    //    }
    //    println!("{GREEN}[trace]{RESET} Skipped: {:#016x}", regs.rip);
    //    ptrace::step(pid, None)?;
    //    waitpid(pid, None)?;
    //}

    // Save context
    let regs = ptrace::getregs(pid)?; // Save current registers
    let buffer = read_memory_vm(pid, seg_x.0 as usize, 128)?; // Save current memory context
    let buffer_rw = read_memory_vm(pid, seg_rw.0 as usize, 128)?; // Save current rw memory

    println!("{GREEN}[trace]{RESET} Seg_x.0 is {:#016x}", seg_x.0);

    write_memory_ptrace(pid, seg_x.0 as usize, &[0x90u8; 128])?;
    ptrace::setregs(pid, user_regs_struct {
        rip: seg_x.0,
        ..regs
    })?;

    // Do inject here

    inject3(
        pid,
        seg_rw,
        user_regs_struct {
            rip: seg_x.0,
            ..regs
        },
    )?;

    // End inject logics

    // Restore context
    ptrace::setregs(pid, regs)?;
    write_memory_ptrace(pid, seg_x.0 as usize, &buffer)?;
    write_memory_vm(pid, seg_rw.0 as usize, &buffer_rw)?;

    ptrace::detach(pid, None)?;

    Ok(())
}
