Windows Kernel

Development Platform:
Visual C++

vrsreg.cpp：Code Content
							//***************************************************************************
//*     Copyright (c) Microsoft Corporation 1995. All rights reserved.      *
//***************************************************************************
//*                                                                         *
//* vrsreg.cpp - stub exe to register the virus scanner engine.             *
//*                                                                         *
//* created 10-15-96 inateeg                                                *
//*                                                                         *
//***************************************************************************
//***************************************************************************
//* INCLUDE FILES                                                           *
//***************************************************************************
#include <windows.h>
#include <windowsx.h>
#include <olectl.h>
#include <olectlid.h>
#include <initguid.h>
#include "vrsscan.h"
#include "resource.h"
#include "advpub.h"
//***************************************************************************
//* GLOBAL VARIABLES                                                        *
//***************************************************************************
BOOL    g_bRegister;
DWORD   g_dwFlags;
CLSID   g_clsid;
HINSTANCE g_hInst;
#define MAX_STRING      512
#define STR_OLESTR      1
#define ARRAYSIZE(a) (sizeof(a)/sizeof(a[0]))
// internal use as ANSI string only
#define VENDOR_REG      "CLSID\%s\VirusScanner"
#define SCANNER_COOKIE  "Cookie"
// predefined built in INF section name,  ANSI string required
#define SETUP_SECTION1  "InstallOCXs"
#define SETUP_SECTION2  "InstallRegs"
#define SETUP_UNINSTALL "Uninstall"
#define SETUP_TITLE     "Virus Scanner Setup"
// default macfee info
#define MACFEE_CLSID    TEXT("{91B0C1B0-0B71-11d0-8217-00A02474294C}")
#define DEFAULT_ENGINE  TEXT("McAfee")
BOOL Init( HINSTANCE hInstance, LPTSTR lpszCmdLine, INT nCmdShow );
BOOL ParseCmdLine( LPTSTR lpszCmdLine );
LPWSTR MakeWideStrFromAnsi( LPSTR psz, BYTE  bType );
LPSTR MakeAnsiStrFromWide( LPWSTR pwsz,  BYTE  bType );
void SaveRestoreCookie( CLSID clsid, DWORD *pdwCookie, BOOL bSave );
void ErrorMsg( HWND hWnd, UINT idErr );
HRESULT RegInstall( HINSTANCE hInst );
BOOL GetWideString( HINSTANCE hInstance, UINT id, LPWSTR *lplpwsStr );
//***************************************************************************
//*                                                                         *
//* NAME:       WinMain                                                     *
//*                                                                         *
//* SYNOPSIS:   Main entry point for the program.                           *
//*                                                                         *
//* REQUIRES:   hInstance:                                                  *
//*             hPrevInstance:                                              *
//*             lpszCmdLine:                                                *
//*             nCmdShow:                                                   *
//*                                                                         *
//* RETURNS:    int:                                                        *
//*                                                                         *
//***************************************************************************
INT WINAPI WinMain( HINSTANCE hInstance, HINSTANCE hPrevInstance,
                    LPTSTR lpszCmdLine, INT nCmdShow )
{
    HRESULT hr = E_FAIL;
    IRegisterVirusScanEngine *pIRegVirusScanner = NULL;
    DWORD   dwCookie = 0;
    LPWSTR  lpwsDesc = NULL;
    if ( Init( hInstance, lpszCmdLine, nCmdShow ) )
    {
        if ( !GetWideString( hInstance, IDS_VRSENG_DESC, &lpwsDesc ) )
            return hr;  //bail out
        CoInitialize( NULL );
        hr = CoCreateInstance( CLSID_VirusScan, NULL, CLSCTX_INPROC_SERVER, IID_IRegisterVirusScanEngine, (void**) &pIRegVirusScanner );
        if( SUCCEEDED(hr) )
        {
            if ( g_bRegister )
            {
                // Do register OCX and add Reg first
                if ( RegInstall( hInstance ) == S_OK )
                {
                    hr = pIRegVirusScanner->RegisterScanEngine( g_clsid, lpwsDesc, g_dwFlags, 0, &dwCookie );
                    SaveRestoreCookie( g_clsid, &dwCookie, TRUE );
                }
            }
            else
            {
                SaveRestoreCookie( g_clsid, &dwCookie, FALSE );
                pIRegVirusScanner->UnRegisterScanEngine( g_clsid, lpwsDesc, g_dwFlags, 0, dwCookie );
                // Do Unregister OCX and del Reg
                hr = RegInstall( hInstance );
            }
        }
        else
        {
             ErrorMsg( NULL, IDS_ERR_CREATINSTANCE );
        }
        CoUninitialize();
    }
    if ( lpwsDesc )
        CoTaskMemFree( lpwsDesc );
    return hr;
}
//***************************************************************************
//*                                                                         *
//* NAME:       Init                                                        *
//*                                                                         *
//* SYNOPSIS:   Initialization for the program is done here.                *
//*                                                                         *
//* REQUIRES:   hInstance:                                                  *
//*             hPrevInstance:                                              *
//*             lpszCmdLine:                                                *
//*             nCmdShow:                                                   *
//*                                                                         *
//* RETURNS:    BOOL:                                                       *
//*                                                                         *
//***************************************************************************
BOOL Init( HINSTANCE hInstance, LPTSTR lpszCmdLine, INT nCmdShow )
{
    TCHAR  szTmpBuf[MAX_PATH];
    LPWSTR pwszClsid;
    g_hInst = hInstance;
    HRESULT hr;
    g_dwFlags = 0;
    if ( !ParseCmdLine(lpszCmdLine) )
    {
        ErrorMsg( NULL, IDS_ERR_BADCMDLINE );
        return FALSE;
    }
    if ( LoadString( hInstance, IDS_VRSENG_CLSID, szTmpBuf, ARRAYSIZE(szTmpBuf) ) )
    {
#ifdef UNICODE
        pwszClsid = szTmpBuf;
#else
        pwszClsid = MakeWideStrFromAnsi( szTmpBuf, STR_OLESTR );
#endif
        hr = CLSIDFromString( pwszClsid, &g_clsid );
#ifndef UNICODE
        CoTaskMemFree( pwszClsid );
#endif
        if ( FAILED( hr) )
            return FALSE;
    }
    if ( lstrcmpi( szTmpBuf, MACFEE_CLSID ) == 0 )
           g_dwFlags |= SFV_DONTDOUI;
    return TRUE;
}
//***************************************************************************
//*                                                                         *
//*  ParseCmdLine()                                                     *
//*                                                                         *
//*  Purpose:    Parses the command line looking for switches               *
//*                                                                         *
//*  Parameters: LPSTR lpszCmdLineOrg - Original command line               *
//*                                                                         *
//*                                                                         *
//*  Return:     (BOOL) TRUE if successful                                  *
//*                     FALSE if an error occurs                            *
//*                                                                         *
//***************************************************************************
BOOL ParseCmdLine( LPTSTR lpszCmdLine )
{
    LPTSTR pArg;
    LPWSTR ;
    if( (!lpszCmdLine) || (lpszCmdLine[0] == 0) )
       return FALSE;
    pArg = strtok( lpszCmdLine, TEXT(" ") );
    while ( pArg )
    {
       if ( lstrcmpi( pArg, TEXT("/U") ) == 0 )
       {
           g_bRegister = FALSE;
       }
       else if ( lstrcmpi( pArg, TEXT("/R") ) == 0 )
       {
           g_bRegister = TRUE;
       }
       else
       {
           return FALSE;
       }
       pArg = strtok( NULL, TEXT(" ") );
    }
    return TRUE;
}
void SaveRestoreCookie( CLSID clsid, DWORD *pdwCookie, BOOL bSave )
{
    char    szBuf[MAX_PATH];
    LPOLESTR  pwsClsid;
    HKEY    hKey;
    LPSTR   pszClsid;
    DWORD   dwSize;
    StringFromCLSID( clsid, &pwsClsid );
    pszClsid = MakeAnsiStrFromWide( pwsClsid, STR_OLESTR );
    CoTaskMemFree( pwsClsid );
    wsprintf( szBuf, VENDOR_REG, pszClsid );
    if ( RegOpenKeyExA( HKEY_CLASSES_ROOT, szBuf, 0, KEY_ALL_ACCESS, &hKey ) == ERROR_SUCCESS )
    {
        if ( bSave )
        {
            RegSetValueExA( hKey, SCANNER_COOKIE, 0, REG_DWORD, (LPBYTE)pdwCookie, sizeof(DWORD) );
        }
        else
        {
            dwSize = sizeof(DWORD);
            *pdwCookie = 0;
            RegQueryValueExA( hKey, SCANNER_COOKIE, NULL, NULL, (LPBYTE)pdwCookie, &dwSize );
        }
        RegCloseKey( hKey );
    }
    CoTaskMemFree( pszClsid );
}
//******************************************************************************
//*  CreateInfFile - Create an INF file from an hmodule
//*
//*  ENTRY
//*      hInst - hmodule that contains the REGINST resource
//*      pszInfFileName - OUT the location to get the INF filename
//*      pszDir - OUT dir directory the INF file
//*
//*  EXIT
//*      HRESULT
//******************************************************************************
HRESULT CreateInfFile( HINSTANCE hInst, LPTSTR pszInfFileName, LPTSTR pszDir )
{
    HRESULT hr = E_FAIL;
    TCHAR szInfFilePath[MAX_PATH];
    LPVOID pvInfData;
    HRSRC hrsrcInfData;
    DWORD cbInfData, cbWritten;
    HANDLE hfileInf = INVALID_HANDLE_VALUE;
    if ( GetTempPath(ARRAYSIZE(szInfFilePath), szInfFilePath) > ARRAYSIZE(szInfFilePath))
    {
        goto Cleanup;
    }
    lstrcpy( pszDir, szInfFilePath );
    if (GetTempFileName(szInfFilePath, TEXT("RGI"), 0, pszInfFileName) == 0)
    {
        goto Cleanup;
    }
    hrsrcInfData = FindResource(hInst, TEXT("REGINST"), TEXT("REGINST"));
    if (hrsrcInfData == NULL)
    {
        goto Cleanup;
    }
    cbInfData = SizeofResource(hInst, hrsrcInfData);
    pvInfData = LockResource(LoadResource(hInst, hrsrcInfData));
    if (pvInfData == NULL)
    {
        goto Cleanup;
    }
    hfileInf = CreateFile(pszInfFileName, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS,
                          FILE_ATTRIBUTE_NORMAL, NULL);
    if (hfileInf == INVALID_HANDLE_VALUE)
    {
        goto Cleanup;
    }
    if ((WriteFile(hfileInf, pvInfData, cbInfData, &cbWritten, NULL) == FALSE) ||
        (cbWritten != cbInfData))
    {
        goto Cleanup;
    }
    hr = S_OK;
Cleanup:
    if (hfileInf != INVALID_HANDLE_VALUE)
    {
        CloseHandle(hfileInf);
    }
    return hr;
}
//********************************************************************************
//*  RegInstall - Install a registry INF
//*
//*
//********************************************************************************
HRESULT RegInstall( HINSTANCE hInst )
{
    HRESULT hr = E_FAIL;
    TCHAR szInfFileName[MAX_PATH];
    HINSTANCE hLib;
    TCHAR szPath[MAX_PATH];
    LPTSTR lpTmp;
    LPSTR  lpInfFile = NULL, lpDir = NULL;
    RUNSETUPCOMMAND pfRunSetupCommand;
    //
    // Load advpack.dll.
    //
    szPath[0] = 0;
    GetSystemDirectory( szPath, ARRAYSIZE(szPath) );
    lpTmp = CharPrev(szPath, szPath+lstrlen(szPath) );
    if ( (lpTmp > szPath) && (*lpTmp != TEXT('\') ) )
    {
       lstrcat( szPath, TEXT("\") );
    }
    lstrcat( szPath, TEXT("advpack.dll") );
    hLib = LoadLibrary( szPath );
    if ( !hLib )
    {
       ErrorMsg( NULL, IDS_ERR_LOADADV );
       hr = S_FALSE;
       goto Cleanup;
    }
    pfRunSetupCommand = (RUNSETUPCOMMAND)GetProcAddress( hLib, TEXT("RunSetupCommand") );
    if ( !pfRunSetupCommand )
    {
       ErrorMsg( NULL, IDS_ERR_GETPROCADD );
       hr = S_FALSE;
       goto Cleanup;
    }
    //
    // Create the INF file.
    //
    szInfFileName[0] = TEXT('');
    szPath[0] = TEXT('');
    hr = CreateInfFile(hInst, szInfFileName, szPath );
    if (FAILED(hr))
    {
        goto Cleanup;
    }
#ifdef UNICODE
    lpInfFile = MakeAnsiStrFromWide( szInfFileName, STR_OLESTR );
    lpDir =  MakeAnsiStrFromWide( szPath, STR_OLESTR );
#else
    lpInfFile = szInfFileName;
    lpDir = szPath;
#endif
    //
    // Execute the INF engine on the INF.
    //
    if ( g_bRegister )
    {
        hr = pfRunSetupCommand( NULL, lpInfFile, SETUP_SECTION1, lpDir,
                                SETUP_TITLE, NULL,
                                RSC_FLAG_INF|RSC_FLAG_QUIET|RSC_FLAG_NGCONV, NULL);
        if (FAILED(hr))
        {
            goto Cleanup;
        }
        hr = pfRunSetupCommand( NULL, lpInfFile, SETUP_SECTION2, lpDir,
                                SETUP_TITLE, NULL,
                                RSC_FLAG_INF|RSC_FLAG_QUIET|RSC_FLAG_NGCONV, NULL);
    }
    else
        hr = pfRunSetupCommand( NULL, lpInfFile, SETUP_UNINSTALL, lpDir,
                                SETUP_TITLE, NULL,
                                RSC_FLAG_INF|RSC_FLAG_QUIET|RSC_FLAG_NGCONV, NULL);
    if (FAILED(hr))
    {
        goto Cleanup;
    }
Cleanup:
    //
    // Delete the INF file, Free library.
    //
    if ( hLib )
        FreeLibrary( hLib );
    if (szInfFileName[0] )
    {
        DeleteFile(szInfFileName);
    }
#ifdef UNICODE
    if ( lpInfFile )
    {
        CoTaskMemFree( lpInfFile );     
    }
    if ( lpDir )
    {
        CoTaskMemFree( lpDir );     
    }
#endif
    return hr;
}
BOOL GetWideString( HINSTANCE hInstance, UINT id, LPWSTR *lplpwsStr )
{
    LPTSTR lpStr;
    
    lpStr = (LPTSTR)CoTaskMemAlloc( MAX_STRING );
    if ( !lpStr )
        return FALSE;
    if ( LoadString( hInstance, id, lpStr, MAX_STRING/sizeof(TCHAR) ) )
    {
#ifdef UNICODE
        *lplpwsStr = lpStr;
#else
        *lplpwsStr = MakeWideStrFromAnsi( lpStr, STR_OLESTR );
#endif
#ifndef UNICODE
        CoTaskMemFree( lpStr );
#endif
        if ( *lplpwsStr == 0 )
            return FALSE;
    }
    return TRUE;
}
//=--------------------------------------------------------------------------=
// given a string, make a BSTR out of it.
//
// Parameters:
//    LPSTR         - [in]
//    BYTE          - [in]
//
// Output:
//    LPWSTR        - needs to be cast to final desired result
//
// Notes:
//
LPWSTR MakeWideStrFromAnsi( LPSTR psz, BYTE  bType )
{
    LPWSTR pwsz;
    int i;
    // arg checking.
    //
    if (!psz)
        return NULL;
    // compute the length of the required BSTR
    //
    i =  MultiByteToWideChar(CP_ACP, 0, psz, -1, NULL, 0);
    if (i <= 0) return NULL;
    // allocate the widestr, +1 for terminating null
    //
    switch (bType) {
    case STR_OLESTR:
        pwsz = (LPWSTR) CoTaskMemAlloc(i * sizeof(WCHAR));
        break;
    default:
        return(NULL);
    }
    if (!pwsz) return NULL;
    MultiByteToWideChar(CP_ACP, 0, psz, -1, pwsz, i);
    pwsz[i - 1] = 0;
    return pwsz;
}
LPSTR MakeAnsiStrFromWide( LPWSTR pwsz,  BYTE  bType )
{
    LPSTR psz;
    int i;
    // arg checking.
    //
    if (!pwsz)
        return NULL;
    // compute the length of the required BSTR
    //
    i =  WideCharToMultiByte(CP_ACP, 0, pwsz, -1, NULL, 0, NULL, NULL);
    if (i <= 0) return NULL;
    // allocate the ansistr, +1 for terminating null
    //
    switch (bType) {
    case STR_OLESTR:
        psz = (LPSTR) CoTaskMemAlloc(i * sizeof(CHAR));
        break;
    default:
        return(NULL);
    }
    if (!psz) return NULL;
    WideCharToMultiByte(CP_ACP, 0, pwsz, -1, psz, i, NULL, NULL);
    psz[i - 1] = 0;
    return psz;
}
void ErrorMsg( HWND hWnd, UINT idErr )
{
    TCHAR szBuf[MAX_STRING];
    if ( !LoadString( g_hInst, idErr, szBuf, ARRAYSIZE(szBuf) ) )
    {
        wsprintf( szBuf, TEXT("LoadString failed on error id: %d"), idErr );
    }
    MessageBox( hWnd, szBuf, TEXT(SETUP_TITLE), MB_ICONSTOP|MB_OK );
}

						