#![allow(unused, nonstandard_style)]

// Utility macro to generata a random seed
//  this has to be done in a macro so that the input values
//      file, line, column
//  use the callers location
#[macro_export]
macro_rules! make_seed {
    () => {
        {
            let seed
                = $crate::lazy_importer::hash_utf8::<0x00>(file!().as_bytes())
                ^ $crate::lazy_importer::Prng::new((column!() << 12) as u64 | (line!() as u64)).next();
            seed
        }
    }
}

pub use make_seed;

#[macro_export]
macro_rules! link {
    ($library:literal $abi:literal $($link_name:literal)? fn $fname:ident ($($name:ident : $arg:ty),*)) => (

        #[allow(non_snake_case)]
        pub mod $fname {
            pub const SEED: u64 = $crate::lazy_importer::make_seed!();
            pub static IMPORT: $crate::lazy_importer::LazyImport<
                { $crate::lazy_importer::hash_utf8::<SEED>($library.as_bytes()) },
                { $crate::lazy_importer::hash_utf8::<SEED>(stringify!($fname).as_bytes()) },
                SEED, ()
            > = $crate::lazy_importer::LazyImport::new();
        }

        #[inline(always)]
        #[allow(non_snake_case)]
        pub fn $fname ($($name: $arg),*) {
            unsafe { core::mem::transmute::<_,extern $abi fn($($arg),*)>($fname::IMPORT.resolve())($($name),*) }
        }
    );
    ($library:literal $abi:literal $($link_name:literal)? fn $fname:ident ($($name:ident : $arg:ty),*) -> $ret:ty) => (

        #[allow(non_snake_case)]
        pub mod $fname {
            pub const SEED: u64 = $crate::lazy_importer::make_seed!();
            pub static IMPORT: $crate::lazy_importer::LazyImport<
                { $crate::lazy_importer::hash_utf8::<SEED>($library.as_bytes()) },
                { $crate::lazy_importer::hash_utf8::<SEED>(stringify!($fname).as_bytes()) },
                SEED, ()
            > = $crate::lazy_importer::LazyImport::new();
        }

        #[inline(always)]
        #[allow(non_snake_case)]
        pub fn $fname ($($name: $arg),*) -> $ret {
            unsafe { core::mem::transmute::<_,extern $abi fn($($arg),*) -> $ret>($fname::IMPORT.resolve())($($name),*) }
        }
    )
}


// ┌─────────────────────────────────────────────────────────────────────────┐
// │                   the main lazy import struct itself                    │
// └─────────────────────────────────────────────────────────────────────────┘


pub struct LazyImport<
    const Module: u64,
    const Import: u64,
    const XorSeed: u64,
    Value
>(
    core::cell::UnsafeCell<u64>,
    core::marker::PhantomData<Value>,
);

unsafe impl<const M: u64, const I: u64, const X: u64, V> Sync for LazyImport<M, I, X, V> {}
// unsafe impl<const M: u64, const I: u64, const X: u64, V> Send for LazyImport<M, I, X, V> {}

impl<const M: u64, const I: u64, const X: u64, V> LazyImport<M, I, X, V> {

    const CIPHER: u64 = const { Prng::new(X).next() };

    pub const fn new() -> Self {
        Self(
            core::cell::UnsafeCell::new(Self::CIPHER),
            core::marker::PhantomData,
        )
    }

    #[inline(always)]
    pub fn resolve(&self) -> usize {
        unsafe {
            if *self.0.get() != Self::CIPHER {
                (*self.0.get() ^ Self::CIPHER) as usize
            } else {
                let ptr = import::<M, I, X, X>();
                *self.0.get() = ptr as u64 ^ Self::CIPHER;
                ptr
            }
        }
    }

}

impl<const M: u64, const I: u64, const X: u64, V> core::ops::Deref for LazyImport<M, I, X, V> {
    type Target = V;

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        unsafe { &*(self.resolve() as *const V) }
    }
}

// ┌─────────────────────────────────────────────────────────────────────────┐
// │            the function actually responsible for everything             │
// └─────────────────────────────────────────────────────────────────────────┘
#[inline(always)]
pub unsafe fn import<
    const Module: u64,
    const Import: u64,
    const XorSeed: u64,
    const XorSeed2: u64,
