

#include "hook.h"

#pragma comment(lib,"libeay32.lib")

typedef INT (WINAPI* CREATEPROCESSW)(LPCWSTR,LPWSTR,LPSECURITY_ATTRIBUTES,LPSECURITY_ATTRIBUTES,
BOOL,
DWORD,
LPVOID,
LPCWSTR,
LPSTARTUPINFOW,
LPPROCESS_INFORMATION);

LPVOID orgCreateProc;

int IsQuoted(LPWSTR lpCommandLine)
{
	if(wcschr(lpCommandLine,L'\"')==lpCommandLine)
		return 1;
	return 0;
}

int IsRemovableDrive(wchar_t *path,int quoted)
{
	wchar_t drv[3];

	if(quoted==0)
		drv[0]=path[0];
	else
		drv[0]=path[1];

	drv[1]=L':';
	drv[2]=0;

	if(GetDriveTypeW(drv)!=DRIVE_REMOVABLE)
		return 0;
	return 1;

}

int IsStopItEnabled()
{
	HKEY hkey;
	DWORD lpdwDisposition,lpcbData,lpData;

	DWORD data=0;

	RegCreateKeyExW(HKEY_LOCAL_MACHINE,L"Software\\StopIt",
		0,L"REG_DWORD",REG_OPTION_NON_VOLATILE,KEY_QUERY_VALUE,0,&hkey,&lpdwDisposition);

	if(hkey==0) // open failed!
		return 0;

	if(RegQueryValueExW(hkey,L"enabled",NULL,&lpData,(BYTE*)&data,&lpcbData)!=ERROR_SUCCESS)
	{
		RegCloseKey(hkey);
		return 0;// open failed!
	}
	else // open ok!
	{
		RegCloseKey(hkey);
		if(data)
			return 1; // task manager disabled!
		return 0;
	}
}

void ShowError(wchar_t* path)
{
	wchar_t *template_msg=L"You are not permitted to execute following program.\n\n%s";
	wchar_t *details=(wchar_t *)malloc((2*lstrlenW(path))+(2*lstrlenW(template_msg))+2);
	swprintf(details,template_msg,path);
	MessageBoxW(0,details,L"StopIt",MB_ICONERROR|MB_TOPMOST);
	free(details);
}

DWORD WINAPI myCreateProcessW(
	LPCWSTR lpApplicationName,
	LPWSTR lpCommandLine, 
	LPSECURITY_ATTRIBUTES lpProcessAttributes,
	LPSECURITY_ATTRIBUTES lpThreadAttributes,
	BOOL bInheritHandles,
	DWORD dwCreationFlags,
	LPVOID lpEnvironment,
	LPCWSTR lpCurrentDirectory,
	LPSTARTUPINFOW lpStartupInfo,
	LPPROCESS_INFORMATION lpProcessInformation
)
{

	CREATEPROCESSW crtProc;
	BOOL ifsuccess;
	DWORD err;
	
	if(IsStopItEnabled())
	{
		if(lstrlenW(lpApplicationName))
		{
			if(IsRemovableDrive((wchar_t*)lpApplicationName,0))
			{
				ShowError((wchar_t*)lpApplicationName);
				return 1;
			}
		}
		else // command line!
		{
			if(IsQuoted(lpCommandLine))
			{
				if(IsRemovableDrive(lpCommandLine,1))
				{
					ShowError(lpCommandLine);
					return 1;
				}
			}else
			{
				if(IsRemovableDrive(lpCommandLine,0))
				{
					ShowError(lpCommandLine);
					return 1;
				}
			}
		}
	}

	crtProc=(CREATEPROCESSW)orgCreateProc;
	
ifsuccess = crtProc(lpApplicationName,
			lpCommandLine, lpProcessAttributes,
			lpThreadAttributes, bInheritHandles, dwCreationFlags, lpEnvironment,
			lpCurrentDirectory, lpStartupInfo, lpProcessInformation);
			err =GetLastError();
			SetLastError(err);
			
	return (DWORD)ifsuccess;		



}


BOOL APIENTRY DllMain( HANDLE hModule, 
                       DWORD  ul_reason_for_call, 
                       LPVOID lpReserved
					 )
{

	if(ul_reason_for_call==DLL_PROCESS_ATTACH){
		orgCreateProc=HookFunctionInCurrentProcess("kernel32.dll","CreateProcessW",&myCreateProcessW);
	}

	if(ul_reason_for_call==DLL_PROCESS_DETACH){
		UnHookFunctionInCurrentProcess("kernel32.dll","CreateProcessW",orgCreateProc);
	}

    return TRUE;
}
