#![allow(unused_imports)]

mod helper;

use nix::sys::ptrace;
use nix::sys::signal::{Signal, kill};
use nix::unistd::Pid;
use std::arch::asm;
use std::ffi::CString;

use helper::*;
use iced_x86::code_asm::{eax, r8, r9, r10, rax, rbp, rcx, rdi, rdx, rsi, rsp};
use libc::{RTLD_NEXT, c_void, dlsym};
use libc::{ptrace, user_regs_struct};
use std::fs;
use std::io::BufRead;

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

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

    let exe = proc.get_exe()?;
    let maps = proc.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 Some(seg_rw) = map.first_rw_segment(&exe) else {
        return Err(Box::new(std::io::Error::new(
            std::io::ErrorKind::Other,
            "first rw segment not found",
        )));
    };

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

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

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

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

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

    // Do inject here

    let c = inject3(&proc, seg_rw)?;

    // End inject logics

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

    ptrace::detach(pid, None)?;
    ptrace::detach(Pid::from_raw(c), None)?;

    Ok(())
}
