#include "stdafx.h"
#include "RemoteAdministrator.h"
#include "GlobalHelperFunc.h"
#include "GlobalMFCHelperFunc.h"
#include "MachineInfo.h"
#include "Resource.h"
#include <winsvc.h>

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif


extern CRITICAL_SECTION g_CriticalSection;

CRemoteAdministrator::CRemoteAdministrator()
{

}

CRemoteAdministrator::~CRemoteAdministrator()
{
	// Clean the connection pending list
	CString* pStr = NULL;
	POSITION pos = m_cplPendingConnectionList.GetHeadPosition();
	while (pos != NULL)
	{
		pStr = m_cplPendingConnectionList.GetNext(pos);

		ASSERT(pStr != NULL);

		if (pStr != NULL)
		{
			delete pStr;
		}
	}
	m_cplPendingConnectionList.RemoveAll();


	// Clean the connected machine list
	CMachineInfo* pMachineInfo = NULL;
	pos = m_milConnectedMachines.GetHeadPosition();
	while (pos != NULL)
	{
		pMachineInfo = m_milConnectedMachines.GetNext(pos);

		ASSERT(pMachineInfo != NULL);

		if (pMachineInfo != NULL)
		{
			delete pMachineInfo;
		}
	}
	m_milConnectedMachines.RemoveAll();
}

BOOL CRemoteAdministrator::EstablishConnection(CString strResource, CString strLogon, CString strRemoteMachineIP, CString strPwd, BOOL bEstablish)
{
	::EnterCriticalSection(&g_CriticalSection);

    TCHAR szRemoteResource[_MAX_PATH];
    DWORD dwRetVal = NO_ERROR;

    // Remote resource, \\remote\ipc$, remote\admin$, ...
    ::sprintf(szRemoteResource, _T("\\\\%s\\%s"), strRemoteMachineIP.GetBuffer(0), strResource.GetBuffer(0));

    //
    // Disconnect or connect to the resource, based on bEstablish
    //
    if (bEstablish) 
    {
        NETRESOURCE nr;
        nr.dwType = RESOURCETYPE_ANY;
        nr.lpLocalName = NULL;
        nr.lpRemoteName = (LPTSTR)&szRemoteResource;
        nr.lpProvider = NULL;
      
        //Establish connection (using username/pwd)
        dwRetVal = ::WNetAddConnection2(
                        &nr, 
                        strPwd.GetBuffer(0), 
                        strLogon.GetBuffer(0), 
                        FALSE
                        );
        
        // Let the caller generate the error message
        /*if (dwRetVal != NO_ERROR)
        {
            //::PopError(dwRetVal);
            /*CString strFromatAPIMsg = ::FormatError(dwRetVal);
            CString strDisplayMsg;
            strDisplayMsg.Format("Machine IP %s: %s", strRemoteMachineIP.GetBuffer(0), strFromatAPIMsg.GetBuffer(0));
            
            ::AfxMessageBox(strDisplayMsg);*/
        //}*/
    }
    else
    {
        // Disconnect
        dwRetVal = ::WNetCancelConnection2(szRemoteResource, 0, NULL);

        // Let the caller generate the error message
        /*if (dwRetVal != NO_ERROR)
        {
            //::PopError(dwRetVal);
            /*CString strFromatAPIMsg = ::FormatError(dwRetVal);
            CString strDisplayMsg;
            strDisplayMsg.Format("Machine IP %s: %s", strRemoteMachineIP.GetBuffer(0), strFromatAPIMsg.GetBuffer(0));
            
            ::AfxMessageBox(strDisplayMsg);*/
        //}*/
     }
      
    // Prepare the return value
    if (dwRetVal == NO_ERROR) 
    {
		::LeaveCriticalSection(&g_CriticalSection);
        return TRUE; // indicate success
    }

	::LeaveCriticalSection(&g_CriticalSection);
    return FALSE;
}


void CRemoteAdministrator::AddMachine(CMachineInfo& miMachineInfo)
{
    CMachineInfo* pmiMachineInfo = new CMachineInfo;
    
    ASSERT(pmiMachineInfo);

    if (pmiMachineInfo != NULL)
    {
        (*pmiMachineInfo) = miMachineInfo;

        // Update the connected machines list
        m_milConnectedMachines.AddHead(pmiMachineInfo);
		::Sleep(1000);
    }
    else
    {
        ::AfxMessageBox(_T("No memory for new machine details !"));
    }
}

