自己封装了Windows内核String

Posted 2021-09-05 冰翔不败传说

#ifndef __WKSTRING_H__
#define __WKSTRING_H__

#include <ntdef.h>

#define UULogInfo(x,...) WKString::log(__FILE__,__LINE__,__FUNCTION__,x,__VA_ARGS__)

/*
	内核String封装, WCHAR*类型,必须以\\0结尾
*/
class WKString
{
private:
	WCHAR* value_;
	int length_;
	int capacity_;
private:
	void initMem(int size);
	bool ensureCapacity(int minimumCapacity);
public:
	WKString();
	WKString(int capacity);
	WKString(const PUNICODE_STRING str);
	~WKString();	
	WKString& append(const PUNICODE_STRING str);
	WKString& append(const WCHAR* buf);
	WKString& append(const char* buf);
	WKString& append(int num);
	void printf();
	const WCHAR* getBuf();
	const UNICODE_STRING getUnicodeString();
	bool equalsIgnoreCase(const WCHAR* other) const;
	bool equalsIgnoreCase(const PUNICODE_STRING other) const;
public:
	//打印日志 对DbgPrint封装
	static void log(char* fileName, int line, char* function, LPCSTR  fmt, ...);
};

#endif // !__WKSTRING_H__


#include "WKString.h"
#include <ntddk.h>
#include <stdio.h>
#include <stdarg.h>

WKString::WKString()
{
	initMem(0);

}
WKString::WKString(int capacity)
{
	initMem(capacity);
}

WKString::WKString(const PUNICODE_STRING str)
{
	if (str == NULL)
	{
		initMem(0);
	}
	else
	{
		int length = str->Length / sizeof(WCHAR);
		if (length <= 0)
		{
			initMem(0);
		}
		else
		{
			initMem(length + 1);
			append(str);
		}
	}
}

WKString::~WKString()
{
	if (value_ != NULL)
		ExFreePool(value_);
	value_ = NULL;
	capacity_ = 0;
	length_ = 0;
}
void WKString::initMem(int size)
{
	capacity_ = 0;
	length_ = 0;
	value_ = NULL;	
	if (size <= 0)
		return;

	value_ = (WCHAR*)ExAllocatePool(NonPagedPool, size * sizeof(WCHAR));
	if (value_ != NULL)
	{
		RtlZeroMemory(value_, size * sizeof(WCHAR));
		length_ = 0;
		capacity_ = size;
	}	
	else
	{
		UULogInfo("初始化内存失败 size=%d",size);
	}
}

WKString& WKString::append(const PUNICODE_STRING str)
{
	if (str->Buffer == NULL)
		return *this;

	if (str->Length == 0)
		return *this;

	int length = str->Length / sizeof(WCHAR);
	if (!ensureCapacity(length + length_ + 1))
	{
		return *this;
	}
	memcpy((void*)(value_ + length_), (void*)(str->Buffer), str->Length);
	length_ += length;
	return *this;	
}

WKString& WKString::append(const WCHAR* buf)
{
	if (buf == NULL)
		return *this;

	int length = wcslen(buf);
	if (length == 0)
		return *this;
	
	if (!ensureCapacity(length + length_ + 1))
	{
		return *this;
	}
	memcpy((void*)(value_ + length_), (void*)(buf), length * sizeof(WCHAR));
	length_ += length;
	return *this;
}

WKString& WKString::append(const char* buf)
{
	if (buf == NULL)
		return *this;

	int length = strlen(buf);
	if (length == 0)
		return *this;

	ANSI_STRING ansiString;
	RtlInitAnsiString(&ansiString, buf);
	UNICODE_STRING unicodeString;
	NTSTATUS status = RtlAnsiStringToUnicodeString(&unicodeString, &ansiString, true);
	if (status == STATUS_SUCCESS)
	{
		append(&unicodeString);
		RtlFreeUnicodeString(&unicodeString);
	}
	else
	{
		UULogInfo("RtlAnsiStringToUnicodeString转化失败:%s", buf);
	}	
	return *this;
}

WKString& WKString::append(int num)
{
	char buf[20] = { 0 };
	sprintf_s(buf, 20, "%d", num);
	return append(buf);
}

void WKString::printf()
{
	if (value_ != NULL)
	{
		DbgPrint("[WKLog] %ws length=%d capacity=%d\\n", value_,length_,capacity_);
	}
	else
	{
		DbgPrint("[WKLog] value == NULL\\n");		
	}		
}

