// Copyright (c) 2021 Lv Decai, Shenzhen, China.
// FileEncryption.cpp: Install function hooks for reading and writing files.
// Encrypt and decrypt the contents while reading and writing files.
// Version: 1.0
// Release date: 2021年09月21日
// License: MIT (http://opensource.org/licenses/MIT)

#include "StdAfx.h"
#include <map>
#include <set>
#include <string>
#include <atlconv.h>
#include <atlexcept.h>
#include <detours.h>
#include "FileEncryption.h"

using namespace std;

// 堆栈锁,析构时自动解锁,便于使用
class CFileEncryptionRangeLock
{
public:
	// 构造函数,加锁
	CFileEncryptionRangeLock(CRITICAL_SECTION& CriticalSection)
		: m_CriticalSection(CriticalSection)
	{
		::EnterCriticalSection(&m_CriticalSection);
	}

	// 析构函数,解锁
	~CFileEncryptionRangeLock()
	{
		::LeaveCriticalSection(&m_CriticalSection);
	}

protected:
	CRITICAL_SECTION& m_CriticalSection;
};

// 监视的目录项结构体
struct CMonitorEntry
{
	// 目录项全路径字符串
	string Path;

	// 加密文件内容的函数指针
	ENCRYPT_FILE_DATA_PROC EncryptFileDataProc;

	// 解密文件内容的函数指针
	DECRYPT_FILE_DATA_PROC DecryptFileDataProc;

	// 重载比较运算符
	bool operator<(const CMonitorEntry& MonitorEntry) const
	{
		return _stricmp(Path.c_str(), MonitorEntry.Path.c_str()) < 0;
	}
};

// 打开文件信息结构体
struct COpenFileInfo : public CMonitorEntry
{
	// 加解密文件的密钥字符串
	string EncryptKey;
};

// 被监视的目录项集合
static map<CMonitorEntry, string> m_mapMonitorEntry;

// 被排除监视的目录项集合
static set<CMonitorEntry> m_setExcludeEntry;

// 被监视的目录下打开的文件集合
static map<HANDLE, COpenFileInfo> m_mapOpenFileInfo;

// 用于同步以上成员变量的充当线程锁的临界区变量
static CRITICAL_SECTION m_CriticalSection = { 0 };

// 指向一系列真实 API 函数的指针

