// ActiveXSurrogate.cpp : Implementation of cActiveXSurrogate
#include "stdafx.h"
#include "Decal.h"
#include "ActiveXSurrogate.h"

// Duplicated from Inject.h to prevent circular dependency
class IPluginSite;

class __declspec(uuid("{BA3E677F-8E44-4829-982E-58BBBC5C5F9B}")) IPlugin 
: public IUnknown
{
public:
   STDMETHOD(Initialize)(IPluginSite *pSite, long) = 0;
   STDMETHOD(Terminate)() = 0;
   STDMETHOD(get_FriendlyName)(BSTR *pstrName) = 0;
};

/////////////////////////////////////////////////////////////////////////////
// cActiveXSurrogate

typedef HRESULT (__stdcall *pfnDllRegisterServer)();
typedef HRESULT (__stdcall *pfnDllUnregisterServer)();

STDMETHODIMP cActiveXSurrogate::Register(BSTR strFilename)
{
   USES_CONVERSION;
   LPCTSTR szFilename = OLE2T( strFilename );

   // First, register the library
   HMODULE hLib = ::LoadLibrary( szFilename );
   if( hLib == NULL )
      return E_FAIL;

   pfnDllRegisterServer pfnReg = reinterpret_cast< pfnDllRegisterServer >( ::GetProcAddress( hLib, _T( "DllRegisterServer" ) ) );
   HRESULT hRes = E_FAIL;
   if( pfnReg != NULL )
      hRes = pfnReg();

   ::FreeLibrary( hLib );

   if( FAILED( hRes ) )
      return hRes;

   // Next open up the type library and look for suitable classes, they have
   // to be creatable - then we make one and see if it implements IPlugin2
   CComPtr< ITypeLib > pTypeLib;
   hRes = ::LoadTypeLib( strFilename, &pTypeLib );

   if( FAILED( hRes ) )
      return hRes;

   UINT nTypeInfoCount = pTypeLib->GetTypeInfoCount();
   int nFound = 0;

   for( UINT i = 0; i < nTypeInfoCount; ++ i )
   {
      TYPEKIND tk;
      pTypeLib->GetTypeInfoType( i, &tk );

      if( tk != TKIND_COCLASS )
         // Only coclasses
         continue;

      CComPtr< ITypeInfo > pClass1;
      pTypeLib->GetTypeInfo( i, &pClass1 );

      TYPEATTR *pta;
      pClass1->GetTypeAttr( &pta );

      if( pta->wTypeFlags & TYPEFLAG_FCANCREATE )
      {
         // Create an instance of the plugin and see if it implements IPlugin
         CComPtr< IUnknown > pUnkPlugin;
         HRESULT hRes = pClass1->CreateInstance ( NULL, IID_IUnknown, reinterpret_cast< LPVOID * > ( &pUnkPlugin ) );
         if ( SUCCEEDED ( hRes ) )
         {
            // Check for the IPlugin2 interface
            CComPtr< IPlugin2 > pPlugin2;
            if ( SUCCEEDED ( pUnkPlugin->QueryInterface ( &pPlugin2 ) ) )
            {
               // This is a plugin (anything that implements IPlugin2 is a plugin afterall, make the registry key
               LPCSTR szCLSID;
               LPOLESTR strCLSID;
               ::StringFromCLSID( pta->guid, &strCLSID );
               szCLSID = OLE2T( strCLSID );
               ::CoTaskMemFree( strCLSID );

               TCHAR szRegKey[ 255 ];
               ::_tcscat ( ::_tcscpy ( szRegKey, _T( "Software\\Decal\\Plugins\\" ) ), szCLSID );

               RegKey key;
               key.Create( HKEY_LOCAL_MACHINE, szRegKey );

               CComBSTR strDescription;
               pClass1->GetDocumentation( MEMBERID_NIL, NULL, &strDescription, NULL, NULL );

               key.SetStringValue( NULL, OLE2T( strDescription ) );
               key.SetStringValue( _T("File"), szFilename );
               key.SetDWORDValue( _T( "Enabled" ), 1 );
               key.SetStringValue( _T("Uninstall"), _T( "{7559F22F-C56F-4621-AE08-9C354D799D4B}" ) );

               ++ nFound;
            }
            else
            {
               CComPtr< IPlugin > pPlugin;
               if ( SUCCEEDED( pUnkPlugin->QueryInterface ( &pPlugin ) ) )
               {
                  // This is a V1 plugin - same registration, but use a surrogate
                  LPCSTR szCLSID;
                  LPOLESTR strCLSID;
                  ::StringFromCLSID( pta->guid, &strCLSID );
                  szCLSID = OLE2T( strCLSID );
                  ::CoTaskMemFree( strCLSID );

                  TCHAR szRegKey[ 255 ];
                  ::_tcscat ( ::_tcscpy ( szRegKey, _T( "Software\\Decal\\Plugins\\" ) ), szCLSID );

                  RegKey key;
                  key.Create( HKEY_LOCAL_MACHINE, szRegKey );

                  CComBSTR strDescription;
                  pPlugin->get_FriendlyName( &strDescription );

                  // Friendly name
                  key.SetStringValue( NULL, OLE2T( strDescription ) );
                  // File for uninstalling
                  key.SetStringValue( _T( "File" ), szFilename );
                  // Enabled by default
                  key.SetDWORDValue( _T( "Enabled" ), 1 );
                  // This object uninstalls
                  key.SetStringValue( _T("Uninstall"), _T( "{7559F22F-C56F-4621-AE08-9C354D799D4B}" ));
                  // The V1 surrogate
                  key.SetStringValue( _T( "Surrogate" ), _T( "{3D837F6E-B5CA-4604-885F-7AB45FCFA62A}" ) );

                  ++ nFound;
               }
            }
         }
      }

      pClass1->ReleaseTypeAttr( pta );
   }

   return ( nFound > 0 ) ? S_OK : E_FAIL;
}

STDMETHODIMP cActiveXSurrogate::Uninstall()
{
   USES_CONVERSION;
   TCHAR szParent[ 255 ];
   ::_tcscat ( ::_tcscpy ( szParent, _T( "Software\\Decal\\Plugins\\" ) ), OLE2T( m_strGroup ) );

   RegKey rk;
   if( rk.Open( HKEY_LOCAL_MACHINE, szParent ) != ERROR_SUCCESS )
   {
      _ASSERT( FALSE );
      return E_FAIL;
   }

   LPOLESTR strCLSID;
   ::StringFromCLSID( m_clsid, &strCLSID );

   LPTSTR szCLSID = OLE2T( strCLSID );

   ::CoTaskMemFree( strCLSID );

   if( rk.RecurseDeleteKey( szCLSID ) != ERROR_SUCCESS )
   {
      _ASSERT( FALSE );
      return E_FAIL;
   }

   // Load the library and call DLL unreg server
   LPCTSTR szFilename = OLE2T( m_strFile );

   // First, register the library
   HMODULE hLib = ::LoadLibrary( szFilename );
   if( hLib == NULL )
      return E_FAIL;

   pfnDllUnregisterServer pfnReg = reinterpret_cast< pfnDllUnregisterServer >( ::GetProcAddress( hLib, _T( "DllUnregisterServer" ) ) );
   HRESULT hRes = E_FAIL;
   if( pfnReg != NULL )
      hRes = pfnReg();

   ::FreeLibrary( hLib );

   // Just call this to make sure the library isn't locked so the user
   // can delete it after this point.
   ::CoFreeUnusedLibraries();

   return hRes;
}