const WCHAR* WKString::getBuf()
{
	if (value_ == NULL)
		return L"";
	return value_;
}

const UNICODE_STRING WKString::getUnicodeString()
{
	UNICODE_STRING string;
	if (value_ == NULL)
	{
		RtlInitUnicodeString(&string, L"");
	}
	else
	{
		string.Buffer = value_;
		string.Length = length_ * sizeof(WCHAR);
		string.MaximumLength = capacity_ * sizeof(WCHAR);
	}	
	return string;
}



bool WKString::equalsIgnoreCase(const WCHAR* other) const
{
	if (other == NULL)
		return false;
	if (length_ != wcslen(other))
		return false;
	return _wcsnicmp(this->value_, other, length_) == 0;	
}

bool WKString::equalsIgnoreCase(const PUNICODE_STRING other) const
{
	WKString s(other);
	return equalsIgnoreCase(s.getBuf());
}

void WKString::log(char* fileName, int line, char* function, LPCSTR fmt, ...)
{
	va_list	va;
	va_start(va, fmt);

	int bufLen = strlen(fmt) + 50;
	int iLen = bufLen;
	char* pCursor = (char*)fmt;

	for (int i = 0; i < bufLen; i++)
	{
		char c = *pCursor;
		if (c == L'%')
		{
			char c1 = *(pCursor + 1);
			char c2 = -1;
			if (i + 2 < bufLen)
			{
				c2 = *(pCursor + 2);
			}
			if (c1 == L's')
			{
				char* pChar = (char*)va_arg(va, char*);
				int n = strlen(pChar);
				iLen += n;
			}
			else if (c1 == L'S')
			{
				WCHAR* pChar = (WCHAR*)va_arg(va, char*);
				int n = wcslen(pChar) * sizeof(WCHAR);
				//DbgPrint("[WKLog] S %d\\n", n);
				iLen += n;
			}
			else if (c1 == L'w' && c2 == L's')
			{
				WCHAR* pChar = (WCHAR*)va_arg(va, char*);
				int n = wcslen(pChar) * sizeof(WCHAR);
				//DbgPrint("[WKLog] ws %d\\n", n);
				iLen += n;
			}
			else if (c1 == L'w' && c2 == L'Z')
			{
				PUNICODE_STRING pChar = (PUNICODE_STRING)va_arg(va, char*);
				int n = pChar->MaximumLength + 1;
				iLen += n;
				//DbgPrint("[WKLog] wZ %d %d\\n", pChar->Length, pChar->MaximumLength);
			}
			else if (c1 == L'd')
			{
				int n1 = (int)va_arg(va, int);
				iLen += 10;
			}
			else if (c1 == 0)
			{
				//什么都不做，结束
			}
			else
			{
				int n1 = (int)va_arg(va, int);
				iLen += 10;
			}
		}
		pCursor++;
	}

	va_start(va, fmt);
	char* pBuf = (char*)ExAllocatePool(NonPagedPool, iLen);
	if (pBuf == NULL)
	{
		va_end(va);
		return;
	}
	memset(pBuf, 0, sizeof(char) * iLen);
	vsprintf_s(pBuf, iLen, (const char*)fmt, va);
	va_end(va);
	int len = strlen(pBuf);
	DbgPrint("[WKLog] [%s] [%s] [No.%d] [内存=%d,长度=%d]\\n", fileName, function, line,iLen,len);
	DbgPrint("[WKLog] %s\\n",pBuf);
	ExFreePool(pBuf);
	pBuf = NULL;
}

bool WKString::ensureCapacity(int minimumCapacity)
{
	if (capacity_ >= minimumCapacity)
		return true;
	int newCapacity = minimumCapacity;
	
	WCHAR* newBuf = (WCHAR*)ExAllocatePool(NonPagedPool, newCapacity * sizeof(WCHAR));
	if (newBuf == NULL)
	{
		UULogInfo("初始化内存失败 size=%d", newCapacity);
		return false;
	}
	RtlZeroMemory(newBuf, newCapacity * sizeof(WCHAR));	
	memcpy((void*)newBuf, (void*)value_, length_ * sizeof(WCHAR));
	if (value_ != NULL)
		ExFreePool(value_);
	value_ = newBuf;
	capacity_ = newCapacity;
	return true;
}