#include "TestDrvice.h"

#define TEST_SERVICE_NAME "aaayhytestservice"
std::thread gWorkThread;
bool gThreadFlags = false;

void workThread(HANDLE eventHandle) {
    while (gThreadFlags) {
        WaitForSingleObject(eventHandle, INFINITE);
        // TODO: 处理内核通知
        printf("Kernel Event!\n");
        ResetEvent(eventHandle); // 重置事件以便下次触发
    }
    // 退出线程时，关闭句柄
    printf("Exit WorkThread! eventHandle: %08x", eventHandle);
    if (eventHandle != NULL) { CloseHandle(eventHandle); }
    return;
}

void ShowError(const char* szName) {
    DWORD dwErrorCode = GetLastError();
    LPVOID lpMsgBuf;
    FormatMessage(
        FORMAT_MESSAGE_ALLOCATE_BUFFER |
        FORMAT_MESSAGE_FROM_SYSTEM |
        FORMAT_MESSAGE_IGNORE_INSERTS,
        NULL,
        dwErrorCode,
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
        (LPTSTR)&lpMsgBuf,
        0,
        NULL
    );
    printf("[%-15s] ErrorCode:%08x ErrorMsg:%s", szName, dwErrorCode, (char*)lpMsgBuf);
    LocalFree(lpMsgBuf);
}

int LoadDriver(const char* szBinPath) {
    using namespace std;
    int Result = -1;

    SC_HANDLE schSCManager = OpenSCManager(
        NULL,                    // local computer
        NULL,                    // ServicesActive database 
        SC_MANAGER_ALL_ACCESS);  // full access rights 
    if (NULL == schSCManager) {
        ShowError("OpenSCManager");
        return Result;
    }
    cout << "[+] OpenSCManager Success!" << endl;

    //创建服务
    SC_HANDLE schService = CreateService(
        schSCManager,              // SCM database 
        TEST_SERVICE_NAME,             // name of service 
        TEST_SERVICE_NAME,             // service name to display 
        SERVICE_ALL_ACCESS,        // desired access 
        SERVICE_KERNEL_DRIVER,     // service type 
        SERVICE_DEMAND_START,      // start type 
        SERVICE_ERROR_NORMAL,      // error control type 
        szBinPath,                 // path to service's binary 
        NULL,                      // no load ordering group 
        NULL,                      // no tag identifier 
        NULL,                      // no dependencies 
        NULL,                      // LocalSystem account 
        NULL);                     // no password 
    if (schService == NULL) {
        if (GetLastError() != ERROR_SERVICE_EXISTS) {
            ShowError("CreateService");
            CloseServiceHandle(schSCManager);
            return Result;
        }

        schService = OpenService(schSCManager, TEST_SERVICE_NAME, SERVICE_ALL_ACCESS);
        if (schService == NULL) {
            ShowError("OpenService");
            CloseServiceHandle(schSCManager);
            return Result;
        }
    }
    cout << "[+] CreateService Success!" << endl;

    SERVICE_STATUS Status;
    ControlService(schService, SERVICE_CONTROL_STOP, &Status);

    //启动服务
    if (!StartService(schService, NULL, NULL)) {
        ShowError("StartService");
    }
    else {
        cout << "[+] StartService Success!" << endl;
        Result = 0;
    }

    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);
    return Result;
}
// 停止并卸载服务
int UnLoadDriver(const char* szServiceName) {
    int Result = -1;

    SC_HANDLE schSCManager = OpenSCManager(
        NULL,                    // local computer
        NULL,                    // ServicesActive database 
        SC_MANAGER_ALL_ACCESS);  // full access rights 
    if (NULL == schSCManager) {
        ShowError("OpenSCManager");
        return Result;
    }

    //打开服务
    SC_HANDLE schService = OpenService(schSCManager, szServiceName, SERVICE_ALL_ACCESS);
    if (schService == NULL) {
        ShowError("OpenService");
        CloseServiceHandle(schSCManager);
        return Result;
    }

    //停止服务
    SERVICE_STATUS Status;
    ControlService(schService, SERVICE_CONTROL_STOP, &Status);

    //删除服务
    if (!DeleteService(schService)) {
        ShowError("DeleteService");
    }
    else {
        Result = 0;
    }

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


void UseDriver() {
    unique_handle hFile = make_unique_handle(CreateFileW(KERNEL_FILE_NAME,
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        OPEN_EXISTING,
        0,
        NULL));
    printf("CreateFile hFile:%p ErrCode:%08x\n", hFile.get(), GetLastError());
    DWORD dwBytes = 0;
    // 1. 发送让驱动创建事件对象的消息
    bool result = DeviceIoControl(hFile.get(), IOCTL_KERNEL_CREATE_EVENT_OBJECT, NULL, 0, NULL, 0, &dwBytes, NULL);
    if (!result) {
        printf("IOCTL_KERNEL_CREATE_EVENT_OBJECT False!\n");
        return;
    }

    // 2. 打开内核驱动的命名同步对象
    HANDLE eventHandle = OpenEventW(EVENT_ALL_ACCESS, FALSE, KERNEL_TO_USER_EVENT_NAME);
    printf("eventHandle: %08x\n", eventHandle);
    gWorkThread = std::thread([=] {
        gThreadFlags = true;
        workThread(eventHandle);
        });


    // 测试ReadFile 功能
    char Buffer[MAXBYTE] = { 0 };
    BOOL Result = ReadFile(hFile.get(), Buffer, 6, &dwBytes, NULL);
    printf("ReadFile Result:%d dwBytes:%d Buffer:%s ErrCode:%08x\n", Result, dwBytes, Buffer, GetLastError());

    //// 测试WriteFile 功能
    //WriteFile(hFile, Buffer, sizeof(Buffer), &dwBytes, NULL);
    //printf("WriteFile Result:%d dwBytes:%d Buffer:%s ErrCode:%08x\n", Result, dwBytes, Buffer, GetLastError());


    // 退出
    gThreadFlags = false;
    result = DeviceIoControl(hFile.get(), IOCTL_CLOSE_EVENT_OBJECT, NULL, 0, NULL, 0, &dwBytes, NULL);
    gWorkThread.join();

    return;
}

int main(int argc, char* argv[])
{
    using namespace std;
    if (argc < 2) {
        std::cerr << "Please drag the driver onto this program icon to test the installation!" << std::endl;
        system("pause");
        return -1;
    }
    const char* driverPath = argv[1];

    cout << "[+] Loading the test driver：" << driverPath << endl;
    if (LoadDriver(driverPath) != 0) {
        cerr << "[x] LoadDriver Error!" << endl;
        system("pause");
        return -1;
    }
    cout << "======== LoadDriver done!   ===========" << endl;
    // TODO: 做一些事情
    UseDriver();

    if (UnLoadDriver(TEST_SERVICE_NAME) != 0) {
        cerr << "[x] UnLoadDriver Faile! " << endl;
        system("pause");
        return -1;
    }
    cout << "======== UnLoadDriver done! ===========" << endl;
    cout << "======== Bye!               ===========" << endl;
    system("pause");
    return 0;
}


//int main() {
//    char data[] = { 0x12, 0x34, 0x56, 0x78 };
//    printf("%p", *(int*)data);
//}