#include "stdafx.h"
#include "GlobalHelperFunc.h"
#include <psapi.h>
#include <afxtempl.h>
#include <afxwin.h>
#include <tchar.h>
#include <winsvc.h>
#include <Tlhelp32.h>
#include <process.h>
#include "psapi.h"

#define  REMOTE_ADMIN_SERVICE              _T("RemoteAdminService")
#define  REMOTE_ADMIN_PIPE                 _T("RempteAdminPipe")
#define  REMOTE_ADMIN_PROCESS_INFO_PIPE    _T("RemoteAdminProcessInfoPipe")
#define  REMOTE_ADMIN_PROCESS_EXECUTE_PIPE _T("RemoteAdminProcessExecutePipe")
#define  REMOTE_ADMIN_PROCESS_KILL_PIPE    _T("RemoteAdminProcessKillPipe")
#define  REMOTE_ADMIN_SYS_SHUTDOWN_PIPE    _T("RemoteAdminSysShutDownPipe")
#define  SERVICENAME                       _T("RemoteAdminService")
#define  LONGSERVICENAME                   _T("RemoteAdminService")

struct SProcessInfo
{
    PROCESSENTRY32 peProcessEntry;
    SIZE_T         stMemUsage;
};

//typedef CTypedPtrList<CPtrList, PROCESSENTRY32*> CProcessInfoList;
typedef CTypedPtrList<CPtrList, SProcessInfo*> CProcessInfoList;

VOID WINAPI StartRemoteAdminService(DWORD, LPTSTR*);
void WINAPI RemoteAdminHandler(DWORD Opcode);
DWORD IsService(BOOL& isService);
void UpdateProcessInfoList(void* pArgument);
void CleanProcessInfoList();
void ServiceMain(void*);
void RemoteAdminThreadProc(void*);
void RemoteAdminThread(void*);
void RemoteAdminProcessInfoThread(void*);
void RemoteAdminExecuteProcessThread(void*);
void RemoteAdminKillProcessThread(void*);
void RemoteAdminSysShutdownThread(void* pParam);
void DeleteTheService();
BOOL SystemShutdown(LPTSTR lpMsg, BOOL bReboot, UINT iTimeOut /*Secs*/);
BOOL PreventSystemShutdown();

BOOL AddAceToWindowStation(HWINSTA hwinsta, PSID psid);
BOOL AddAceToDesktop(HDESK hdesk, PSID psid);
VOID FreeLogonSID (PSID *ppsid);
BOOL GetLogonSID (HANDLE hToken, PSID *ppsid);
BOOL StartInteractiveClientProcess (
    LPTSTR lpszUsername,    // client to log on
    LPTSTR lpszDomain,      // domain of client's account
    LPTSTR lpszPassword,    // client's password
    LPTSTR lpCommandLine    // command line to execute
); 


SIZE_T GetPhysicalMemUsage(DWORD dwProcessID);

#define DESKTOP_ALL (DESKTOP_READOBJECTS | DESKTOP_CREATEWINDOW |DESKTOP_CREATEMENU | DESKTOP_HOOKCONTROL |  DESKTOP_JOURNALRECORD | DESKTOP_JOURNALPLAYBACK | DESKTOP_ENUMERATE | DESKTOP_WRITEOBJECTS | DESKTOP_SWITCHDESKTOP | STANDARD_RIGHTS_REQUIRED)
#define WINSTA_ALL (WINSTA_ENUMDESKTOPS | WINSTA_READATTRIBUTES | WINSTA_ACCESSCLIPBOARD | WINSTA_CREATEDESKTOP | WINSTA_WRITEATTRIBUTES | WINSTA_ACCESSGLOBALATOMS | WINSTA_EXITWINDOWS | WINSTA_ENUMERATE | WINSTA_READSCREEN | STANDARD_RIGHTS_REQUIRED)
#define GENERIC_ACCESS (GENERIC_READ | GENERIC_WRITE | GENERIC_EXECUTE | GENERIC_ALL)



struct SCommand
{
    BOOL m_bThreadExit;
};

struct SExecuteCommand
{
    TCHAR m_szProcessPath[_MAX_PATH]; // Process to start 
    TCHAR m_szUsername[_MAX_PATH];    // client to log on
    TCHAR m_szDomain[_MAX_PATH];      // domain of client's account
    TCHAR m_szPassword[_MAX_PATH];    // client's password
};

struct SSysShutDownInfo
{
    BOOL bShutDown;   // TRUE if you want to shutdown,FALSE if want to cancel shutdown
    BOOL bReboot;     // Reboot if TRUE, else HALT if FALSE
    UINT iTimeToShutDown; // Time given to user before shutdown in secs
};

CProcessInfoList pilProcessInfoList;
HANDLE hStopServiceEvent;
long lServicePipeInstanceCount = 0;
SERVICE_STATUS          ServiceStatus; 
SERVICE_STATUS_HANDLE   ServiceStatusHandle;
CRITICAL_SECTION g_CriticalSection;
 

