mirror of
https://github.com/notscimmy/libcapcom
synced 2026-06-22 07:30:16 +00:00
113 lines
3.2 KiB
C++
113 lines
3.2 KiB
C++
#pragma once
|
|
|
|
#include <string>
|
|
#define WIN32_NO_STATUS
|
|
#include <Windows.h>
|
|
#include <Winternl.h>
|
|
#undef WIN32_NO_STATUS
|
|
#include <ntstatus.h>
|
|
|
|
namespace native
|
|
{
|
|
extern "C" NTSTATUS NTAPI ZwLoadDriver(PUNICODE_STRING str);
|
|
extern "C" NTSTATUS NTAPI ZwUnloadDriver(PUNICODE_STRING str);
|
|
|
|
namespace internal
|
|
{
|
|
LSTATUS prepare_driver_registry(const wchar_t *svcName, const wchar_t *svcDrv, const wchar_t *group, int startupType)
|
|
{
|
|
HKEY key;
|
|
HKEY subkey;
|
|
|
|
DWORD type = 1;
|
|
DWORD err = 0;
|
|
LSTATUS status = 0;
|
|
|
|
wchar_t path[MAX_PATH];
|
|
swprintf_s(path, ARRAYSIZE(path), L"\\??\\%s", svcDrv);
|
|
|
|
status = RegOpenKeyW(HKEY_LOCAL_MACHINE, L"system\\CurrentControlSet\\Services", &key);
|
|
if (status)
|
|
return status;
|
|
|
|
status = RegCreateKeyW(key, svcName, &subkey);
|
|
if (status)
|
|
{
|
|
RegCloseKey(key);
|
|
return status;
|
|
}
|
|
|
|
status |= RegSetValueExW(subkey, L"DisplayName", 0, REG_SZ, (const BYTE *)svcName, (DWORD)(sizeof(WCHAR) * wcslen(svcName) + 1));
|
|
status |= RegSetValueExW(subkey, L"ErrorControl", 0, REG_DWORD, (const BYTE *)&err, sizeof(err));
|
|
status |= RegSetValueExW(subkey, L"Group", 0, REG_SZ, (const BYTE *)group, sizeof(WCHAR) * ((DWORD)wcslen(group) + 1));
|
|
status |= RegSetValueExW(subkey, L"ImagePath", 0, REG_SZ, (const BYTE *)path, (sizeof(WCHAR) * ((DWORD)wcslen(path) + 1)));
|
|
status |= RegSetValueExW(subkey, L"Start", 0, REG_DWORD, (const BYTE *)&startupType, sizeof(startupType));
|
|
status |= RegSetValueExW(subkey, L"Type", 0, REG_DWORD, (const BYTE*)&type, sizeof(type));
|
|
|
|
RegCloseKey(subkey);
|
|
|
|
return status;
|
|
}
|
|
|
|
std::wstring make_path(const std::wstring& svcName)
|
|
{
|
|
std::wstring path = L"\\registry\\machine\\SYSTEM\\CurrentControlSet\\Services\\";
|
|
path += svcName;
|
|
return path;
|
|
}
|
|
|
|
bool set_privilege(HANDLE hToken, LPCTSTR lpszPrivilege, BOOL bEnablePrivilege)
|
|
{
|
|
TOKEN_PRIVILEGES tp;
|
|
LUID luid;
|
|
|
|
if (!LookupPrivilegeValue(NULL, lpszPrivilege, &luid))
|
|
return false;
|
|
|
|
tp.PrivilegeCount = 1;
|
|
tp.Privileges[0].Luid = luid;
|
|
if (bEnablePrivilege)
|
|
tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
|
|
else
|
|
tp.Privileges[0].Attributes = 0;
|
|
|
|
// Enable the privilege or disable all privileges.
|
|
if (!AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES),
|
|
(PTOKEN_PRIVILEGES)NULL, (PDWORD)NULL))
|
|
return false;
|
|
|
|
return GetLastError() != ERROR_NOT_ALL_ASSIGNED;
|
|
}
|
|
}
|
|
|
|
bool load_driver(const std::wstring& path, const std::wstring& service)
|
|
{
|
|
LSTATUS stat = internal::prepare_driver_registry(service.c_str(), path.c_str(), L"Base", 1);
|
|
if (stat != ERROR_SUCCESS)
|
|
return stat;
|
|
|
|
UNICODE_STRING str;
|
|
auto wpath = internal::make_path(service);
|
|
RtlInitUnicodeString(&str, wpath.c_str());
|
|
|
|
HANDLE token;
|
|
if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &token))
|
|
return false;
|
|
|
|
bool done = internal::set_privilege(token, TEXT("SeLoadDriverPrivilege"), TRUE);
|
|
CloseHandle(token);
|
|
if (!done)
|
|
return false;
|
|
|
|
return ZwLoadDriver(&str) >= 0;
|
|
}
|
|
|
|
bool unload_driver(const std::wstring& service)
|
|
{
|
|
UNICODE_STRING str;
|
|
auto wservice = internal::make_path(service);
|
|
RtlInitUnicodeString(&str, wservice.c_str());
|
|
|
|
return ZwUnloadDriver(&str) >= 0;
|
|
}
|
|
} |