static HANDLE(WINAPI *CreateFileAPtr)(
	_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 HANDLE(WINAPI *CreateFileWPtr)(
	_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 BOOL(WINAPI *CloseHandlePtr)(
	_In_ HANDLE hObject
	) = CloseHandle;

static _Must_inspect_result_ BOOL(WINAPI *ReadFilePtr)(
	_In_ HANDLE hFile,
	_Out_writes_bytes_to_opt_(nNumberOfBytesToRead, *lpNumberOfBytesRead) __out_data_source(FILE) LPVOID lpBuffer,
	_In_ DWORD nNumberOfBytesToRead,
	_Out_opt_ LPDWORD lpNumberOfBytesRead,
	_Inout_opt_ LPOVERLAPPED lpOverlapped
	) = ReadFile;

static _Must_inspect_result_ BOOL(WINAPI *ReadFileExPtr)(
	_In_ HANDLE hFile,
	_Out_writes_bytes_opt_(nNumberOfBytesToRead) __out_data_source(FILE) LPVOID lpBuffer,
	_In_ DWORD nNumberOfBytesToRead,
	_Inout_ LPOVERLAPPED lpOverlapped,
	_In_ LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine
	) = ReadFileEx;

static BOOL(WINAPI *WriteFilePtr)(
	_In_ HANDLE hFile,
	_In_reads_bytes_opt_(nNumberOfBytesToWrite) LPCVOID lpBuffer,
	_In_ DWORD nNumberOfBytesToWrite,
	_Out_opt_ LPDWORD lpNumberOfBytesWritten,
	_Inout_opt_ LPOVERLAPPED lpOverlapped
	) = WriteFile;

static BOOL(WINAPI *WriteFileExPtr)(
	_In_ HANDLE hFile,
	_In_reads_bytes_opt_(nNumberOfBytesToWrite) LPCVOID lpBuffer,
	_In_ DWORD nNumberOfBytesToWrite,
	_Inout_ LPOVERLAPPED lpOverlapped,
	_In_ LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine
	) = WriteFileEx;

// 判断字符串的有效性
BOOL ValidString(const char* s)
{
	return s != NULL
		&& s[0] != '0円';
}

// 添加要监视的目录
// 返回值:0 表示监视成功,-1 表示参数错误
int WINAPI MonitorEntry(const char* sPath, const char* sEncryptKey,
	ENCRYPT_FILE_DATA_PROC EncryptFileDataProc, DECRYPT_FILE_DATA_PROC DecryptFileDataProc)
{
	//LTRACE("sPath = %s, sEncryptKey = %s", sPath, sEncryptKey);

	if (!ValidString(sPath))
	{
		return -1;
	}

	CMonitorEntry MonitorEntry;
	MonitorEntry.Path.resize(strlen(sPath) + 1);
	GetShortPathNameA(sPath, (char*) MonitorEntry.Path.c_str(), (DWORD) MonitorEntry.Path.size());
	MonitorEntry.EncryptFileDataProc = EncryptFileDataProc;
	MonitorEntry.DecryptFileDataProc = DecryptFileDataProc;

	CFileEncryptionRangeLock Lock(m_CriticalSection);

	if (!ValidString(sEncryptKey))
	{
		// 如果密码为空,则表示排除监视
		m_setExcludeEntry.insert(MonitorEntry);

		m_mapMonitorEntry.erase(MonitorEntry);
	}
	else
	{
		// 密码不为空,表示要监视该目录
		m_mapMonitorEntry[MonitorEntry] = sEncryptKey;

		m_setExcludeEntry.erase(MonitorEntry);
	}

	return 0;
}

// 判断文件是否被监视,如是,将打开的文件句柄加入被监视文件集合
BOOL AddMonitorFile(const char* sFilePath, HANDLE hFile)
{
	if (!ValidString(sFilePath))
	{
		return FALSE;
	}

	CMonitorEntry MonitorEntry;

	MonitorEntry.Path.resize(strlen(sFilePath) + 1);
	GetShortPathNameA(sFilePath, (char*) MonitorEntry.Path.c_str(), (DWORD) MonitorEntry.Path.size());

	CFileEncryptionRangeLock Lock(m_CriticalSection);

	auto itInclude = m_mapMonitorEntry.find(MonitorEntry);

	if (itInclude != m_mapMonitorEntry.end())
	{
		// 如果文件全路径在被监视目录项集合中,要监视该文件

		//LTRACE("Hook CreateFile %p %s", hResult, lpFileName);

		// 将文件加入被监视文件集合
		COpenFileInfo OpenFileInfo;
		OpenFileInfo.Path = MonitorEntry.Path;
		OpenFileInfo.EncryptKey = itInclude->second;
		OpenFileInfo.EncryptFileDataProc = itInclude->first.EncryptFileDataProc;
		OpenFileInfo.DecryptFileDataProc = itInclude->first.DecryptFileDataProc;
		m_mapOpenFileInfo[hFile] = OpenFileInfo;

		return TRUE;
	}

	if (m_setExcludeEntry.find(MonitorEntry) != m_setExcludeEntry.end())
	{
		// 如果文件全路径在被排除监视目录项集合中,不监视该文件

		//LTRACE("Exclude hook CreateFile %s", lpFileName);

		return FALSE;
	}

	// 判断文件所在的目录是否位于被监视或被排除监视的目录项集合中
	// 算法缺陷:如果调用 MonitorEntry 添加的目录项不是完整的目录,而是目录字符串的一部分,会影响判断结果
	// 因此要求调用 MonitorEntry 添加的目录项时传入完整的路径

	size_t iExcludeMaxLength = 0;

	// 遍历被排除监视的目录项集合,找出包含当前被打开的文件路径的最长的项的长度
	for (auto& ExcludeEntry : m_setExcludeEntry)
	{
		if (_strnicmp(MonitorEntry.Path.c_str(), ExcludeEntry.Path.c_str(), ExcludeEntry.Path.size()) == 0)
		{
			if (ExcludeEntry.Path.size() > iExcludeMaxLength)
			{
				iExcludeMaxLength = ExcludeEntry.Path.size();
			}
		}
	}

	// 遍历被监视的目录项集合,找出包含当前被打开的文件路径的最长的项及其长度
	size_t iIncludeMaxLength = 0;
	map<CMonitorEntry, string>::iterator itIncludeMaxLength = m_mapMonitorEntry.end();

	for (itInclude = m_mapMonitorEntry.begin(); itInclude != m_mapMonitorEntry.end(); itInclude++)
	{
		auto& IncludeEntry = itInclude->first;

		if (_strnicmp(MonitorEntry.Path.c_str(), IncludeEntry.Path.c_str(), IncludeEntry.Path.size()) == 0)
		{
			if (IncludeEntry.Path.size() > iIncludeMaxLength)
			{
				iIncludeMaxLength = IncludeEntry.Path.size();
				itIncludeMaxLength = itInclude;
			}
		}
	}

	// 如果当前被打开的文件位于被监视的目录项集合中,且目录深度深于被排除监视的目录项,则要监视该文件
	if (itIncludeMaxLength != m_mapMonitorEntry.end() && iIncludeMaxLength > iExcludeMaxLength)
	{
		//LTRACE("Hook CreateFile %p %s", hResult, lpFileName);

		// 将文件加入被监视文件集合
		COpenFileInfo OpenFileInfo;
		OpenFileInfo.Path = MonitorEntry.Path;
		OpenFileInfo.EncryptKey = itIncludeMaxLength->second;
		OpenFileInfo.EncryptFileDataProc = itIncludeMaxLength->first.EncryptFileDataProc;
		OpenFileInfo.DecryptFileDataProc = itIncludeMaxLength->first.DecryptFileDataProc;
		m_mapOpenFileInfo[hFile] = OpenFileInfo;

		return TRUE;
	}
	else
	{
		//LTRACE("Exclude hook CreateFile %s", lpFileName);

		return FALSE;
	}
}

HANDLE WINAPI MyCreateFileA(
	__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 hResult = CreateFileAPtr(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes,
		dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);

	if (hResult != INVALID_HANDLE_VALUE)
	{
		AddMonitorFile(lpFileName, hResult);
	}

	return hResult;
}

HANDLE WINAPI MyCreateFileW(
	__in LPWSTR lpFileName,
	__in DWORD dwDesiredAccess,
	__in DWORD dwShareMode,
	__in_opt LPSECURITY_ATTRIBUTES lpSecurityAttributes,
	__in DWORD dwCreationDisposition,
	__in DWORD dwFlagsAndAttributes,
	__in_opt HANDLE hTemplateFile)
{
	HANDLE hResult = CreateFileWPtr(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes,
		dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);

	if (hResult != INVALID_HANDLE_VALUE)
	{
		AddMonitorFile(ATL::CW2A(lpFileName), hResult);
	}

	return hResult;
}

BOOL WINAPI MyCloseHandle(__in HANDLE hObject)
{
	BOOL bResult = CloseHandlePtr(hObject);

	CFileEncryptionRangeLock Lock(m_CriticalSection);

	// 从被监控的文件集合删除已关闭的文件
	if (m_mapOpenFileInfo.find(hObject) != m_mapOpenFileInfo.end())
	{
		//LTRACE("Hook CloseHandle %p", hObject);

		m_mapOpenFileInfo.erase(hObject);
	}

	return bResult;
}

BOOL WINAPI MyReadFile(
	__in HANDLE hFile,
	__out_bcount_part(nNumberOfBytesToRead, *lpNumberOfBytesRead) LPVOID lpBuffer,
	__in DWORD nNumberOfBytesToRead,
	__out_opt LPDWORD lpNumberOfBytesRead,
	__inout_opt LPOVERLAPPED lpOverlapped)
{
	BOOL bResult = FALSE;

	CFileEncryptionRangeLock Lock(m_CriticalSection);

	if (m_mapOpenFileInfo.find(hFile) != m_mapOpenFileInfo.end())
	{
		DWORD iOffset = SetFilePointer(hFile, 0, 0, FILE_CURRENT);

		//LTRACE("Hook ReadFile %p %u bytes at %u", hFile, nNumberOfBytesToRead, iOffset);

		bResult = ReadFilePtr(hFile, lpBuffer, nNumberOfBytesToRead, lpNumberOfBytesRead, lpOverlapped);

		if (bResult
			&& lpNumberOfBytesRead != NULL
			&& *lpNumberOfBytesRead > 0)
		{
			// 对读到的内容进行解密
			COpenFileInfo& OpenFileInfo = m_mapOpenFileInfo[hFile];

			if (OpenFileInfo.DecryptFileDataProc != NULL)
			{
				UINT iBufferLength = OpenFileInfo.DecryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
					(UINT) OpenFileInfo.EncryptKey.size(), iOffset, (PCBYTE) lpBuffer, *lpNumberOfBytesRead, NULL, 0);

				UCHAR* pBuffer = new UCHAR[iBufferLength];

				OpenFileInfo.DecryptFileDataProc(OpenFileInfo.EncryptKey.c_str(), (UINT) OpenFileInfo.EncryptKey.size(),
					iOffset, (PCBYTE) lpBuffer, *lpNumberOfBytesRead, pBuffer, iBufferLength);

				memcpy(lpBuffer, pBuffer, min(iBufferLength, *lpNumberOfBytesRead));

				delete[] pBuffer;
			}
		}
	}
	else
	{
		bResult = ReadFilePtr(hFile, lpBuffer, nNumberOfBytesToRead, lpNumberOfBytesRead, lpOverlapped);
	}

	return bResult;
}

BOOL WINAPI MyReadFileEx(
	__in HANDLE hFile,
	__out_bcount(nNumberOfBytesToRead) LPVOID lpBuffer,
	__in DWORD nNumberOfBytesToRead,
	__inout LPOVERLAPPED lpOverlapped,
	__in_opt LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	BOOL bResult = FALSE;

	CFileEncryptionRangeLock Lock(m_CriticalSection);

	if (m_mapOpenFileInfo.find(hFile) != m_mapOpenFileInfo.end())
	{
		DWORD iOffset = SetFilePointer(hFile, 0, 0, FILE_CURRENT);

		//LTRACE("Hook ReadFileEx %p %u bytes at %u", hFile, nNumberOfBytesToRead, iOffset);

		bResult = ReadFileExPtr(hFile, lpBuffer, nNumberOfBytesToRead, lpOverlapped, lpCompletionRoutine);

		if (bResult)
		{
			// 对读到的内容进行解密
			COpenFileInfo& OpenFileInfo = m_mapOpenFileInfo[hFile];

			if (OpenFileInfo.DecryptFileDataProc != NULL)
			{
				UINT iBufferLength = OpenFileInfo.DecryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
					(UINT) OpenFileInfo.EncryptKey.size(), iOffset, (PCBYTE) lpBuffer, nNumberOfBytesToRead, NULL, 0);

				BYTE* pBuffer = new BYTE[iBufferLength];

				iBufferLength = OpenFileInfo.DecryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
					(UINT) OpenFileInfo.EncryptKey.size(), iOffset,
					(PCBYTE) lpBuffer, nNumberOfBytesToRead, pBuffer, iBufferLength);

				memcpy(lpBuffer, pBuffer, min(iBufferLength, nNumberOfBytesToRead));

				delete[] pBuffer;
			}
		}
	}
	else
	{
		bResult = ReadFileExPtr(hFile, lpBuffer, nNumberOfBytesToRead, lpOverlapped, lpCompletionRoutine);
	}

	return bResult;
}