int main(int argc, char* argv[])
{
    SERVICE_TABLE_ENTRY DispatchTable[] = { 
        {REMOTE_ADMIN_SERVICE, StartRemoteAdminService}, 
        {NULL, NULL}
    }; 

    BOOL bIsService = FALSE;
    ::IsService(bIsService);
    
    if (bIsService)
    {
		::InitializeCriticalSection(&g_CriticalSection);
        return StartServiceCtrlDispatcher(DispatchTable);
    }
    
   	return 0;
}


// This process is a service or is not ?
DWORD IsService(BOOL& isService)
{
    DWORD pID = GetCurrentProcessId(); 
	HANDLE hProcessToken = NULL;
	DWORD groupLength = 50;
	PTOKEN_GROUPS groupInfo = NULL;

	SID_IDENTIFIER_AUTHORITY siaNt = SECURITY_NT_AUTHORITY;
	PSID pInteractiveSid = NULL;
	PSID pServiceSid = NULL;

	DWORD dwRet = NO_ERROR;
    
    // reset flags
	BOOL isInteractive = FALSE;
	isService = FALSE;

	DWORD ndx;

    HANDLE hProcess = ::OpenProcess(PROCESS_ALL_ACCESS, FALSE, pID);

	// open the token
	if (!::OpenProcessToken(hProcess, TOKEN_QUERY, &hProcessToken))
	{
		dwRet = ::GetLastError();
		goto closedown;
	}

	// allocate a buffer of default size
	groupInfo = (PTOKEN_GROUPS)::LocalAlloc(0, groupLength);
	if (groupInfo == NULL)
	{
		dwRet = ::GetLastError();
		goto closedown;
	}

	// try to get the info
	if (!::GetTokenInformation(hProcessToken, TokenGroups, groupInfo, groupLength, &groupLength))
	{
		// if buffer was too small, allocate to proper size, otherwise error
		if (::GetLastError() != ERROR_INSUFFICIENT_BUFFER)
		{
			dwRet = ::GetLastError();
			goto closedown;
		}

		::LocalFree(groupInfo);

		groupInfo = (PTOKEN_GROUPS)::LocalAlloc(0, groupLength);
		if (groupInfo == NULL)
		{
			dwRet = ::GetLastError();
			goto closedown;
		}

		if (!GetTokenInformation(hProcessToken, TokenGroups, groupInfo, groupLength, &groupLength))
		{
			dwRet = ::GetLastError();
			goto closedown;
		}
	}

	// create comparison sids
	if (!AllocateAndInitializeSid(&siaNt, 1, SECURITY_INTERACTIVE_RID, 0, 0, 0, 0, 0, 0, 0, &pInteractiveSid))
	{
		dwRet = ::GetLastError();
		goto closedown;
	}

	if (!AllocateAndInitializeSid(&siaNt, 1, SECURITY_SERVICE_RID, 0, 0, 0, 0, 0, 0, 0, &pServiceSid))
	{
		dwRet = ::GetLastError();
		goto closedown;
	}

	// try to match sids
	for (ndx = 0; ndx < groupInfo->GroupCount ; ndx += 1)
	{
		SID_AND_ATTRIBUTES sanda = groupInfo->Groups[ndx];
		PSID pSid = sanda.Sid;

		if (::EqualSid(pSid, pInteractiveSid))
		{
			isInteractive = TRUE;
			isService = FALSE;
			break;
		}
		else if (::EqualSid(pSid, pServiceSid))
		{
			isService = TRUE;
			isInteractive = FALSE;
			break;
		}
	}

   if ( !(isService || isInteractive ))
		isService = TRUE;
        
closedown:
		if (pServiceSid)
        {
            ::FreeSid(pServiceSid);
        }
		
		if (pInteractiveSid)
        {
            ::FreeSid(pInteractiveSid);
        }
		
		if (groupInfo)
        {
            ::LocalFree(groupInfo);
        }
			

		if (hProcessToken)
        {
            ::CloseHandle(hProcessToken);
        }
			

		if (hProcess)
        {
            ::CloseHandle(hProcess);
        }
			

	return dwRet;
}

