//------------------------------------------------------------------------
// PtInRect Logger
//
// * User-friendly GUI that logs addresses that call PtInRect and allows 
//   the user to filter out specific calls.
// * Hopefully thread-safe
//                                                   Coded by Franc[e]sco
//------------------------------------------------------------------------
// 04/02/2012 - rev1
// [*] First Revision
//------------------------------------------------------------------------

#include "GUI.h"
#include "resource.h"
#include "PtInRectLog.hpp"
#include <CommCtrl.h>
#include <algorithm>
#include <list>
#include <sstream>

#pragma region STRUCTS_AND_MACROS
struct CALL_ITEM
{
	LPVOID lpvAddress;
	int index;
};

extern HWND hDialog;
extern BOOL PIRLOGAPI SendDlgItemMessage(UINT uItem, UINT uMessage, WPARAM wParam, LPARAM lParam);
#define PERROR(szMessage) MessageBoxW(hDialog ? hDialog : NULL, szMessage, L"PtInRectLog", 0);
#define PDIE() TerminateProcess(GetCurrentProcess(), 0);
BOOL IS_CHECKED(UINT uCheckBox) { return (SendDlgItemMessage(uCheckBox, BM_GETCHECK, 0, 0) == BST_CHECKED); };

DWORD wshextoi(LPWSTR lpsz)
{
	DWORD res = 0;
	wstringstream ss(lpsz);
	ss << hex << lpsz;
	ss >> res;
	return res;
}
#pragma endregion

HWND hDialog = NULL;
HINSTANCE hInstance = NULL;
list<CALL_ITEM *> lvItems;
HANDLE Mutex = NULL;
extern CPtInRectLog *pLogger;

LPCWSTR LV_COLUMN_TEXT = L"Return Address";
LPCWSTR APP_VERSION = L"PtInRectLog rev1 (04/02/2012)\nCoded by Franc[e]sco";

BOOL PIRLOGAPI SendMessage(UINT uMessage, WPARAM wParam, LPARAM lParam)
{
	return SendMessage(hDialog, uMessage, wParam, lParam);
}

BOOL PIRLOGAPI SendDlgItemMessage(UINT uItem, UINT uMessage, WPARAM wParam, LPARAM lParam)
{
	return SendDlgItemMessage(hDialog, uItem, uMessage, wParam, lParam);
}

VOID PIRLOGAPI OnCreate(HWND hDialog)
{
	HFONT hFont = NULL;
	HICON hIcon = NULL;
	LVCOLUMN Column = {0};

	::hDialog = hDialog;

	Column.pszText = const_cast<LPWSTR>(LV_COLUMN_TEXT);
	Column.mask = LVCF_TEXT;
	Column.fmt = LVCFMT_LEFT;
	Column.cx = -1;
	Column.cchTextMax = wcslen(LV_COLUMN_TEXT);

	hIcon = LoadIcon(hInstance, MAKEINTRESOURCE(IDI_ICON1));
	hFont = CreateFont(14, 0, 0, 0, FW_DONTCARE, FALSE, FALSE, FALSE, ANSI_CHARSET, 
		OUT_TT_PRECIS, CLIP_DEFAULT_PRECIS, CLEARTYPE_QUALITY, DEFAULT_PITCH, L"Consolas");

	SendMessage(WM_SETICON, ICON_BIG, reinterpret_cast<LPARAM>(hIcon));
	SendMessage(WM_SETICON, ICON_SMALL, reinterpret_cast<LPARAM>(hIcon));
	SendDlgItemMessage(IDC_LIST1, WM_SETFONT, reinterpret_cast<WPARAM>(hFont), TRUE);
	SendDlgItemMessage(IDC_LIST2, WM_SETFONT, reinterpret_cast<WPARAM>(hFont), TRUE);
	SendDlgItemMessage(IDC_CHECK1, BM_SETCHECK, BST_CHECKED, 0);
	SendDlgItemMessage(IDC_CHECK2, BM_SETCHECK, BST_CHECKED, 0);
	SendDlgItemMessage(IDC_CHECK3, BM_SETCHECK, BST_CHECKED, 0);
	SendDlgItemMessage(IDC_LIST1, LVM_SETEXTENDEDLISTVIEWSTYLE, 0, LVS_EX_FULLROWSELECT);
	SendDlgItemMessage(IDC_LIST1, LVM_INSERTCOLUMN, 0, reinterpret_cast<LPARAM>(&Column));
	SendDlgItemMessage(IDC_LIST2, LVM_SETEXTENDEDLISTVIEWSTYLE, 0, LVS_EX_FULLROWSELECT);
	SendDlgItemMessage(IDC_LIST2, LVM_INSERTCOLUMN, 0, reinterpret_cast<LPARAM>(&Column));

	CPtInRectLog::Hook();
	ShowWindow(hDialog, SW_SHOW);
}

