WinDriverCode/Registry.c at master · Vinx911/WinDriverCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
#include "Registry.h"

NTSTATUS Reg_OpenKey(OUT PHANDLE phKey, IN LPCWSTR pszValueName)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING RegistryPath;
	OBJECT_ATTRIBUTES objectAttributes;
	RtlInitUnicodeString(&RegistryPath, pszValueName);
	InitializeObjectAttributes(&objectAttributes, &RegistryPath, OBJ_CASE_INSENSITIVE, NULL, NULL);
	status = ZwOpenKey(phKey, KEY_ALL_ACCESS, &objectAttributes);
	return status;
}

NTSTATUS Reg_GetDwordValue(IN HANDLE hKey, IN LPCWSTR pszValueName, OUT PULONG32 pValue)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	struct {
		KEY_VALUE_PARTIAL_INFORMATION Info;
		ULONG32 Extra;
	} Buffer;
	ULONG ResultLength;
	status = ZwQueryValueKey(hKey, &ValueName, KeyValuePartialInformation, &Buffer, sizeof(Buffer), &ResultLength);
	if (NT_SUCCESS(status))
	{
		if (REG_DWORD == Buffer.Info.Type)
		{
			ASSERT(Buffer.Info.DataLength == sizeof(ULONG32));
			*pValue = *(PULONG32)Buffer.Info.Data;
		}
		else
		{
			status = STATUS_INVALID_BUFFER_SIZE;
		}
	}
	return status;
}

NTSTATUS Reg_SetDwordValue(IN HANDLE hKey, IN LPCWSTR pszValueName, IN ULONG32 Value)
{
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	return ZwSetValueKey(hKey, &ValueName, 0, REG_DWORD, &Value, sizeof(ULONG32));
}

NTSTATUS Reg_GetQwordValue(IN HANDLE hKey, IN LPCWSTR pszValueName, OUT PULONG64 pValue)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	struct {
		KEY_VALUE_PARTIAL_INFORMATION Info;
		ULONG64 Extra;
	} Buffer;
	ULONG ResultLength;
	status = ZwQueryValueKey(hKey, &ValueName, KeyValuePartialInformation, &Buffer, sizeof(Buffer), &ResultLength);
	if (NT_SUCCESS(status))
	{
		if (REG_QWORD == Buffer.Info.Type) {
			ASSERT(Buffer.Info.DataLength == sizeof(ULONG64));
			*pValue = *(PULONG64)Buffer.Info.Data;
		}
		else
			status = STATUS_INVALID_BUFFER_SIZE;
	}
	return status;
}

NTSTATUS Reg_SetQwordValue(IN HANDLE hKey, IN LPCWSTR pszValueName, IN ULONG64 Value)
{
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	return ZwSetValueKey(hKey, &ValueName, 0, REG_QWORD, &Value, sizeof(ULONG64));
}

NTSTATUS Reg_GetStringValue(IN HANDLE hKey, IN LPCWSTR pszValueName, OUT PWSTR *ppValue)
{
	NTSTATUS status = STATUS_SUCCESS;
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	ULONG ResultLength = 0;
	status = ZwQueryValueKey(hKey, &ValueName, KeyValuePartialInformation, NULL, 0, &ResultLength);

	if (status == STATUS_OBJECT_NAME_NOT_FOUND || ResultLength == 0)
	{
		KdPrint(("注册表键值不存在！\n"));
		return CUSTOM_STATUS_REGISTRY_KEY_NOTEXIST;
	}
	KEY_VALUE_PARTIAL_INFORMATION *pBuffer =
		(PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePool(PagedPool, ResultLength);
	if (pBuffer)
	{
		status = ZwQueryValueKey(hKey, &ValueName, KeyValuePartialInformation, pBuffer, ResultLength, &ResultLength);
		if (pBuffer->Type != REG_SZ)
		{
			KdPrint(("注册表键值类型不匹配！\n"));
			return CUSTOM_STATUS_REGISTRY_KEYVALUE_TYPEMISMATCH;
		}
		*ppValue = (PWSTR)ExAllocatePool(PagedPool, pBuffer->DataLength);
		if (*ppValue)
		{
			RtlMoveMemory(*ppValue, pBuffer->Data, pBuffer->DataLength);
		}
		else
		{
			return STATUS_INSUFFICIENT_RESOURCES;
		}
		ExFreePool(pBuffer);
	}

	return STATUS_SUCCESS;
}

NTSTATUS Reg_SetStringValue(IN HANDLE hKey, IN LPCWSTR pszValueName, IN LPCWSTR pszValue)
{
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	return ZwSetValueKey(hKey, &ValueName, 0, REG_SZ, (PVOID)pszValue, wcslen(pszValue) * 2 + 2);
}

NTSTATUS Reg_GetBinaryValue(IN HANDLE hKey, IN LPCWSTR pszValueName, OUT PVOID *ppValue, OUT PULONG ValueLength)
{
	NTSTATUS status = STATUS_SUCCESS;
	ULONG ResultLength = 0;
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	status = ZwQueryValueKey(hKey, &ValueName, KeyValuePartialInformation, NULL, 0, &ResultLength);

	if (status == STATUS_OBJECT_NAME_NOT_FOUND || ResultLength == 0)
	{
		KdPrint(("注册表键值不存在！\n"));
		return CUSTOM_STATUS_REGISTRY_KEY_NOTEXIST;
	}

	KEY_VALUE_PARTIAL_INFORMATION *pBuffer =
		(PKEY_VALUE_PARTIAL_INFORMATION)ExAllocatePool(PagedPool, ResultLength);
	if (pBuffer)
	{
		status = ZwQueryValueKey(hKey, &ValueName, KeyValuePartialInformation, pBuffer, ResultLength, &ResultLength);
		if (pBuffer->Type != REG_BINARY)
		{
			KdPrint(("注册表键值类型不匹配！\n"));
			return CUSTOM_STATUS_REGISTRY_KEYVALUE_TYPEMISMATCH;
		}
		*ppValue = (PWSTR)ExAllocatePool(PagedPool, pBuffer->DataLength);
		if (*ppValue)
		{
			RtlMoveMemory(*ppValue, pBuffer->Data, pBuffer->DataLength);
			*ValueLength = pBuffer->DataLength;
		}
		else
		{
			return STATUS_INSUFFICIENT_RESOURCES;
		}
		ExFreePool(pBuffer);
	}
	return STATUS_SUCCESS;
}

NTSTATUS Reg_SetBinaryValue(IN HANDLE hKey, IN LPCWSTR pszValueName, IN PVOID pValue, IN ULONG ValueLength)
{
	UNICODE_STRING ValueName;
	RtlInitUnicodeString(&ValueName, pszValueName);
	return ZwSetValueKey(hKey, &ValueName, 0, REG_BINARY, pValue, ValueLength);
}