use std::error::Error;
use std::ffi::CString;
use anyhow::Context;
use iced_x86::code_asm::*;
use crate::map::MemoryMap;
use crate::processes::{get_pid_by_name, Process};
use nix::unistd::Pid;

use libc::user_regs_struct;
use nix::sys::ptrace;

use crate::asm::assemble;

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

mod disassembly;
mod elf;
mod map;
mod processes;
mod asm;

pub fn plt_hook<F>(proc: &Process, map: &MemoryMap, symbol: &str, payload: F)
    -> Result<(), Box<dyn Error>>
where
    F: Fn(&Process, u64, u64) -> Vec<u8>
{
    let bias = map.module_base_address(&proc.get_exe()?).unwrap_or(0);
    let got_item_ptr = proc.find_got_pointer_plt(symbol).context("Unable to find symbol")?;

    println!("{GREEN}[memory map]{RESET} Bias is {:#016x}", bias);

    let got_item_byte: [u8; 8] = proc.read_memory_vm(got_item_ptr as usize, 8)?
        .try_into()
        .map_err(|_| "Failed to convert Vec to array")?;
    let got_item = u64::from_le_bytes(got_item_byte);
    println!(
        "{GREEN}[memory map]{RESET} got_item = {:#016x}", got_item
    );

    // Alloc r-x private memory
    let regs = ptrace::getregs(proc.get_pid())?;
    let r = proc.execute_once_inplace(&assemble(regs.rip as u64, |asm| {
        asm.mov(rax, 9u64)?; // Syscall 9 (mmap)

        asm.mov(rdi, 0u64)?; // Addr
        asm.mov(rsi, 0x1000u64)?; // Length, we alloc a page (4K)
        asm.mov(
            rdx,
            (libc::PROT_READ | 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
        Ok(())
    })?)?;

    let page_addr = r.rax as u64;
    println!("{GREEN}[trace]{RESET} allocated page is at {:#016x}", page_addr);

    let payload = payload(&proc, page_addr, got_item);
    if payload.len() > 0x1000 { return Err(Box::<dyn Error>::from("payload exceeds 0x1000 bytes")); }

    proc.write_memory_ptrace(page_addr as usize, &payload)?;
    println!("{GREEN}[trace]{RESET} wrote {} bytes of instructions/data to {:#016x}", payload.len(), page_addr);
    proc.write_memory_ptrace(got_item_ptr as usize, &page_addr.to_le_bytes())?;
    println!("{GREEN}[trace]{RESET} rewrote got item at {:#016x} to {:#016x}, old value is {:#016x}",
        got_item_ptr, page_addr, got_item);

    Ok(())
}

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

    let exe = process.get_exe()?;
    let maps = process.get_map_str()?;
    let lines: Vec<&str> = maps.lines().filter(|&line| !line.is_empty()).collect();

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

    let map = MemoryMap::new(&lines);

    let seg_x = map.first_exec_segment(&exe).context("Can't find first exec segment")?;

    ptrace::attach(pid)?;
    process.wait();
    ptrace::step(pid, None)?;
    process.wait();

    // Save context
    let regs = ptrace::getregs(pid)?; // Save current registers

    ptrace::setregs(
        pid,
        user_regs_struct {
            rip: seg_x.0,
            ..regs
        },
    )?;

    // Do inject here

    plt_hook(&process, &map, "write", |proc, addr, old| {
        let inst = assemble(addr, |asm| {
            asm.push(rdi)?;
            asm.push(rsi)?;
            asm.push(rdx)?;
            asm.mov(rdi, 1u64)?;
            asm.mov(rsi, addr + 0x800)?;
            asm.mov(rdx, 7u64)?;
            asm.call(old)?;
            asm.pop(rdx)?;
            asm.pop(rsi)?;
            asm.pop(rdi)?;

            asm.call(old)?;

            asm.mov(rdi, 1u64)?;
            asm.mov(rsi, addr + 0x810)?;
            asm.mov(rdx, 2u64)?;
            asm.call(old)?;

            asm.ret()?;
            Ok(())
        }).unwrap();

        let mut payload = Vec::from([0u8; 0x1000]);
        payload.splice(0..inst.len(), inst);
        payload.splice(0x800..(0x800 + 7), Vec::from("[hook] ".as_bytes()));
        payload.splice(0x810..(0x810 + 2), Vec::from("\r\n".as_bytes()));

        return payload;
    })?;

    // End inject logics

    // Restore context
    ptrace::setregs(pid, regs)?;
    ptrace::detach(pid, None)?;
    
    Ok(())
}