///////////////////////////////////////////////////////////////
//  registry.cpp - create registry entries for server        //
//                                                           //
//  Platform:      Micron Dual Pentium 200, WinNT 4.0, SP3   //
//  Application:   Demonstration for CSE791 - Distrib. Objs  //
//  Authors:       Eddon & Eddon, Inside Distributed COM,    //
//                 Microsoft Press, 1998                     //
//  Modified:      Jim Fawcett, Syracuse Univ, CST 2-187     //
//                 (315) 443-3948, fawcett@ecs.syr.edu       //
///////////////////////////////////////////////////////////////

#include <objbase.h>
#include <assert.h>
#include "registry.h"

// Forward declarations for private functions
//
// Set the given key and its value.
BOOL setKeyAndValue(
       const char* pszPath, 
       const char* szSubkey, 
       const char* szValue
     );
//
// Open a key and set a value.
BOOL setValueInKey(
       const char* szKey, 
       const char* szNamedValue, 
       const char* szValue
     );
//
// Convert a CLSID into a char string.
void CLSIDtochar(REFCLSID clsid, char* szCLSID, int length);
//
// Delete szKeyChild and all of its descendents.
LONG recursiveDeleteKey(HKEY hKeyParent, const char* szKeyChild);
//
// Size of a CLSID as a string
const int CLSID_STRING_SIZE = 39;

//
//----< Register the component in the registry. >--------------

HRESULT RegisterServer(
          const char* szModuleName,     // DLL module handle
          REFCLSID clsid,               // Class ID
          const char* szFriendlyName,   // Friendly Name
          const char* szVerIndProgID,   // Programmatic
          const char* szProgID,         // IDs
					const char* szThreadingModel  // ThreadingModel
        ) {
  // Get server location.
  char szModule[512];
  HMODULE hModule = GetModuleHandle(szModuleName);
  DWORD dwResult = GetModuleFileName(
                     hModule, szModule, sizeof(szModule)/sizeof(char)
                   );
  assert(dwResult != 0);

  // Convert the CLSID into a char.
  char szCLSID[CLSID_STRING_SIZE];
  CLSIDtochar(clsid, szCLSID, sizeof(szCLSID));

  // Build the key CLSID\\{...}
  char szKey[64];
  strcpy(szKey, "CLSID\\"); strcat(szKey, szCLSID);
  
  // Add the CLSID to the registry.
  setKeyAndValue(szKey, NULL, szFriendlyName);

  // Add the server filename subkey under the CLSID key.
  if(strstr(szModuleName, ".exe") == NULL) {
    setKeyAndValue(szKey, "InprocServer32", szModule);
    char szInproc[64];
    strcpy(szInproc, szKey); strcat(szInproc, "\\InprocServer32");
    setValueInKey(szInproc, "ThreadingModel", szThreadingModel);
  }
  else
    setKeyAndValue(szKey, "LocalServer32", szModule);

  // Add the ProgID subkey under the CLSID key.
  setKeyAndValue(szKey, "ProgID", szProgID);

  // Add the version-independent ProgID subkey under CLSID key.
  setKeyAndValue(szKey, "VersionIndependentProgID", szVerIndProgID);

  // Add the version-independent ProgID subkey under HKEY_CLASSES_ROOT.
  setKeyAndValue(szVerIndProgID, NULL, szFriendlyName); 
  setKeyAndValue(szVerIndProgID, "CLSID", szCLSID);
  setKeyAndValue(szVerIndProgID, "CurVer", szProgID);

  // Add the versioned ProgID subkey under HKEY_CLASSES_ROOT.
  setKeyAndValue(szProgID, NULL, szFriendlyName); 
  setKeyAndValue(szProgID, "CLSID", szCLSID);
  return S_OK;
}
//

//----< Remove component from the registry. >------------------