VOID PIRLOGAPI OnClose()
{
	int mbResult = MessageBox(hDialog, L"Do you also want to terminate the application" 
		L" this was injected to?", L"PtInRectLog", MB_YESNO | MB_ICONINFORMATION);

	DestroyGUI();

	if (mbResult == IDYES)
		PDIE();
}

VOID PIRLOGAPI About()
{
	MessageBox(hDialog, APP_VERSION, L"PtInRectLog", MB_OK | MB_ICONINFORMATION);
}

VOID PIRLOGAPI AddToIgnoreList()
{
	WCHAR szItemText[MAX_PATH] = {0};
	LVITEM lvItem = {0};
	LVITEM lvDeselectItem = {0};
	lvDeselectItem.stateMask = LVIF_STATE;
	lvDeselectItem.state = 0;
	INT iSelected;

	if (SendDlgItemMessage(IDC_LIST2, LVM_GETNEXTITEM, -1, LVNI_SELECTED) < 0)
		return;

	WaitForSingleObject(Mutex, INFINITE);
	do
	{
		iSelected = SendDlgItemMessage(IDC_LIST2, LVM_GETNEXTITEM, -1, LVNI_SELECTED);
		
		if (iSelected >= 0)
		{
			list<CALL_ITEM *>::iterator it;
			SendDlgItemMessage(IDC_LIST2, LVM_SETITEMSTATE, iSelected, 
				reinterpret_cast<LPARAM>(&lvDeselectItem));
			SendDlgItemMessage(IDC_LIST2, LVM_DELETEITEM, iSelected, 0);

			for (it = lvItems.begin(); it != lvItems.end() && 
				(*it)->index != iSelected; it++);

			if (it != lvItems.end())
				pLogger->Ignore((*it)->lpvAddress);
			else
				continue;

			for each(CALL_ITEM *Call in lvItems)
				if (Call->index > iSelected)
					Call->index--;

			wsprintf(szItemText, L"0x%.8X", (*it)->lpvAddress);
			lvItem.pszText = szItemText;
			lvItem.mask = LVIF_TEXT;
			lvItem.iItem = 0x7FFFFFFF;
			SendDlgItemMessage(IDC_LIST1, LVM_INSERTITEM, 0, reinterpret_cast<LPARAM>(&lvItem));

			delete *it;
			lvItems.erase(it);
		}
	}
	while (iSelected >= 0);
	ReleaseMutex(Mutex);
}

VOID PIRLOGAPI RemoveFromIgnoreList()
{
	INT iSelected;
	DWORD dwAddress = 0;
	WCHAR szItemText[MAX_PATH] = {0};
	LVITEM lvItem = {0};

	WaitForSingleObject(Mutex, INFINITE);
	do
	{
		iSelected = SendDlgItemMessage(IDC_LIST1, LVM_GETNEXTITEM, -1, LVNI_SELECTED);
		
		if (iSelected >= 0)
		{
			lvItem.pszText = szItemText;
			lvItem.cchTextMax = MAX_PATH;
			SendDlgItemMessage(IDC_LIST1, LVM_GETITEMTEXT, 
				iSelected, reinterpret_cast<LPARAM>(&lvItem));
			dwAddress = wshextoi(szItemText+2);
			SendDlgItemMessage(IDC_LIST1, LVM_DELETEITEM, iSelected, 0);
			pLogger->UnIgnore(reinterpret_cast<LPVOID>(dwAddress));
		}
	}
	while (iSelected >= 0);
	ReleaseMutex(Mutex);
}

VOID PIRLOGAPI ClearLogList()
{
	WaitForSingleObject(Mutex, INFINITE);
	SendDlgItemMessage(IDC_LIST2, LVM_DELETEALLITEMS, 0, 0);
	lvItems.clear();
	pLogger->ResetCallCount();
	ReleaseMutex(Mutex);
}

