#ifndef _WIN64
#error x64 builds supported only!
#endif

#include <Windows.h>
#include <map>
#include <string>
#include <memory>
#include <thread>
#include "scanner.h"

bool WinDefScanner::m_isAPIloaded = false;

namespace windefapi
{
	MpManagerOpen_t*		MpManagerOpen;
	MpHandleClose_t*		MpHandleClose;
	MpErrorMessageFormat_t* MpErrorMessageFormat;
	MpFreeMemory_t*			MpFreeMemory;
	MpThreatOpen_t*			MpThreatOpen;
	MpThreatEnumerate_t*	MpThreatEnumerate;
	MpScanStartEx_t*		MpScanStartEx;
	WDStatus_t*				MpWDStatus;
	MpScanResult_t*			MpScanResult;
	MpScanControl_t*		MpScanControl;
}

using namespace windefapi;

static std::wstring AnsiToWstring(const std::string& input, DWORD locale /*= CP_ACP*/)
{
	wchar_t buf[8192] = { 0 };
	MultiByteToWideChar(locale, 0, input.c_str(), (int) input.length(), buf, ARRAYSIZE(buf));
	return buf;
}

static std::string WstringToAnsi(const std::wstring& input, DWORD locale /*= CP_ACP*/)
{
	char buf[8192] = { 0 };
	WideCharToMultiByte(locale, 0, input.c_str(), (int) input.length(), buf, ARRAYSIZE(buf), nullptr, nullptr);
	return buf;
}

static void event_handler(void *pUserData, MPCALLBACK_DATA *pData)
{
	/*std::wcout << L"\t--- callback invoked ---\n";
	std::wcout << L" pData->Type= " << pData->Type << "\n";
	std::wcout << L" pData->Notify= " << pData->Notify << "\n";
	std::wcout << L" unk0= " << (ULONG64)unk0 << "\n";*/

	HANDLE *pEventStop = (HANDLE*) pUserData;

	switch (pData->Notify)
	{
	/*
	case MPNOTIFY::MPNOTIFY_SCAN_INFECTED:
			std::wcout << L"Callback: MPNOTIFY_SCAN_INFECTED!\n";
		break;

		case MPNOTIFY::MPNOTIFY_SCAN_START:
			std::wcout << L"Callback: MPNOTIFY_SCAN_START\n";
		break;
	*/
	case MPNOTIFY::MPNOTIFY_SCAN_COMPLETE:
	{
	/*
		std::wcout << L"------------------------\n";
		std::wcout << L"Callback: MPNOTIFY_SCAN_COMPLETE\n";
		std::wcout << L" ThreatCount: " << ((MPSCAN_DATA*) pData->Data.pEngineData)->ThreatStats.ThreatCount << "\n";
		std::wcout << L" SuspiciousThreatCount: " << ((MPSCAN_DATA*) pData->Data.pEngineData)->ThreatStats.SuspiciousThreatCount << "\n";
		std::wcout << L" Timestamp: " << pData->TimeStamp.LowPart << "\n\n";
		std::wcout << L">>> scan complete <<< " << "\n";
		*/

		SetEvent(*pEventStop);
		break;
	}

	case 8193: 	// occurs when no file at the resource path/scan abort
	{
	//	std::wcout << L"Callback: possible no file at the target path/no access/scan aborted!\n";
		SetEvent(*pEventStop);
		break;
	}

	//	default:
	//	std::wcout << L"[Callback invoked with unhandled reason]\n";
	}

}

bool WinDefScanner::load_api()
{
	if (WinDefScanner::m_isAPIloaded)
		return true;

	HMODULE winDefLibHandle = LoadLibraryA(WINDEF_CLIENT_DLL);

	MpManagerOpen = (MpManagerOpen_t*) GetProcAddress(winDefLibHandle, "MpManagerOpen");
	MpHandleClose = (MpHandleClose_t*) GetProcAddress(winDefLibHandle, "MpHandleClose");
	MpErrorMessageFormat = (MpErrorMessageFormat_t*) GetProcAddress(winDefLibHandle, "MpErrorMessageFormat");
	MpFreeMemory = (MpFreeMemory_t*) GetProcAddress(winDefLibHandle, "MpFreeMemory");
	MpScanStartEx = (MpScanStartEx_t*) GetProcAddress(winDefLibHandle, "MpScanStartEx");
	MpThreatOpen = (MpThreatOpen_t*) GetProcAddress(winDefLibHandle, "MpThreatOpen");
	MpThreatEnumerate = (MpThreatEnumerate_t*) GetProcAddress(winDefLibHandle, "MpThreatEnumerate");
	MpWDStatus = (WDStatus_t*) GetProcAddress(winDefLibHandle, "WDStatus");
	MpScanResult = (MpScanResult_t*) GetProcAddress(winDefLibHandle, "MpScanResult");
	MpScanControl = (MpScanControl_t*) GetProcAddress(winDefLibHandle, "MpScanControl");

	if (MpManagerOpen &&
		MpHandleClose &&
		MpErrorMessageFormat &&
		MpFreeMemory &&
		MpScanStartEx &&
		MpThreatOpen &&
		MpThreatEnumerate &&
		MpWDStatus &&
		MpScanResult &&
		MpScanControl)
	{
		WinDefScanner::m_isAPIloaded = true;
		return true;
	}

	return false;
}

