#include "stdafx.h"

#include "define.h"

#include "Hook.h"
#include "zpzDll.h"

#include "Detours/detours.h"
#include "Detours/detver.h"

#include <algorithm>

extern HWND g_hWnd;
extern char g_sComNum[MAX_PATH];
extern wchar_t g_wsComNum[MAX_PATH];
extern char g_data[MAX_DATA_LENGTH];
extern int data_length;

//这个是打开COM1的hport
HANDLE g_hPort = NULL;

//把这个被hook进程的工作的端口号保存下来
std::string g_sworkComNum;
std::wstring g_wsworkComNum;

static HANDLE (WINAPI * Real_CreateFileW)(
	__in     LPCWSTR lpFileName,
	__in     DWORD dwDesiredAccess,
	__in     DWORD dwShareMode,
	__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
	__in     DWORD dwCreationDisposition,
	__in     DWORD dwFlagsAndAttributes,
	__in_opt HANDLE hTemplateFile
	) = CreateFileW;

static HANDLE(WINAPI * Real_CreateFileA)(
	__in     LPCSTR lpFileName,
	__in     DWORD dwDesiredAccess,
	__in     DWORD dwShareMode,
	__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
	__in     DWORD dwCreationDisposition,
	__in     DWORD dwFlagsAndAttributes,
	__in_opt HANDLE hTemplateFile
	) = CreateFileA;

static BOOL(WINAPI * Real_WriteFile)(
	__in        HANDLE hFile,
	__in_bcount_opt(nNumberOfBytesToWrite) LPCVOID lpBuffer,
	__in        DWORD nNumberOfBytesToWrite,
	__out_opt   LPDWORD lpNumberOfBytesWritten,
	__inout_opt LPOVERLAPPED lpOverlapped
	) = WriteFile;

HANDLE WINAPI Mine_CreateFileW(
	__in     LPCWSTR lpFileName,
	__in     DWORD dwDesiredAccess,
	__in     DWORD dwShareMode,
	__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
	__in     DWORD dwCreationDisposition,
	__in     DWORD dwFlagsAndAttributes,
	__in_opt HANDLE hTemplateFile
)
{
	HANDLE hPort = Real_CreateFileW(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);

	std::wstring wsFileName = lpFileName;
	std::transform(wsFileName.begin(), wsFileName.end(), wsFileName.begin(), ::toupper);

	//if (wsFileName.find(L"COM1") != wsFileName.npos)
	//{
	//	int a = 1;
	//}

	std::wstring wsComNum = g_wsComNum;
	if (wsComNum != L"" && wsFileName.find(g_wsComNum) != wsFileName.npos)
	{
		//打开的是目标端口
		if (hPort != INVALID_HANDLE_VALUE)
		{
			//成功打开端口，保存端口句柄，开始监听
			g_hPort = hPort;

			//把这个进程监听的端口号，存起来；
			g_wsworkComNum = g_wsComNum;

			g_sworkComNum = "";
		}
	}

	return hPort;
}

HANDLE WINAPI Mine_CreateFileA(
	__in     LPCSTR lpFileName,
	__in     DWORD dwDesiredAccess,
	__in     DWORD dwShareMode,
	__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
	__in     DWORD dwCreationDisposition,
	__in     DWORD dwFlagsAndAttributes,
	__in_opt HANDLE hTemplateFile
)
{
	HANDLE hPort = Real_CreateFileA(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);

	std::string FileName = lpFileName;
	std::transform(FileName.begin(), FileName.end(), FileName.begin(), ::toupper);

	//if (FileName.find("COM1") != FileName.npos)
	//{
	//	int a = 1;
	//}

	std::string sComNum = g_sComNum;
	if (sComNum != "" && FileName.find(g_sComNum) != FileName.npos)
	{
		//打开的是目标端口
		if (hPort != INVALID_HANDLE_VALUE)
		{
			//成功打开端口，保存端口句柄，开始监听
			g_hPort = hPort;

			//把这个进程监听的端口号，存起来；
			g_sworkComNum = g_sComNum;

			g_wsworkComNum = L"";
		}
	}

	return hPort;
}

BOOL WINAPI Mine_WriteFile(
	__in        HANDLE hFile,
	__in_bcount_opt(nNumberOfBytesToWrite) LPCVOID lpBuffer,
	__in        DWORD nNumberOfBytesToWrite,
	__out_opt   LPDWORD lpNumberOfBytesWritten,
	__inout_opt LPOVERLAPPED lpOverlapped
)
{
	BOOL ret = Real_WriteFile(hFile, lpBuffer, nNumberOfBytesToWrite, lpNumberOfBytesWritten, lpOverlapped);

	//同时满足发送数据的句柄和创建目标端口时候的句柄相等，以及工作的端口号等于最新的端口号，才进行监听
	if (hFile == g_hPort && (g_sworkComNum == g_sComNum || g_wsworkComNum == g_wsComNum))
	{
		//捕捉到一条客显消息，先把消息内容截取下来
		DWORD nWirten = nNumberOfBytesToWrite;

		if (nWirten > MAX_DATA_LENGTH)
		{
			//消息太长了，忽略掉
			return ret;
		}

		memcpy_s(g_data + data_length, nWirten, (char*)lpBuffer, nWirten);
		data_length += nWirten;

		//再发一个消息通知插件，去读取消息（//用这个类型，是为了绕过UAC，低权限可以用这个类型像管理员权限的进程POST消息）
		::PostMessage(g_hWnd, WM_GETHOTKEY, 0, 0);
	}
	
	return ret;
}

void Hook()
{
	LONG error;

	DetourRestoreAfterWith();
	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());

	DetourAttach(&(PVOID&)Real_CreateFileW, Mine_CreateFileW);
	DetourAttach(&(PVOID&)Real_CreateFileA, Mine_CreateFileA);

	DetourAttach(&(PVOID&)Real_WriteFile, Mine_WriteFile);

	error = DetourTransactionCommit();

	if (error == NO_ERROR) {
		printf("echofx" DETOURS_STRINGIFY(DETOURS_BITS) ".dll:"
			" Detoured Echo().\n");
	}
	else {
		printf("echofx" DETOURS_STRINGIFY(DETOURS_BITS) ".dll:"
			" Error detouring Echo(): %d\n", error);
	}
}

void UnHook()
{
	LONG error;

	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());

	DetourDetach(&(PVOID&)Real_CreateFileW, Mine_CreateFileW);
	DetourDetach(&(PVOID&)Real_CreateFileA, Mine_CreateFileA);

	DetourDetach(&(PVOID&)Real_WriteFile, Mine_WriteFile);

	error = DetourTransactionCommit();

	printf("echofx" DETOURS_STRINGIFY(DETOURS_BITS) ".dll:"
		" Removed Echo() (result=%d)\n", error);
	fflush(stdout);
}