BOOL WINAPI MyWriteFile(
	__in HANDLE hFile,
	__in_bcount(nNumberOfBytesToWrite) LPCVOID lpBuffer,
	__in DWORD nNumberOfBytesToWrite,
	__out_opt LPDWORD lpNumberOfBytesWritten,
	__inout_opt LPOVERLAPPED lpOverlapped)
{
	CFileEncryptionRangeLock Lock(m_CriticalSection);

	if (m_mapOpenFileInfo.find(hFile) != m_mapOpenFileInfo.end())
	{
		DWORD iOffset = SetFilePointer(hFile, 0, 0, FILE_CURRENT);

		//LTRACE("Hook WriteFile %p %u bytes at %u", hFile, nNumberOfBytesToWrite, iOffset);

		// 对要写入的内容进行加密
		COpenFileInfo& OpenFileInfo = m_mapOpenFileInfo[hFile];

		if (OpenFileInfo.EncryptFileDataProc != NULL)
		{
			UINT iBufferLength = OpenFileInfo.EncryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
				(UINT) OpenFileInfo.EncryptKey.size(), iOffset, (PCBYTE) lpBuffer, nNumberOfBytesToWrite, NULL, 0);

			BYTE* pBuffer = new BYTE[iBufferLength];

			iBufferLength = OpenFileInfo.EncryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
				(UINT) OpenFileInfo.EncryptKey.size(), iOffset,
				(PCBYTE) lpBuffer, nNumberOfBytesToWrite, pBuffer, iBufferLength);

			BOOL bResult = WriteFilePtr(hFile, pBuffer, iBufferLength, lpNumberOfBytesWritten, lpOverlapped);

			delete[] pBuffer;

			return bResult;
		}
	}

	return WriteFilePtr(hFile, lpBuffer, nNumberOfBytesToWrite, lpNumberOfBytesWritten, lpOverlapped);
}