VOID PIRLOGAPI DispatchCall(LPVOID lpvReturnAddress)
{
	WCHAR szItemText[MAX_PATH] = {0};
	LVITEM lvItem = {0};
	list<CALL_ITEM *>::iterator it;

	WaitForSingleObject(Mutex, INFINITE);
	wsprintf(szItemText, L"0x%.8X (%d)", lpvReturnAddress, pLogger->GetCallCount(lpvReturnAddress));
	lvItem.pszText = szItemText;
	lvItem.mask = LVIF_TEXT;
	lvItem.iItem = 0x7FFFFFFF;
	
	for (it = lvItems.begin(); it != lvItems.end() && (*it)->lpvAddress != lpvReturnAddress; it++);

	if (it != lvItems.end())
		SendDlgItemMessage(IDC_LIST2, LVM_SETITEMTEXT, (*it)->index, reinterpret_cast<LPARAM>(&lvItem));
	else
	{
		INT Index;
		CALL_ITEM *pCallItem = new CALL_ITEM;
		pCallItem->lpvAddress = lpvReturnAddress;
		
		Index = SendDlgItemMessage(IDC_LIST2, LVM_INSERTITEM, 0, reinterpret_cast<LPARAM>(&lvItem));
		pCallItem->index = Index;
		lvItems.push_back(pCallItem);
	}
	ReleaseMutex(Mutex);
	Sleep(1);
}

BOOL CALLBACK DlgProc(HWND hWnd, UINT Message, WPARAM wParam, LPARAM lParam)
{
	switch(Message)
	{
#pragma region SYSTEM_MESSAGES
	case WM_INITDIALOG:
		OnCreate(hWnd);
		return TRUE;

	case WM_CLOSE:
		OnClose();
		return TRUE;

	case WM_DESTROY:
		PostQuitMessage(0);
		return TRUE;
#pragma endregion

#pragma region CONTROLS
	case WM_COMMAND:
		switch(LOWORD(wParam))
		{
#pragma region BUTTONS
		case IDC_BUTTON1:
			AddToIgnoreList();
			return TRUE;

		case IDC_BUTTON2:
			RemoveFromIgnoreList();
			return TRUE;

		case IDC_BUTTON3:
			ClearLogList();
			return TRUE;
#pragma endregion

#pragma region MENU
		case ID_FILE_EXIT:
			SendMessage(WM_CLOSE, 0, 0);
			return TRUE;

		case ID_ABOUT_EXIT:
			About();
			return TRUE;
#pragma endregion

#pragma region CHECKBOXES
		case IDC_CHECK1:
			if(IS_CHECKED(IDC_CHECK1))
				CPtInRectLog::Hook();
			else
				CPtInRectLog::UnHook();
			return TRUE;

		case IDC_CHECK2:
			pLogger->ToggleIgnoring(IS_CHECKED(IDC_CHECK2));
			return TRUE;

		case IDC_CHECK3:
			pLogger->ToggleExternCalls(!IS_CHECKED(IDC_CHECK3));
			return TRUE;
#pragma endregion
		}
		break;
#pragma endregion
	}

	return FALSE;
}

VOID PIRLOGAPI CreateGUI(HINSTANCE hInstance)
{
	Mutex = CreateMutex(NULL, FALSE, NULL);
	::hInstance = hInstance;
	hDialog = CreateDialog(hInstance, MAKEINTRESOURCE(IDD_DIALOG1), NULL, DlgProc);
	SendDlgItemMessage(IDC_LIST1, LVM_SETCOLUMNWIDTH, 0, 100);
	SendDlgItemMessage(IDC_LIST2, LVM_SETCOLUMNWIDTH, 0, 240);
}

VOID PIRLOGAPI DestroyGUI()
{
	for each (CALL_ITEM *pItem in lvItems)
		delete pItem;

	lvItems.clear();
	CloseHandle(Mutex);
	DestroyWindow(hDialog);
}

VOID PIRLOGAPI MessageLoop()
{
    MSG msg;
	INT status;

    while ((status = GetMessage(&msg, 0, 0, 0)) != 0)
    {
        if (status == -1)
            return;
        if (!IsDialogMessage(hDialog, &msg))
        {
            TranslateMessage(&msg);
            DispatchMessage(&msg);
        }
    }
}