class DriverControl
{
public:
    DriverControl(std::wstring serviceName, std::wstring deviceName, std::wstring binPath)
        : _driverHandle(INVALID_HANDLE_VALUE), _serviceName(std::move(serviceName))
    {
        wchar_t pathBuffer[MAX_PATH];
        _wfullpath(pathBuffer, std::data(binPath), MAX_PATH);

        _driverPath         = pathBuffer;
        _registryKeyName    = L"SYSTEM\\CurrentControlSet\\Services\\" + _serviceName;
        _registryKeyNtName  = L"\\registry\\machine\\SYSTEM\\CurrentControlSet\\Services\\" + _serviceName;
        _symbolicName       = L"\\Device\\" + deviceName;
    }

    ~DriverControl()
    {
        if(IsLoaded()) {
            UnloadDriver();
            _driverHandle = INVALID_HANDLE_VALUE;
        }
    }

    long LoadDriver()
    {
        BOOLEAN wasEnabled;
        UNICODE_STRING ustr;
        RtlInitUnicodeString(&ustr, std::data(_registryKeyNtName));

        auto status = RtlAdjustPrivilege(10L /*SE_LOAD_DRIVER_PRIVILEGE*/, TRUE, FALSE, &wasEnabled);

        if(!NT_SUCCESS(status)) {
            return RtlNtStatusToDosError(status);
        }

        auto hr = CreateRegistryEntry();

        if(FAILED(hr)) {
            return hr;
        }

        status = NtLoadDriver(&ustr);

        if(!NT_SUCCESS(status)) {
            RemoveRegistryEntry();
            return RtlNtStatusToDosError(status);
        }

        return OpenDevice();
    }
    long UnloadDriver()
    {
        UNICODE_STRING ustr;
        RtlInitUnicodeString(&ustr, std::data(_registryKeyNtName));

        auto status = NtUnloadDriver(&ustr);

        if(status < 0) {
            return RtlNtStatusToDosError(status);
        }

        return RemoveRegistryEntry();
    }
    long OpenDevice()
    {
        NTSTATUS status;
        UNICODE_STRING symName;
        OBJECT_ATTRIBUTES obj;
        IO_STATUS_BLOCK ioStatus;
        ULONG tryCount = 0;

        RtlInitUnicodeString(&symName, std::data(_symbolicName));

        InitializeObjectAttributes(&obj, &symName, 0, NULL, NULL);

        do {
            status = NtOpenFile(
                &_driverHandle,
                GENERIC_READ | GENERIC_WRITE,
                &obj,
                &ioStatus,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                OPEN_EXISTING);

            if(NT_SUCCESS(status))
                break;
            Sleep(100);
        } while(tryCount++ < 5);

        return RtlNtStatusToDosError(status);
    }
    bool IsLoaded()
    {
        return _driverHandle != INVALID_HANDLE_VALUE;
    }

    template<typename T, typename... Args>
    BOOL KaDeviceIoControl(DWORD ioctl, Args... args)
    {
        T params = {args...};
        return DeviceIoControl(_driverHandle, ioctl, &params, sizeof(params), nullptr, 0, nullptr, nullptr);
    }

private:
    HRESULT CreateRegistryEntry()
    {
        HRESULT      Status;
        HKEY         KeyService;
        ULONG        DriverType = 1;                     /*Kernel*/
        std::wstring NtPath = L"\\??\\" + _driverPath;   /*Dos -> NT conversion*/

        Status = RegCreateKeyW(HKEY_LOCAL_MACHINE, std::data(_registryKeyName), &KeyService);

        if(FAILED(Status))
            return Status;

        Status = RegSetValueExW(KeyService, L"ImagePath", 0, REG_EXPAND_SZ, reinterpret_cast<LPCBYTE>(std::data(NtPath)), NtPath.size() * sizeof(WCHAR));
        if(FAILED(Status)) goto Cleanup;

        Status = RegSetValueExW(KeyService, L"Type", 0, REG_DWORD, reinterpret_cast<LPCBYTE>(&DriverType), sizeof(DriverType));
        if(FAILED(Status)) goto Cleanup;

    Cleanup:
        RegCloseKey(KeyService);
        return Status;
    }

    HRESULT RemoveRegistryEntry()
    {
        return RegDeleteKeyW(HKEY_LOCAL_MACHINE, std::data(_registryKeyName));
    }
private:
    HANDLE          _driverHandle;
    std::wstring    _serviceName;
    std::wstring    _driverPath;
    std::wstring    _registryKeyName;
    std::wstring    _registryKeyNtName;
    std::wstring    _symbolicName;
};