void UpdateProcessInfoList(void* pArgument)
{
    for (;;)
    {
        CleanProcessInfoList();

		::EnterCriticalSection(&g_CriticalSection);
		
        DWORD dwCurrentProcessId = ::GetCurrentProcessId();
    
        // Take this process's snapshot
        HANDLE hProcessSnapShot = ::CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, dwCurrentProcessId);

        //PROCESSENTRY32* pPe = new PROCESSENTRY32;
		//pPe->dwSize = sizeof(PROCESSENTRY32); 
        SProcessInfo* pPi = new SProcessInfo;
        BOOL bOk = ::Process32First(hProcessSnapShot, &pPi->peProcessEntry);
        
        // Get the process memeory usage
        pPi->stMemUsage = GetPhysicalMemUsage(pPi->peProcessEntry.th32ProcessID);
        
        if (!bOk)
        {
            delete pPi;
            pPi = NULL;
        }
        
        while (bOk)
        {
            pilProcessInfoList.AddTail(pPi);

            //pPe = new PROCESSENTRY32;
            pPi = new SProcessInfo;
            bOk = ::Process32Next(hProcessSnapShot, &pPi->peProcessEntry);
            
            // Get the process memeory usage
            pPi->stMemUsage = GetPhysicalMemUsage(pPi->peProcessEntry.th32ProcessID);

            if (!bOk)
            {
                delete pPi;
                pPi = NULL;
            }
        }

        // We didn't forget to clean up the snapshot object.
        ::CloseHandle(hProcessSnapShot);
        ::Sleep(2000);

		::LeaveCriticalSection(&g_CriticalSection);
    }

    return ;
}

void CleanProcessInfoList()
{
	//PROCESSENTRY32* pPe = NULL;
    SProcessInfo* pPi = NULL;

	::EnterCriticalSection(&g_CriticalSection);

    POSITION pos = pilProcessInfoList.GetHeadPosition();
    while (pos)
    {
        pPi = pilProcessInfoList.GetNext(pos);

		if (pPi != NULL)
		{
			delete pPi;
		}
    }
            
    pilProcessInfoList.RemoveAll();
	
	::LeaveCriticalSection(&g_CriticalSection);

    return;
}


// Start service
VOID WINAPI StartRemoteAdminService(DWORD, LPTSTR*) 
{
   DWORD status = 0; 
   DWORD specificError = 0;

   // Prepare the ServiceStatus structure that will be used for the
   // comunication with SCM(Service Control Manager).
   // If you fully under stand the members of this structure, feel
   // free to change these values :o)
   ServiceStatus.dwServiceType        = SERVICE_WIN32; 
   ServiceStatus.dwCurrentState       = SERVICE_START_PENDING; 
   ServiceStatus.dwControlsAccepted   = SERVICE_ACCEPT_STOP; 
   ServiceStatus.dwWin32ExitCode      = 0; 
   ServiceStatus.dwServiceSpecificExitCode = 0; 
   ServiceStatus.dwCheckPoint         = 0; 
   ServiceStatus.dwWaitHint           = 0; 

   // Here we register the control handler for our service.
   // We tell the SCM about a call back function that SCM will
   // call when user tries to Start, Stop or Pause your service.
   ServiceStatusHandle = RegisterServiceCtrlHandler( 
         TEXT("Service"), RemoteAdminHandler); 

   if (ServiceStatusHandle == (SERVICE_STATUS_HANDLE)0) 
      return; 
   
   // Handle error condition 
   if (status != NO_ERROR) 
   { 
      ServiceStatus.dwCurrentState       = SERVICE_STOPPED; 
      ServiceStatus.dwCheckPoint         = 0; 
      ServiceStatus.dwWaitHint           = 0; 
      ServiceStatus.dwWin32ExitCode      = status; 
      ServiceStatus.dwServiceSpecificExitCode = specificError; 

      SetServiceStatus (ServiceStatusHandle, &ServiceStatus); 
      return; 
   } 

   // Initialization complete - report running status. 
   ServiceStatus.dwCurrentState       = SERVICE_RUNNING; 
   ServiceStatus.dwCheckPoint         = 0; 
   ServiceStatus.dwWaitHint           = 0; 

   if (!SetServiceStatus (ServiceStatusHandle, &ServiceStatus)) 
      status = GetLastError(); 
   else
   {
      // Start the main thread
      hStopServiceEvent = CreateEvent( NULL, FALSE, FALSE, NULL );
      _beginthread(ServiceMain, 0, NULL );
   }
   
   return; 
} 


void ServiceMain(void*)
{
    ::_beginthread(UpdateProcessInfoList, 0, NULL);

    // Start CommunicationPoolThread, which handles the incoming instances
    ::_beginthread(RemoteAdminThread, 0, NULL );

    // Waiting for stop the service
    while(::WaitForSingleObject(hStopServiceEvent, 10) != WAIT_OBJECT_0)
    {
    }
   
    // Let's delete itself, after the service stopped
    ::DeleteTheService();

    ::CloseHandle(hStopServiceEvent);
}

void RemoteAdminThread(void*)
{
    HANDLE hPipe = NULL;
    
    for (;;)
    {
        SECURITY_ATTRIBUTES SecAttrib = {0};
        SECURITY_DESCRIPTOR SecDesc;
        InitializeSecurityDescriptor(&SecDesc, SECURITY_DESCRIPTOR_REVISION);
        SetSecurityDescriptorDacl(&SecDesc, TRUE, NULL, TRUE);

        SecAttrib.nLength = sizeof(SECURITY_ATTRIBUTES);
        SecAttrib.lpSecurityDescriptor = &SecDesc;;
        SecAttrib.bInheritHandle = TRUE;

        // Create communication pipe
        hPipe = ::CreateNamedPipe(