LONG UnregisterServer(REFCLSID clsid,             // Class ID
                      const char* szVerIndProgID, // Programmatic
                      const char* szProgID)       // IDs
{
  // Convert the CLSID into a char.
  char szCLSID[CLSID_STRING_SIZE];
  CLSIDtochar(clsid, szCLSID, sizeof(szCLSID));

  // Build the key CLSID\\{...}
  char szKey[64];
  strcpy(szKey, "CLSID\\");
  strcat(szKey, szCLSID);

  // Delete the CLSID Key - CLSID\{...}
  LONG lResult = recursiveDeleteKey(HKEY_CLASSES_ROOT, szKey);
  assert((lResult == ERROR_SUCCESS) || 
         (lResult == ERROR_FILE_NOT_FOUND)); // Subkey may not exist.

  // Delete the version-independent ProgID Key.
  lResult = recursiveDeleteKey(HKEY_CLASSES_ROOT, szVerIndProgID);
  assert((lResult == ERROR_SUCCESS) || 
         (lResult == ERROR_FILE_NOT_FOUND)); // Subkey may not exist.

  // Delete the ProgID key.
  lResult = recursiveDeleteKey(HKEY_CLASSES_ROOT, szProgID);
  assert((lResult == ERROR_SUCCESS) || 
         (lResult == ERROR_FILE_NOT_FOUND)); // Subkey may not exist.

  return S_OK;
}

//----< Convert CLSID to char string. >------------------------

void CLSIDtochar(REFCLSID clsid, char* szCLSID, int length) {

  assert(length >= CLSID_STRING_SIZE);
  // Get CLSID
  LPOLESTR wszCLSID = NULL;
  HRESULT hr = StringFromCLSID(clsid, &wszCLSID);
  assert(SUCCEEDED(hr));

  // Covert from wide characters to non-wide.
  wcstombs(szCLSID, wszCLSID, length);

  // Free memory.
  CoTaskMemFree(wszCLSID);
}

//

//----< Delete a key and all of its descendents. >-------------

LONG recursiveDeleteKey(
       HKEY hKeyParent,           // Parent of key to delete
       const char* lpszKeyChild   // Key to delete
     )
{
  // Open the child.
  HKEY hKeyChild;
  LONG lRes = RegOpenKeyEx(
                hKeyParent, lpszKeyChild, 0, KEY_ALL_ACCESS, &hKeyChild
              );
  if(lRes != ERROR_SUCCESS)
    return lRes;

  // Enumerate all of the decendents of this child.
  FILETIME time;
  char szBuffer[256];
  DWORD dwSize = 256;
  while(
         RegEnumKeyEx(hKeyChild,0,szBuffer,&dwSize,NULL,NULL,NULL,&time)
           == S_OK
       )
  {
    // Delete the decendents of this child.
    lRes = recursiveDeleteKey(hKeyChild, szBuffer);
    if(lRes != ERROR_SUCCESS)
    {
      // Cleanup before exiting.
      RegCloseKey(hKeyChild);
      return lRes;
    }
    dwSize = 256;
  }

  // Close the child.
  RegCloseKey(hKeyChild);

  // Delete this child.
  return RegDeleteKey(hKeyParent, lpszKeyChild);
}

//

//----< Create a key and set its value. >----------------------

BOOL setKeyAndValue(
       const char* szKey, const char* szSubkey, const char* szValue
     )
{
  HKEY hKey;
  char szKeyBuf[1024];

  // Copy keyname into buffer.
  strcpy(szKeyBuf, szKey);

  // Add subkey name to buffer.
  if(szSubkey != NULL) {

    strcat(szKeyBuf, "\\");
    strcat(szKeyBuf, szSubkey );
  }

  // Create and open key and subkey.
  long lResult = RegCreateKeyEx(
                   HKEY_CLASSES_ROOT, szKeyBuf, 0, NULL, 
                   REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, 
                   NULL, &hKey, NULL
                 );
  if(lResult != ERROR_SUCCESS)
    return FALSE;

  // Set the Value.
  if(szValue != NULL)
    RegSetValueEx(
      hKey, NULL, 0, REG_SZ, (BYTE *)szValue, strlen(szValue)+1
    );

  RegCloseKey(hKey);
  return TRUE;
}

//

//----< Open a key and set a value. >--------------------------

BOOL setValueInKey(
       const char* szKey, const char* szNamedValue, const char* szValue
     )
{
  HKEY hKey;
  char szKeyBuf[1024];

  // Copy keyname into buffer.
  strcpy(szKeyBuf, szKey);

  // Create and open key and subkey.
  long lResult = RegOpenKeyEx(
                   HKEY_CLASSES_ROOT, szKeyBuf, 0, KEY_SET_VALUE, &hKey
                 );
  if(lResult != ERROR_SUCCESS)
    return FALSE;

  // Set the Value.
  if(szValue != NULL)
    RegSetValueEx(
      hKey, szNamedValue, 0, REG_SZ, (BYTE*)szValue, strlen(szValue)+1
    );

  RegCloseKey(hKey);
  return TRUE;
}