void CRemoteAdministrator::DeleteMachine(CMachineInfo& miMachineInfo)
{
    CMachineInfo* pMachineInfo = NULL;
    POSITION pos = m_milConnectedMachines.GetHeadPosition();
    
	::EnterCriticalSection(&g_CriticalSection);

    while (pos != (POSITION)0xcdcdcdcd && pos != NULL)
    {
        pMachineInfo = m_milConnectedMachines.GetAt(pos);
        
        ASSERT (pMachineInfo);

        if (miMachineInfo == *pMachineInfo) 
        {
            // delete the machine from list
            m_milConnectedMachines.RemoveAt(pos);

            delete pMachineInfo;
            pMachineInfo = NULL;

			::LeaveCriticalSection(&g_CriticalSection);
            return;

        }

        m_milConnectedMachines.GetNext(pos);
    }

	::LeaveCriticalSection(&g_CriticalSection);
}


BOOL CRemoteAdministrator::CheckIfMachinePresent(CString strIP)
{
    // Machine with same IP cannot be added to the monitoring list
    CMachineInfo* pMachineInfo = NULL;
    POSITION pos = m_milConnectedMachines.GetHeadPosition();
    
    while (pos != (POSITION)0xcdcdcdcd && pos != NULL)
    {
        pMachineInfo = m_milConnectedMachines.GetNext(pos);
        
        ASSERT (pMachineInfo);

        if (strIP == pMachineInfo->GetIP()) 
        {
            return TRUE;
        }
    }

    return FALSE;
}


void CRemoteAdministrator::RefreshProcessList(CString strIP, CProcessInfoList& pilProcessList)
{
    CMachineInfo* pMachineInfo = GetMachineInfo(strIP);
    
    ASSERT(pMachineInfo);
    
    if (pMachineInfo != NULL)
    {
        pMachineInfo->RefreshProcessList(pilProcessList);
    }
}

CMachineInfo* CRemoteAdministrator::GetMachineInfo(CString strIP)
{
    CMachineInfo* pMachineInfo = NULL;
    POSITION pos = NULL;

	::EnterCriticalSection(&g_CriticalSection);
    pos = m_milConnectedMachines.GetHeadPosition();

    while (pos != (POSITION)0xcdcdcdcd && pos != NULL)
    {
        pMachineInfo = m_milConnectedMachines.GetNext(pos);

        if (pMachineInfo->GetIP() == strIP)
        {
			::LeaveCriticalSection(&g_CriticalSection);
            return pMachineInfo;
        }
    }

	::LeaveCriticalSection(&g_CriticalSection);

    // No machine with IP strIP is present
    ASSERT (NULL);

    return NULL;
}

BOOL CRemoteAdministrator::EstablishAdminConnection(CString strRemoteMachineIP, CString strPwd, BOOL bEstablish)
{
	CString strResource = _T("ADMIN$");
	CString strLogon    = _T("Administrator");

	BOOL bConnectionSuccess = EstablishConnection(strResource, strLogon, strRemoteMachineIP, strPwd, bEstablish);
    
    if (!bConnectionSuccess)
    {
        CString strFormattedErrorMsg = ErrorHandling::ConvertStringTableIDToErrorMsg(strRemoteMachineIP, IDS_NO_ADMIN_CONNECTION);

        //::AfxMessageBox(strFormattedErrorMsg);
		//::GetMainFrame()->ShowBalloonMsgInTray(_T("Warning"), strFormattedErrorMsg);
    }
    return bConnectionSuccess;
}

BOOL CRemoteAdministrator::EstablishIPCConnection(CString strRemoteMachineIP, CString strPwd, BOOL bEstablish)
{
	CString strResource = _T("IPC$");
	CString strLogon    = _T("Administrator");

    BOOL bConnectionSuccess = EstablishConnection(strResource, strLogon, strRemoteMachineIP, strPwd, bEstablish);

    if (!bConnectionSuccess)
    {
        CString strFormattedErrorMsg = ErrorHandling::ConvertStringTableIDToErrorMsg(strRemoteMachineIP, IDS_NO_IPC_CONNECTION);
        
        //::AfxMessageBox(strFormattedErrorMsg);
		//::GetMainFrame()->ShowBalloonMsgInTray(_T("Warning"), strFormattedErrorMsg);
    }
    return bConnectionSuccess;
}


