#include <iostream>
#include <cstdlib>
#include <Windows.h>
#include "../KernelCommunication/ioctl.h"
#include <memory>

#define TEST_SERVICE_NAME "aaayhytestservice"


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() {
    HANDLE hFile = CreateFileW(L"\\\\.\\" MY_DRVICE_NAME,
        GENERIC_READ | GENERIC_WRITE,
        0,
        NULL,
        OPEN_EXISTING,
        0,
        NULL);
    printf("CreateFile hFile:%p ErrCode:%08x\n", hFile, GetLastError());
    // 测试ReadFile 功能
    char Buffer[MAXBYTE] = { 0 };
    DWORD dwBytes = 0;
    BOOL Result = ReadFile(hFile, Buffer, sizeof(Buffer), &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());
    // 测试DeviceIoControl 功能
    char data[4] = { 0x78, 0x56, 0x34, 0x12 };
    printf("Data value: %p\n", *(int*)data);
    DeviceIoControl(hFile, IOCTL_READ_MEMORY, &data, sizeof(data),
        data, sizeof(data), &dwBytes, NULL);
    printf("DeviceIoControl Result:%d dwBytes:%d Data:%s ErrCode:%08x\n", Result, dwBytes, data, GetLastError());
    printf("Data newValue: %p\n", *(int*)data);

    // 关闭文件
    CloseHandle(hFile);
    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);
//}