// ==========================================================================
// Author:  Yee Hsu
// Date:    6/8/2007
//
// Desc:    Service Enumerator. Lists all the Windows services that are
//          currently running.
// ==========================================================================

#include "stdafx.h"
#include "CRegKey.h"
#include "MD5.h"
#include "IntCommon.h"

bool GetServList(lServiceList& lServices)
{
    DWORD dwBytesNeeded         = 0;
    DWORD dwServicesReturned    = 0;
    DWORD dwResumedHandle       = 0;
    ENUM_SERVICE_STATUS service;

    ServiceBlock ServData;

    // Opens a handle to the service manager
    SC_HANDLE hHandle = OpenSCManager(NULL, NULL, SC_MANAGER_ENUMERATE_SERVICE);

    if (!hHandle)
        return false;

    // Query services
    BOOL retVal = EnumServicesStatus(hHandle, SERVICE_WIN32 | SERVICE_DRIVER, SERVICE_STATE_ALL, 
        &service, sizeof(ENUM_SERVICE_STATUS), &dwBytesNeeded, &dwServicesReturned,
        &dwResumedHandle);

    if (!retVal)
    {
        // Set the buffer
        DWORD dwBytes = sizeof(ENUM_SERVICE_STATUS) + dwBytesNeeded;
        ENUM_SERVICE_STATUS* pServices = NULL;
        pServices = new ENUM_SERVICE_STATUS [dwBytes];

        // Now query again for services
        EnumServicesStatus(hHandle, SERVICE_WIN32 | SERVICE_DRIVER, SERVICE_STATE_ALL, 
            pServices, dwBytes, &dwBytesNeeded, &dwServicesReturned, &dwResumedHandle);

        // now traverse and save each services in a structure
        for (unsigned iIndex = 0; iIndex < dwServicesReturned; iIndex++)
        {
            ServData.sDisplayName = (pServices + iIndex)->lpDisplayName;
            ServData.sServiceName = (pServices + iIndex)->lpServiceName;
            ServData.sServiceStat = GetCurrentStatus((pServices + iIndex)->ServiceStatus.dwCurrentState);
            ServData.sServiceType = GetTypeOfService((pServices + iIndex)->ServiceStatus.dwServiceType);
            
            GetServiceMD5andPath(ServData);
            lServices.push_back(ServData);

            //////////////////////////////////////////////////////////////////// FLUSH TO STDOUT for real time parsing
            fprintf(stderr, "%s\n", ServData.sServiceName.c_str());
            //////////////////////////////////////////////////////////////////////////////////////////////////////////    
        }
    }
    CloseServiceHandle(hHandle);

    return true;
}

void GetServiceMD5andPath(ServiceBlock& ServData)
{
    TCHAR strWindowsPath[MAX_PATH+1];
    std::string service_reg_path;
    CRegKey test;
    char str_path[1200];
    unsigned long str_path_size=sizeof(str_path);
    std::string objImagePath;

    service_reg_path =    "SYSTEM\\CurrentControlSet\\Services\\";
    service_reg_path +=    ServData.sServiceName;

    if(    test.Open(HKEY_LOCAL_MACHINE,(LPCTSTR)service_reg_path.c_str(),KEY_QUERY_VALUE|KEY_ENUMERATE_SUB_KEYS) == ERROR_SUCCESS)
    {
        if ( FindServiceImagePath( test.m_hKey, objImagePath ) )
        {
            size_t j = objImagePath.find('"');
            while ( j != -1 )
            {
                objImagePath = objImagePath.erase(j,1);
                j = objImagePath.find('"');
            }

            std::string::size_type pos = objImagePath.find("\\??\\");
            if( pos != std::string::npos )
            {
                objImagePath = objImagePath.erase(pos,4);
            }

            if ( 
                    ( !strcmp(objImagePath.substr(0,sizeof("System32\\")-1).c_str(),"System32\\") )  ||
                    ( !strcmp(objImagePath.substr(0,sizeof("SYSTEM32\\")-1).c_str(),"SYSTEM32\\") )  ||
                    ( !strcmp(objImagePath.substr(0,sizeof("system32\\")-1).c_str(),"system32\\") )
                )
            {
                if ( GetWindowsDirectory( strWindowsPath, MAX_PATH ) )
                {
                    objImagePath.insert(0,"\\");
                    objImagePath.insert(0,strWindowsPath);
                }
            }
            else
                if ( !strcmp(objImagePath.substr(0,sizeof("\\SystemRoot\\")-1).c_str(),"\\SystemRoot\\") )
                {
                    if ( GetWindowsDirectory( strWindowsPath, MAX_PATH ) )
                    {
                        objImagePath = objImagePath.substr( sizeof("\\SystemRoot")-1 );
                        objImagePath.insert(0,strWindowsPath);
                    }
                }

            size_t n;
            if ( ( (n = objImagePath.find(".exe")) != -1 ) || ( (n = objImagePath.find(".EXE")) != -1) )
            {
                if ( ( n + 4 ) !=  objImagePath.size() )
                    objImagePath = objImagePath.substr(0,n+4);

            }
        }
        else
            objImagePath = "";    

//            test.Close();

        if( objImagePath.empty() )
            ServData.sMD5 = "";
        else
        {
            ServData.sServicePath = RemovePathTaggingFileName(objImagePath);
            ServData.sServiceProg = objImagePath.substr(objImagePath.rfind("\\")+1, objImagePath.npos);
            ServData.sServicePath = ConvertPathToShortPath(ServData.sServicePath);

            //Get MD5 Signature
            char chBuff[256]; 
            memset( chBuff, 0, sizeof(chBuff) );              
            if( GetMD5Signature( objImagePath.c_str(), chBuff, sizeof(chBuff) ) )
            {
                ServData.sMD5 = chBuff;
            }
            else
                ServData.sMD5 = "";
        }
    }
    else
        ServData.sMD5 = "";
}