BOOL WINAPI MyWriteFileEx(
	__in HANDLE hFile,
	__in_bcount(nNumberOfBytesToWrite) LPCVOID lpBuffer,
	__in DWORD nNumberOfBytesToWrite,
	__inout LPOVERLAPPED lpOverlapped,
	__in_opt LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	CFileEncryptionRangeLock Lock(m_CriticalSection);

	if (m_mapOpenFileInfo.find(hFile) != m_mapOpenFileInfo.end())
	{
		DWORD iOffset = SetFilePointer(hFile, 0, 0, FILE_CURRENT);

		//LTRACE("Hook WriteFileEx %p %u bytes at %u", hFile, nNumberOfBytesToWrite, iOffset);

		// 对要写入的内容进行加密
		COpenFileInfo& OpenFileInfo = m_mapOpenFileInfo[hFile];

		if (OpenFileInfo.EncryptFileDataProc != NULL)
		{
			UINT iBufferLength = OpenFileInfo.EncryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
				(UINT)OpenFileInfo.EncryptKey.size(), iOffset, (PCBYTE) lpBuffer, nNumberOfBytesToWrite, NULL, 0);

			BYTE* pBuffer = new BYTE[iBufferLength];

			iBufferLength = OpenFileInfo.EncryptFileDataProc(OpenFileInfo.EncryptKey.c_str(),
				(UINT) OpenFileInfo.EncryptKey.size(),
				iOffset, (PCBYTE) lpBuffer, nNumberOfBytesToWrite, pBuffer, iBufferLength);

			BOOL bResult = WriteFileExPtr(hFile, pBuffer, iBufferLength, lpOverlapped, lpCompletionRoutine);

			delete[] pBuffer;

			return bResult;
		}
	}

	return WriteFileExPtr(hFile, lpBuffer, nNumberOfBytesToWrite, lpOverlapped, lpCompletionRoutine);
}