BOOL CRemoteAdministrator::EstablishAllConnections(CString strRemoteMachineIP, CString strPwd, BOOL bEstablish)
{
    BOOL bSuccessAdminConnection = EstablishAdminConnection(strRemoteMachineIP, strPwd, bEstablish);
    BOOL bSuccessIPCConnection   = EstablishIPCConnection(strRemoteMachineIP, strPwd, bEstablish);
    
    return (bSuccessAdminConnection && bSuccessIPCConnection);
}


BOOL CRemoteAdministrator::CopyServiceExeToRemoteMachine(CString strRemoteMachineIP)
{
   DWORD dwWritten = 0;

   HMODULE hInstance = ::GetModuleHandle(NULL);

   // Find the binary file in resources
   HRSRC hServiceExecutableRes = ::FindResource( 
                                    hInstance, 
                                    MAKEINTRESOURCE(IDR_REMOTEADMIN), 
                                    _T("EXECUTABLES")
                                    );

   HGLOBAL hServiceExecutable = ::LoadResource( 
                                    hInstance, 
                                    hServiceExecutableRes
                                    );

   LPVOID pServiceExecutable = ::LockResource(hServiceExecutable);

   if (pServiceExecutable == NULL)
      return FALSE;

   DWORD dwServiceExecutableSize = ::SizeofResource(
                                   hInstance,
                                   hServiceExecutableRes
                                   );

   TCHAR szServiceExePath[_MAX_PATH];

   ::sprintf(
       szServiceExePath, 
       _T("\\\\%s\\ADMIN$\\System32\\%s"), 
       strRemoteMachineIP.GetBuffer(0), 
       REMOTE_ADMIN_SERVICE_EXE
       );

   // Copy binary file from resources to \\remote\ADMIN$\System32
   HANDLE hFileServiceExecutable = ::CreateFile( 
                                        szServiceExePath,
                                        GENERIC_WRITE,
                                        0,
                                        NULL,
                                        CREATE_ALWAYS,
                                        FILE_ATTRIBUTE_NORMAL,
                                        NULL
                                        );

   if (hFileServiceExecutable == INVALID_HANDLE_VALUE)
   {
       return FALSE;
   }
         
   ::WriteFile(hFileServiceExecutable, pServiceExecutable, dwServiceExecutableSize, &dwWritten, NULL);

   ::CloseHandle(hFileServiceExecutable);
   
   return (dwWritten == dwServiceExecutableSize);
}


BOOL CRemoteAdministrator::ConnectToRemoteService(CString strRemoteMachineIP, DWORD dwRetry, DWORD dwRetryTimeOut)
{
    TCHAR szRemoteAdminPipeName[_MAX_PATH]               = _T("");
    TCHAR szRemoteAdminProcessInfoPipeName[_MAX_PATH]    = _T("");
    TCHAR szRemoteAdminProcessExecutePipeName[_MAX_PATH] = _T("");
    TCHAR szRemoteAdminProcessKillPipeName[_MAX_PATH]    = _T("");
    TCHAR szRemoteAdminSysShutdownPipe[_MAX_PATH]        = _T("");
    
    HANDLE hCommandPipe = INVALID_HANDLE_VALUE;

    // Remote service communication pipe name
    ::sprintf(
        szRemoteAdminPipeName, 
        _T("\\\\%s\\pipe\\%s"), 
        strRemoteMachineIP.GetBuffer(0), 
        REMOTE_ADMIN_PIPE
        );

    // Remote service communication pipe name
    ::sprintf(
        szRemoteAdminProcessInfoPipeName, 
        _T("\\\\%s\\pipe\\%s"), 
        strRemoteMachineIP.GetBuffer(0), 
        REMOTE_ADMIN_PROCESS_INFO_PIPE
        );

    // Remote service communication pipe name
    ::sprintf(
        szRemoteAdminProcessExecutePipeName, 
        _T("\\\\%s\\pipe\\%s"),