#include "stdafx.h"

#include <conio.h>

#include <shlwapi.h>

#include <Tlhelp32.h>



//=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+

//return PID from a process name

ULONG     GetProcessID(char *szProcessName)

{

	PROCESSENTRY32      ProcessEntry32;

	HANDLE              hSnap;

	int                    ret;



	ProcessEntry32.dwSize = sizeof(PROCESSENTRY32);



	hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

	if (hSnap == 0)

		return 0;



	ret = Process32First(hSnap, &ProcessEntry32);

	while (ret)

	{

		if (StrStrI(ProcessEntry32.szExeFile, szProcessName))

		{

			return ProcessEntry32.th32ProcessID;

		}

		ret = Process32Next(hSnap, &ProcessEntry32);

	}



	CloseHandle(hSnap);

	return 0;

}

//=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+



void BuilHoldFile(CString strHostProcessName, CString strFileName)

{

	DWORD pid = GetProcessID((char*)(LPCTSTR)strHostProcessName);

	if (pid == 0)

	{

		printf_s("can not get the pid of the process:%s", strHostProcessName);

		return;

	}

	HANDLE hProcess = OpenProcess(PROCESS_DUP_HANDLE, FALSE, pid);

	if (hProcess == NULL)

	{

		printf_s("can not open process:%s[%d]\n", strHostProcessName, (DWORD)pid);

		return;

	}



	HANDLE hFile = CreateFile(strFileName, GENERIC_READ, 0, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL | FILE_FLAG_DELETE_ON_CLOSE, NULL);

	if (hFile == INVALID_HANDLE_VALUE)

	{

		CloseHandle(hProcess);

		printf_s("open file failed:%s\n", strFileName);

		return;

	}



	if (!DuplicateHandle(GetCurrentProcess(), hFile, hProcess, &hFile, 0,

		FALSE, DUPLICATE_CLOSE_SOURCE | DUPLICATE_SAME_ACCESS))

	{

		CloseHandle(hFile);

		CloseHandle(hProcess);

		printf_s("DuplicateHandle failed\n");

		return;

	}



	MoveFileEx(strFileName, NULL, MOVEFILE_DELAY_UNTIL_REBOOT);

	CloseHandle(hFile);

	CloseHandle(hProcess);

	printf_s("build hold file: %s OK!\n", strFileName);

}


// 提高本进程权限,以取得系统进程的信息
BOOL WINAPI EnableDebugPrivilege(BOOL bEnable)

{

	// 附给本进程特权，用以取得系统进程的信息

	BOOL bOk = FALSE;    // Assume function fails

	HANDLE hToken;



	// 打开一个进程的访问令牌

	if (OpenProcessToken(::GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES,

		&hToken))

	{

		TOKEN_PRIVILEGES tp;

		tp.PrivilegeCount = 1;

		LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tp.Privileges[0].Luid);

		tp.Privileges[0].Attributes = bEnable ? SE_PRIVILEGE_ENABLED : 0;

		AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(tp), NULL, NULL);

		bOk = (GetLastError() == ERROR_SUCCESS);

		::CloseHandle(hToken);

	}

	return(bOk);

}

int _tmain(int argc, _TCHAR* argv[]) {
	if (argc > 1) {
		EnableDebugPrivilege(TRUE);

		if (PathIsDirectory(argv[1])) {
			CString strFolder = argv[1];
			if (strFolder.Right(1) != "\\") {
				strFolder += "\\";
			}
			CFileFind finder;
			BOOL bWorking = finder.FindFile(strFolder + "*.*");
			while (bWorking) {
				bWorking = finder.FindNextFile();
				//skip . and ..
				if (finder.IsDots())
					continue;
				else if (finder.IsDirectory()) {
					continue;
				} else {
					BuilHoldFile("winlogon.exe", finder.GetFilePath());
				}
			}
			finder.Close();
		} else {
			BuilHoldFile("winlogon.exe", argv[1]);
		}
	} else {
		printf_s(
			"参数为[目录名]将该目录下的文件批量占坑\n"
			"参数为[文件名]将该文件占坑\n");
	}

	getch();
	return 0;
}