bool FindServiceImagePath( HKEY hSrvKey, std::string& objImagePath )
{
    // Futher code optimization recommended in that function.
    CRegKey                                   objRegInterface;
    LONG                                   numFuncRes;  
    DWORD                                   numInd;
    TCHAR                                   strKeyName[MAX_PATH]; // Declare as TCHAR - call to API
    DWORD                                   numKeySize;

    // go recursion for every sub-key

    // check gor ImagePath in this key
    if ( GetValueAsString( hSrvKey,"ImagePath",objImagePath) )
        return true;

     for ( numInd = 0, numFuncRes = ERROR_SUCCESS; numFuncRes == ERROR_SUCCESS; numInd++ ) 
         if ( ( numFuncRes = RegEnumKeyEx( hSrvKey,
                                           numInd, 
                                           strKeyName, 
                                           &(numKeySize = MAX_PATH),
                                           NULL, 
                                           NULL, 
                                           NULL, 
                                           NULL ) ) == ERROR_SUCCESS )
         {
            if ( objRegInterface.Open( hSrvKey, ( LPCTSTR ) strKeyName, KEY_ENUMERATE_SUB_KEYS | KEY_QUERY_VALUE ) == ERROR_SUCCESS )
            {
                // Go recursion for that subkey
                return FindServiceImagePath( objRegInterface.m_hKey, objImagePath );
            }
            else
                return false;
         }
         return false;
}

bool GetValueAsString( HKEY hRegKey, const _TCHAR* strValueName, std::string& objValueData )
{
    LONG                    numFuncRes;  
    DWORD                    numBuffLen;
    LPBYTE                    pBuff;
    bool                    numResult = false;
    DWORD                    numType;
    _TCHAR*                 strItoABuffer = NULL;
    DWORD                   numExpStringSize;

    // Querry value type
    if( ( numFuncRes = RegQueryValueEx( hRegKey, ( LPCTSTR) strValueName, 0, &numType, NULL , &numBuffLen ) ) == ERROR_SUCCESS ) 
    {
        // Allocate memory
        if ( ( pBuff = new BYTE[numBuffLen] ) != NULL ) 
        {
            // Read data
            if ( ( numFuncRes = RegQueryValueEx( hRegKey, ( LPCTSTR) strValueName, 0, NULL, ( LPBYTE ) pBuff , &numBuffLen ) ) == ERROR_SUCCESS ) 
                // Determine the need of convertion
                switch( numType )
                {
                    case REG_DWORD:
                        // Convert!
                        if ( ( strItoABuffer = new _TCHAR[11] ) != NULL ) 
                        {
                            objValueData = itoa( *( ( LPDWORD ) pBuff ), ( _TCHAR* ) strItoABuffer, 10 );
                            numResult = true;
                        }
                        break;

                    case REG_EXPAND_SZ:
                        if ( ( numExpStringSize = ExpandEnvironmentStrings((LPCSTR)pBuff,NULL,0) ) )
                            if (  ( strItoABuffer = new _TCHAR[numExpStringSize] ) != NULL )
                                if ( ExpandEnvironmentStrings((LPCSTR) pBuff, strItoABuffer, numExpStringSize ) )
                                {
                                   objValueData = strItoABuffer;
                                     numResult = true;
                                }
                        break;


                    case REG_SZ:
                        // Read directly
                        objValueData = ( _TCHAR* ) pBuff;
                        // Indicate success
                        numResult = true;
                    break;

                    default:
                        // Not supported value type is being read
                        numResult = false;
                }
            // Clean up
            delete[] pBuff;
            if ( strItoABuffer )
                delete[] strItoABuffer;
        }
    }
    return numResult;
}

bool GetMD5Signature( const char* filename, char* pszMD5Buff, int iMD5BuffLen )
{
    char chBuff[256];
    memset( chBuff, 0, sizeof(chBuff) );

    if( MD5File( filename, chBuff, sizeof(chBuff) ) )
    {
        return false;
    }
    else
    {
        strncpy(pszMD5Buff, chBuff, iMD5BuffLen);
        return true;
    }
}

// get current status of the services
string GetCurrentStatus(const DWORD p_dwType)
{
    string sStatus = "";

    switch (p_dwType)
    {
        case SERVICE_STOPPED:               sStatus = "Off";                                    break;
        case SERVICE_START_PENDING:         sStatus = "Starting";                               break;
        case SERVICE_STOP_PENDING:          sStatus = "Stopping";                               break;
        case SERVICE_RUNNING:               sStatus = "On";                                     break;
        case SERVICE_CONTINUE_PENDING:      sStatus = "Continue is pending";                    break;
        case SERVICE_PAUSE_PENDING:         sStatus = "Pause is pending";                       break;
        case SERVICE_PAUSED:                sStatus = "Paused";                                 break;
    }
    return sStatus;
}

// get type of the service
string GetTypeOfService(const DWORD p_dwType)
{
    string sType = "";

    switch (p_dwType)
    {    
        case SERVICE_WIN32_OWN_PROCESS:     sType = "Run its own process";                      break;
        case SERVICE_WIN32_SHARE_PROCESS:   sType = "Share a process with other application";   break;
        case SERVICE_KERNEL_DRIVER:         sType = "Device driver";                            break;
        case SERVICE_FILE_SYSTEM_DRIVER:    sType = "File system driver";                       break;
        case SERVICE_INTERACTIVE_PROCESS:   sType = "Service can interact with desktop";        break;
    }
    return sType;
}