>() -> usize {
    unsafe {

        // get the ldr data table entries
        let module_link: *const usize;
        core::arch::asm!(
        "mov {x}, gs:[60h]",    // TEB->PEB
        "mov {x}, [{x} + 18h]", // PEB->LDR
        "lea {x}, [{x} + 10h]", // LDR->InLoadOrderModuleList
        x = out(reg) module_link,
        );

        // 0x0 = next, 0x8 = prev, use the xor seed to flip the direction of the iterator
        let offsets = const { XorSeed & 1 } as usize;

        // in case of forwarded imports we will change these...
        let mut module_hash = Module;
        let mut import_hash = Import;

        'search_for_import: loop {

            let mut cursor = module_link as usize;
            let end = (cursor as *const usize).add(offsets ^ 1).read();
            let mut module: usize = 0usize;

            #[cfg(debug_assertions)]
            let mut found = false;

            while cursor != end {
                cursor = (cursor as *const usize).add(offsets).read();

                // extract the appropriate fields
                let name_len = ((cursor + 0x58) as *const u16).read();
                let name_ptr = ((cursor + 0x60) as *const *const u16).read();
                module = ((cursor + 0x30) as *const usize).read();

                // calculate the hash for the module
                let hash = hash_utf16::<XorSeed>(core::slice::from_raw_parts(
                    name_ptr, (name_len >> 1) as usize));

                if hash == module_hash {

                    #[cfg(debug_assertions)]
                    if true {
                        found = true;
                    }
                    break;
                }
            }

            #[cfg(debug_assertions)]
            debug_assert!(found, "Failed to find module");
            let module = &*(module as *const ImageBase);

            match module.exports() {
                None => {
                    debug_assert!(false, "Module has no exports");
                    core::arch::asm!("", options(noreturn));
                }
                Some(exports) => {
                    for (vaddr, name, forwarded) in exports {

                        if hash_utf8::<XorSeed2,1>([name]) == import_hash {

                            // the import is forwarded... oh boy...
                            if forwarded {

                                // strlen and convert to a slice
                                let fwd = {
                                    let ptr = vaddr as *const u8;
                                    let len = (0usize..)
                                        .position(|i| ptr.add(i).read() == 0)
                                        .unwrap_or(0);
                                    core::slice::from_raw_parts(ptr, len)
                                };


                                let Some(split) = fwd.iter().position(|i| *i == b'.') else {
                                    debug_assert!(false, "Forwarded input lacks . separator");
                                    core::arch::asm!("", options(noreturn))
                                };

                                // split our bytes into two slices at the .
                                let (module, import) = (&fwd[..split], &fwd[split+1..]);

                                // to prevent bullshit where people can sig scan for .dll or xref it
                                // I do this hacky encrypt
                                let dll_crypt = const {
                                    let mut p = Prng::new(XorSeed);
                                    p.skip(4);
                                    u32::from_le_bytes(*b".dll") ^ p.next() as u32
                                };

                                // use blackbox so it's not decrypted at compile time
                                let dll = core::hint::black_box(dll_crypt) ^ const {
                                    let mut p = Prng::new(XorSeed);
                                    p.skip(4);
                                    p.next() as u32
                                };

                                let dll = dll.to_le_bytes();

                                module_hash = hash_utf8::<XorSeed,2>([module, &dll]);
                                import_hash = hash_utf8::<XorSeed2, 1>([import]);

                                // restart the loop but with our new hashes now...
                                continue 'search_for_import;
                            }

                            return vaddr;
                        }
                    }
                }
            }
            debug_assert!(false, "function not found");
            core::arch::asm!("", options(noreturn))
        }
    }
}



// ┌─────────────────────────────────────────────────────────────────────────┐
// │                          Helper PRNG functions                          │
// └─────────────────────────────────────────────────────────────────────────┘

pub struct Prng(u64);

impl Prng {
    pub const fn new(seed: u64) -> Self {
        Self(seed)
    }

    pub const fn next(&mut self) -> u64 {
        self.skip(1);
        let result = self.0.wrapping_mul(0x2545F4914F6CDD1D);
        result
    }

    pub const fn size(&mut self, limit: usize) -> usize {
        self.next() as usize % limit
    }

    pub const fn skip(&mut self, mut cycles: usize) {
        let mut state = self.0;
        while cycles > 0 {
            cycles = cycles.saturating_sub(1);
            state ^= state >> 12;
            state ^= state << 25;
            state ^= state >> 27;
        }
        self.0 = state;
    }
}