bool WinDefScanner::execute()
{
	if (!load_api())
		return false;

	return true;
}

void WinDefScanner::stop()
{
	SetEvent(m_hEvtExecutionRestricted);
}

void WinDefScanner::scan_folder(const std::shared_ptr<IScanObj>& scanObj)
{
	ResetEvent(m_hEvtExecutionRestricted);
	std::thread t(&WinDefScanner::proceed_single, this, scanObj);
	t.detach();
}

ESCANSTAT WinDefScanner::proceed_single(std::shared_ptr<IScanObj>& item) //thread
{
	if (item->is_compleat())
		return ESCANSTAT::SCANNEDPREVIOSLY;

	MPHANDLE managerHandle = NULL;
	MPHANDLE scanHandle = NULL;
	HRESULT hres = NULL;
	MPRESOURCE_INFO resInfo = { 0 };
	MPSCAN_RESOURCES res;

	hres = MpManagerOpen(0, &managerHandle);

	if (FAILED(hres)) return ESCANSTAT::E_CANTOPENMANAGER;

	resInfo.Class = MP_RESOURCE_CLASS::MP_RESOURCE_CLASS_UNKNOWN;
	resInfo.Scheme = L"folder";

	wchar_t path[256];
	wsprintfW(path, AnsiToWstring(item->folderForScan(), CP_ACP).c_str());
	resInfo.Path = path;

	res.dwResourceCount = 1;
	res.pResourceList = &resInfo;

	HANDLE hCurrentComplete = CreateEvent(nullptr, TRUE, FALSE, nullptr);
	char custom_data[256];

	memset(custom_data, 0, 256);
	memcpy(custom_data, &hCurrentComplete, sizeof(HANDLE));

	struct {
		void *cbHandler;
		void *pUserData;
	}param = { &event_handler, custom_data };

	const DWORD MpCmdRunOpts = 0x02001031; // set minor bit for async call to MpScanStartEx

	size_t startTicks = GetTickCount();

	hres = MpScanStartEx(
		managerHandle,
		MPSCAN_TYPE::MPSCAN_TYPE_RESOURCE,
		MpCmdRunOpts,
		0,
		&res,
		(PMPCALLBACK_INFO) &param,
		&scanHandle
	);

	if (FAILED(hres))
	{
		CloseHandle(hCurrentComplete);
		MpHandleClose(scanHandle);
		MpHandleClose(managerHandle);
		return ESCANSTAT::E_SCANSTARTEXFAILED;
	}

	HANDLE h_run_conditions []{ hCurrentComplete , m_hEvtExecutionRestricted };

	DWORD waitstate = WaitForMultipleObjects(2, h_run_conditions, FALSE, INFINITE);

	// if waitstate == 1, then it means that m_hEvtExecutionRestricted
	//	has been triggered (event index in h_run_conditions array)
	if (waitstate == 1)
	{
		CloseHandle(hCurrentComplete);
		MpHandleClose(scanHandle);
		MpHandleClose(managerHandle);
		return ESCANSTAT::E_SCANFORCETERMINATED;
	}

	BYTE arg1[0x90] = { 0 };
	hres = MpScanResult(scanHandle, arg1); // mpclient!MpScanResult(__int64 a1, void *a2): memset_0(v2, 0, 0x90ui64);

	if (FAILED(hres)) {
		CloseHandle(hCurrentComplete);
		MpHandleClose(scanHandle);
		MpHandleClose(managerHandle);
		return ESCANSTAT::E_SCANRESULTFAILED;
	}

	auto store_scan_data = [scanHandle, &item]() -> bool
	{
		MPHANDLE hThreatEnumHandle;
		HRESULT rc;

		rc = MpThreatOpen(scanHandle, nullptr, nullptr, &hThreatEnumHandle);

		if (FAILED(rc)) return false;

		PMPTHREAT_INFO pInfo;
		std::map<std::string, std::string> m_rc;

		while (MpThreatEnumerate(hThreatEnumHandle, &pInfo) == S_OK)
		{
			for (size_t i = 0; i < pInfo->ResourceCount; i++)
			{
				m_rc[WstringToAnsi(std::wstring(pInfo->ResourceList[i]->Path), CP_ACP)] =
					WstringToAnsi(std::wstring(pInfo->Name), CP_ACP);
			}
		}
		rc = MpHandleClose(hThreatEnumHandle);
		item->submitResult(m_rc);
		return true;
	};

	bool b_dstored = store_scan_data();

	CloseHandle(hCurrentComplete);
	MpHandleClose(scanHandle);
	MpHandleClose(managerHandle);

	item->setElapsedTime(GetTickCount() - startTicks);
	item->finish();

	return b_dstored ? ESCANSTAT::SUCCEEDED : ESCANSTAT::FAILURE;
}