// 安装操作文件的 API 函数钩子
BOOL InstallFileEncryptionHook()
{
	DetourRestoreAfterWith();
	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());
	DetourAttach(&(PVOID&) CreateFileAPtr, MyCreateFileA);
	DetourAttach(&(PVOID&) CreateFileWPtr, MyCreateFileW);
	DetourAttach(&(PVOID&) CloseHandlePtr, MyCloseHandle);
	DetourAttach(&(PVOID&) ReadFilePtr, MyReadFile);
	DetourAttach(&(PVOID&) ReadFileExPtr, MyReadFileEx);
	DetourAttach(&(PVOID&) WriteFilePtr, MyWriteFile);
	DetourAttach(&(PVOID&) WriteFileExPtr, MyWriteFileEx);
	long iResult = DetourTransactionCommit();

	// 初始化临界区变量
	::InitializeCriticalSection(&m_CriticalSection);

	return TRUE;
}

// 卸载操作文件的 API 函数钩子
void UninstallFileEncryptionHook()
{
	DetourTransactionBegin();
	DetourUpdateThread(GetCurrentThread());
	DetourDetach(&(PVOID&) CreateFileAPtr, MyCreateFileA);
	DetourDetach(&(PVOID&) CreateFileWPtr, MyCreateFileW);
	DetourDetach(&(PVOID&) CloseHandlePtr, MyCloseHandle);
	DetourDetach(&(PVOID&) ReadFilePtr, MyReadFile);
	DetourDetach(&(PVOID&) ReadFileExPtr, MyReadFileEx);
	DetourDetach(&(PVOID&) WriteFilePtr, MyWriteFile);
	DetourDetach(&(PVOID&) WriteFileExPtr, MyWriteFileEx);
	long iResult = DetourTransactionCommit();

	// 释放临界区变量
	::DeleteCriticalSection(&m_CriticalSection);
}