// ┌─────────────────────────────────────────────────────────────────────────┐
// │                             Hash Algorithm                              │
// └─────────────────────────────────────────────────────────────────────────┘

const INITIAL_STATE: u64 = 0xcbf29ce484222325;
const PRIME: u64 = 0x100000001b3;

//noinspection DuplicatedCode
#[inline(always)]
pub const fn hash_utf8<const SALT: u64, const S: usize>(bytes: [&[u8]; S]) -> u64 {

    // make the initial state based on the salt
    let mut hash = const {
        let mut hash = INITIAL_STATE;
        // inject pre_seed
        let b = SALT.to_le_bytes();
        let mut i = 0;
        while i < 7 {
            hash ^= b[i] as u64;
            hash = hash.wrapping_mul(PRIME);
            i += 1;
        }
        i = 0;
        hash
    };


    let mut j = 0;
    while j < bytes.len() {
        let bytes = bytes[j];
        let mut i = 0;
        while i < bytes.len() {
            let char = match bytes[i] {
                0x40..=0x5A => bytes[i] + 0x20,
                _ => bytes[i],
            } as u64;

            hash = hash ^ (char);
            hash = hash.wrapping_mul(PRIME);
            i += 1;
        }
        j += 1;
    }

    hash
}

//noinspection DuplicatedCode
pub const fn hash_utf16<const SALT: u64>(bytes: &[u16]) -> u64 {

    // make the initial state based on the salt
    let mut hash = const {
        let mut hash = INITIAL_STATE;
        // inject pre_seed
        let b = SALT.to_le_bytes();
        let mut i = 0;
        while i < 7 {
            hash ^= b[i] as u64;
            hash = hash.wrapping_mul(PRIME);
            i += 1;
        }
        i = 0;
        hash
    };

    let mut i = 0;
    while i < bytes.len() {
        let char = match bytes[i] {
            0x40..=0x5A => bytes[i] + 0x20,
            _ => bytes[i],
        } as u64;

        hash = hash ^ (char);
        hash = hash.wrapping_mul(PRIME);
        i += 1;
    }
    hash
}

// ┌─────────────────────────────────────────────────────────────────────────┐
// │                           Pe Data Structures                            │
// └─────────────────────────────────────────────────────────────────────────┘

#[repr(C)]
pub struct ImageBase(());

impl ImageBase {
    pub fn offset<'a, T>(&'a self, offset: T) -> *const u8
    where
        &'a Self: core::ops::Add<T, Output=usize>,
    {
        (self + offset) as *const u8
    }

    pub fn dos(&self) -> &ImageDOSHeader {
        unsafe { &*(self as *const _ as *const ImageDOSHeader) }
    }

    pub fn nt_header(&self) -> &'static ImageNTHeaders64 {
        unsafe { &*((self + self.dos().e_lfanew) as *const ImageNTHeaders64) }
    }
}

impl core::ops::Add<usize> for &ImageBase {
    type Output = usize;
    fn add(self, rhs: usize) -> Self::Output {
        (self as *const _ as usize) + rhs
    }
}

impl core::ops::Add<u32> for &ImageBase {
    type Output = usize;
    fn add(self, rhs: u32) -> Self::Output {
        (self as *const _ as usize) + (rhs as usize)
    }
}

#[repr(C)]
pub struct ImageDOSHeader {
    pub e_magic: u16,
    pub e_cblp: u16,
    pub e_cp: u16,
    pub e_crlc: u16,
    pub e_cparhdr: u16,
    pub e_minalloc: u16,
    pub e_maxalloc: u16,
    pub e_ss: u16,
    pub e_sp: u16,
    pub e_csum: u16,
    pub e_ip: u16,
    pub e_cs: u16,
    pub e_lfarlc: u16,
    pub e_ovno: u16,
    pub e_res: [u16; 4],
    pub e_oemid: u16,
    pub e_oeminfo: u16,
    pub e_res2: [u16; 10],
    pub e_lfanew: u32,
}

#[repr(C)]
pub struct ImageNTHeaders64 {
    pub signature: u32,
    pub file_header: ImageFileHeader,
    pub optional_header: ImageOptionalHeader64,
}

#[repr(C)]
pub struct ImageFileHeader {
    pub machine: u16,
    pub number_of_sections: u16,
    pub timestamp: u32,
    pub pointer_to_symbol_table: u32,
    pub number_of_symbols: u32,
    pub size_of_optional_header: u16,
    pub characteristics: u16,
}

#[repr(C)]
pub struct ImageOptionalHeader64 {
    pub magic: u16,
    pub major_linker_version: u8,
    pub minor_linker_version: u8,
    pub size_of_code: u32,
    pub size_of_initialized_data: u32,
    pub size_of_uninitialized_data: u32,
    pub address_of_entry_point: u32,
    pub base_of_code: u32,
    pub image_base: u64,
    pub section_alignment: u32,
    pub file_alignment: u32,
    pub major_operating_system_version: u16,
    pub minor_operating_system_version: u16,
    pub major_image_version: u16,
    pub minor_image_version: u16,
    pub major_subsystem_version: u16,
    pub minor_subsystem_version: u16,
    pub win32_version_value: u32,
    pub size_of_image: u32,
    pub size_of_headers: u32,
    pub checksum: u32,
    pub subsystem: u16,
    pub dll_characteristics: u16,
    pub size_of_stack_reserve: u64,
    pub size_of_stack_commit: u64,
    pub size_of_heap_reserve: u64,
    pub size_of_heap_commit: u64,
    pub loader_flags: u32,
    pub number_of_rva_and_sizes: u32,
    pub data_directory: [ImageDataDirectory; 16],
}

#[repr(C)]
pub struct ImageDataDirectory {
    pub virtual_address: u32,
    pub size: u32,
}

#[repr(C)]
pub struct ImageExportDirectory {
    /* 0x00 */pub export_flags: u32,
    /* 0x04 */
    pub timestamp: u32,
    /* 0x06 */
    pub major_version: u16,
    /* 0x08 */
    pub minor_version: u16,
    /* 0x0C */
    pub name_rva: u32,
    /* 0x10 */
    pub ordinal_base: u32,
    /* 0x14 */
    pub address_table_entries: u32,
    /* 0x18 */
    pub number_of_name_pointers: u32,
    /* 0x1C */
    pub export_address_table_rva: u32,
    /* 0x20 */
    pub name_pointer_rva: u32,
    /* 0x24 */
    pub ordinal_table_rva: u32,
}


pub struct ExportIter<'a> {
    image: &'a ImageBase,
    export_dir: &'a ImageExportDirectory,
    export_index: usize,
    ext_range: core::ops::Range<usize>,
}

impl<'a> Iterator for ExportIter<'a> {
    type Item = (usize, &'static [u8], bool);

    fn next(&mut self) -> Option<Self::Item> {
        match self.export_index < self.export_dir.number_of_name_pointers as usize {
            true => unsafe {
                #[inline(always)]
                unsafe fn u8_nul_terminated(ptr: *const u8) -> &'static [u8] {
                    unsafe {
                        let mut end = ptr;
                        while *end != 0 { end = end.add(1) }
                        let len = (end as usize) - (ptr as usize);
                        &*core::ptr::slice_from_raw_parts(ptr, len)
                    }
                }

                let export_functions = self.image.offset(self.export_dir.export_address_table_rva) as *const u32;
                let export_names = self.image.offset(self.export_dir.name_pointer_rva) as *const u32;
                let export_ordinals = self.image.offset(self.export_dir.ordinal_table_rva) as *const u16;

                let export_name = self.image.offset(*export_names.add(self.export_index));
                let export_ordinal = *export_ordinals.add(self.export_index);
                let export_rva = *export_functions.add(export_ordinal as usize);
                let export_va = self.image.offset(export_rva) as usize;

                self.export_index += 1;

                let is_forwarded = self.ext_range.contains(&(export_rva as usize));
                Some((export_va, u8_nul_terminated(export_name), is_forwarded))
            }
            false => None,
        }
    }
}

impl ImageBase {
    pub fn exports(&self) -> Option<ExportIter> {
        let directory = &self.nt_header().optional_header.data_directory[0];
        if directory.size == 0 || directory.virtual_address == 0 { return None; }
        let export_directory = unsafe {
            &*((self + directory.virtual_address) as *const ImageExportDirectory)
        };

        let start = directory.virtual_address as usize;
        let end = start + directory.size as usize;

        Some(ExportIter {
            image: self,
            export_dir: export_directory,
            export_index: 0,
            ext_range: start..end,
        })
    }
}