// NetService.cpp : Implementation of cNetService
#include "stdafx.h"
#include "DecalNet.h"
#include "NetService.h"

#include <time.h>
#include "Message.h"
#include <ApiHook.h>
#include "ProtocolHook.h"
#include "FilterAdapterV1.h"


#define _OUTGOING_PACKETS_
//#define _LOGGING

extern DWORD HookCall (DWORD dwCallAddress, DWORD dwReplacement);
extern void ReceiveBlob();

bool	(__cdecl*LockMemoryPages)(void *);
bool	g_bLockMemoryPages = false;
DWORD	g_dwReceiveBlobReadyCall = 0;
DWORD	g_dwReceiveBlobReadyProc = 0;
DWORD	g_dwNetBlobObjectDataOffset = 0;
DWORD	g_dwNetBlobObjectSizeOffset = 0;

/////////////////////////////////////////////////////////////////////////////
// cNetService

// Functions for hooking from wsock32.dll
int PASCAL recvfromF( SOCKET s, char FAR* buf, int len, int flags, struct sockaddr FAR* from, int FAR* fromlen );
int PASCAL sendtoF( SOCKET s, const char FAR* buf, int len, int flags, const struct sockaddr FAR* to, int tolen );
static cHookDescriptor _hooks[] = {
	{ eByOrdinal, _T( "fsock32.dll" ), _T( "recvfrom" ), 17, reinterpret_cast< DWORD >( recvfromF ), 0 },
    { eByOrdinal, _T( "wsock32.dll" ), _T( "recvfrom" ), 17, reinterpret_cast< DWORD >( recvfromF ), 0 },
};
static cHookDescriptor _hooksOut[] = {
	{ eByOrdinal, _T( "wsock32.dll" ), _T( "sendto" ), 20, reinterpret_cast< DWORD >( sendtoF ), 0 },
};

HRESULT cNetService::onInitialize()
{
	if( g_pService == NULL )
	{
		g_pService = this;

		m_pDecal->get_Hooks( &m_pHooks );

		long lTemp = 0;

		m_pHooks->QueryMemLoc( _bstr_t( "ReceiveBlobReadyCall" ), &lTemp );
		g_dwReceiveBlobReadyCall = lTemp, lTemp = 0;

		m_pHooks->QueryMemLoc( _bstr_t( "NetBlobObjectDataOffset" ), &lTemp );
		g_dwNetBlobObjectDataOffset = lTemp, lTemp = 0;

		m_pHooks->QueryMemLoc( _bstr_t( "NetBlobObjectSizeOffset" ), &lTemp );
		g_dwNetBlobObjectSizeOffset = lTemp, lTemp = 0;


		m_pHooks->QueryMemLoc( _bstr_t( "LockMemoryPages" ), &lTemp );
		if( lTemp )
			g_bLockMemoryPages = true;
		LockMemoryPages = (bool(__cdecl*)(void*)) lTemp, lTemp = 0;



		if( g_dwReceiveBlobReadyCall && g_dwNetBlobObjectDataOffset && g_dwNetBlobObjectSizeOffset )
			g_dwReceiveBlobReadyProc = HookCall( g_dwReceiveBlobReadyCall, (DWORD) ReceiveBlob );
		else
		{
			g_dwReceiveBlobReadyCall = 0; // just in case.

			hookFunctions( _hooks, 2, true );
			if( _hooks[ 0 ].m_pOldFunction != 0 )
				g_fn_recvfrom = reinterpret_cast< fn_recvfrom >( _hooks[ 0 ].m_pOldFunction );
			else if( _hooks[ 1 ].m_pOldFunction != 0 )
				g_fn_recvfrom = reinterpret_cast< fn_recvfrom >( _hooks[ 1 ].m_pOldFunction );
			else
				_ASSERTE( false );
		}
	}

#ifdef _OUTGOING_PACKETS_
		hookFunctions( _hooksOut, 1, true );

		if( _hooksOut[ 0 ].m_pOldFunction != 0 )
			g_fn_sendto = reinterpret_cast< fn_sendto >( _hooksOut[ 0 ].m_pOldFunction );
		else
			_ASSERTE( false );
#endif

	m_stack.start( this );

   // Start all of the network filters
   {
      CComPtr< IDecalEnum > pEnum;
      HRESULT hRes = m_pDecal->get_Configuration( _bstr_t( _T( "NetworkFilters" ) ), GUID_NULL, &pEnum );
      if( FAILED( hRes ) )
         return hRes;

      while( pEnum->Next() == S_OK )
      {
         // Unlike services, network filters are allowed to fail
         VARIANT_BOOL bEnabled;
         HRESULT hRes = pEnum->get_Enabled ( &bEnabled );
         _ASSERTE ( SUCCEEDED ( hRes ) );

         if ( !bEnabled )
            continue;

         cFilter f;
         pEnum->get_ComClass( &f.m_clsid );

         hRes = pEnum->CreateInstance( __uuidof( INetworkFilter2 ), reinterpret_cast< void ** >( &f.m_p ) );
         if( FAILED( hRes ) )
         {
            // Heads up that a filter is failing to init
            _ASSERT( FALSE );
            continue;
         }

         hRes = f.m_p->Initialize( this );

         if( SUCCEEDED( hRes ) )
         {
            f.m_ver = eVersion2;
            m_filters.push_back( f );
         }
      }
   }

   // Prepare the message container
   CComObject< cMessage > *pMessage;
   CComObject< cMessage >::CreateInstance( &pMessage );

   pMessage->AddRef();
   pMessage->m_pService = this;
   pMessage->init();

   m_pMessage = pMessage;

   return S_OK;
}

void cNetService::onTerminate()
{
	m_pMessage->term();

	// The release should destroy it
	m_pMessage->Release();
	m_pMessage = NULL;

	// Kill all of the network filters in reverse order
	while( !m_filters.empty() )
	{
		cFilterList::iterator i_end = ( -- m_filters.end() );
		i_end->m_p->Terminate();
		m_filters.erase( i_end );
	}

	m_filters.clear();

	m_stack.stop();

	if( g_pService == this )
	{
#ifdef _OUTGOING_PACKETS_
		hookFunctions( _hooksOut, 1, false );
#endif
		if( g_dwReceiveBlobReadyCall )
			HookCall( g_dwReceiveBlobReadyCall, g_dwReceiveBlobReadyProc );
		else
			hookFunctions( _hooks, 2, false );

		g_pService = NULL;
	}

	m_pHooks.Release();
}

int PASCAL recvfromF( SOCKET s, char FAR* buf, int len, int flags, struct sockaddr FAR* from, int FAR* fromlen )
{
	_ASSERTE( cNetService::g_pService != NULL );

	int iRes = cNetService::g_fn_recvfrom( s, buf, len, flags, from, fromlen );

#ifdef _LOGGING_
	if( iRes != SOCKET_ERROR && cNetService::g_pService != NULL )
	{
		char *harro = new char[256];
		_snprintf(harro, 256, "recvfrom( %d, %08X, %d, %d, %08X, %d ) = %d", s, buf, len, flags, from, fromlen, iRes);
		m_pHooks->ChatOut(_bstr_t (harro), 0);
		delete[] harro;
	}
#endif

	if( iRes != SOCKET_ERROR && cNetService::g_pService != NULL )
		cNetService::g_pService->m_stack.processPacket( static_cast< DWORD >( iRes ), reinterpret_cast< BYTE * >( buf ) );

	return iRes;
}

int PASCAL sendtoF( SOCKET s, const char FAR* buf, int len, int flags, const struct sockaddr FAR* to, int tolen )
{
   _ASSERTE( cNetService::g_pService != NULL );

#ifdef _LOGGING_
   char *harro = new char[256];
   char *buffer = new char[512];
   sprintf(harro, "sendto( %d, %08X, %d, %d, %s:%d, %d )",
	   s, buf, len, flags,
	   inet_ntoa( ((sockaddr_in *)(to))->sin_addr ),
	   ntohs( ((sockaddr_in *)(to))->sin_port ),
	   tolen);
   m_pHooks->ChatOut(_bstr_t (harro), 0);
   delete[] buffer;
   delete[] harro;
#endif

   int iRes = cNetService::g_fn_sendto( s, buf, len, flags, to, tolen );

   if( iRes != SOCKET_ERROR && cNetService::g_pService != NULL )
      cNetService::g_pService->m_stack.processPacketOut( static_cast< DWORD >( len ), reinterpret_cast< const BYTE * >( buf ) );

	return iRes;
}

cNetService *cNetService::g_pService = NULL;
cNetService::fn_recvfrom cNetService::g_fn_recvfrom = NULL;
cNetService::fn_sendto cNetService::g_fn_sendto = NULL;

void cNetService::onMessage( ACMessage &msg )
{
	_ASSERTE( m_pMessage != NULL );

	DWORD dwMessageCode = msg.getType();
	if( dwMessageCode == 0xF7C7 )
		m_pDecal->StartPlugins();

	else if(( dwMessageCode == 0xF653 ) || ( dwMessageCode == 0xF659 ))
		m_pDecal->StopPlugins();

	m_pMessage->crackMessage( msg.getData (), msg.getSize () );

//	if( g_bLockMemoryPages )
//		if( !LockMemoryPages( reinterpret_cast< void * >(m_pMessage) ) )
//			return;

	// Dispatch the message to all of the network filters
	for( cFilterList::iterator i = m_filters.begin(); i != m_filters.end(); ++ i )
		i->m_p->DispatchServer( m_pMessage );
}

void cNetService::onMessageOut( ACMessage &msg )
{
   _ASSERTE( m_pMessage != NULL );

   DWORD dwMessageCode = msg.getType();

   m_pMessage->crackMessage( msg.getData (), msg.getSize () );

   // Dispatch the message to all of the network filters
   for( cFilterList::iterator i = m_filters.begin(); i != m_filters.end(); ++ i )
      i->m_p->DispatchClient( m_pMessage );
}

STDMETHODIMP cNetService::BeforePlugins()
{
   USES_CONVERSION;

   // Startup version 1 network filters
   RegKey hkeyNF;
   if ( hkeyNF.Open ( HKEY_LOCAL_MACHINE, _T( "Software\\Decal\\NetworkFilters" ), KEY_READ ) == ERROR_SUCCESS )
   {
      // Enum the values
      DWORD dwSize = 64;
      TCHAR szCLSID[ 64 ];

      for ( int i = 0; ::RegEnumValue ( hkeyNF, i, szCLSID, &dwSize, NULL, NULL, NULL, NULL ) == ERROR_SUCCESS; ++ i )
      {
         // Reset the sizes
         dwSize = 64;

         // Try to convert the string
         cFilter f;
         HRESULT hRes = ::CLSIDFromString ( T2OLE ( szCLSID ), &f.m_clsid );
         if ( FAILED ( hRes ) )
            // This is not a CLSID
            continue;

         // Check for enablement
         DWORD dwEnabled;
         hkeyNF.QueryDWORDValue ( szCLSID, dwEnabled );
         if ( !dwEnabled )
            continue;

         // We have a CLSID - try to create the object and the adapter
         CComObject< cFilterAdapterV1 > *pAdapter;
         hRes = CComObject< cFilterAdapterV1 >::CreateInstance ( &pAdapter );

         if ( FAILED ( hRes ) )
         {
            _ASSERT ( FALSE );
            continue;
         }

         f.m_p = pAdapter;
         hRes = pAdapter->init ( f.m_clsid );
         if( SUCCEEDED( hRes ) )
         {
            f.m_ver = eVersion1;
            m_filters.push_back( f );
         }
      }
   }

   return S_OK;
}

STDMETHODIMP cNetService::AfterPlugins()
{
   // Remove all version 1 network filters
   for ( cFilterList::iterator i = m_filters.begin (); i != m_filters.end (); )
   {
      if ( i->m_ver == eVersion1 )
         i = m_filters.erase ( i );
      else
         ++ i;
   }

   return S_OK;
}

STDMETHODIMP cNetService::Lookup( BSTR strName, IUnknown **ppvItf )
{
   // Attempt to convert the name to a GUID
   CLSID clsid;
   HRESULT hRes;
   if( strName[ 0 ] == OLECHAR( '{' ) )
      hRes = ::CLSIDFromString( strName, &clsid );
   else
      hRes = ::CLSIDFromProgID( strName, &clsid );

   if( FAILED( hRes ) )
   {
      _ASSERT( FALSE );
      return E_INVALIDARG;
   }

   // Walk through the list of filters to find our match
   for( cFilterList::iterator i = m_filters.begin(); i != m_filters.end(); ++ i )
   {
      if( i->m_clsid == clsid )
         return i->m_p->QueryInterface( IID_IUnknown, reinterpret_cast< void ** >( ppvItf ) );
   }

   // Could not find a network filter with that class ID
   _ASSERT( FALSE );
   return E_INVALIDARG;
}

STDMETHODIMP cNetService::get_Decal(IDecal **pVal)
{
	return m_pDecal->QueryInterface( pVal );
}

STDMETHODIMP cNetService::get_Filter(REFCLSID clsid, REFIID iid, LPVOID *pVal)
{
   for( cFilterList::iterator i = m_filters.begin(); i != m_filters.end(); ++ i )
   {
      if( i->m_clsid == clsid )
         return i->m_p->QueryInterface( iid, pVal );
   }

	return E_INVALIDARG;
}

STDMETHODIMP cNetService::get_FilterVB(BSTR strProgID, LPDISPATCH *pVal)
{
   _ASSERTE( strProgID != NULL );
   _ASSERTE( pVal != NULL );

   // Prepend a fully-qualified service id
   _bstr_t strPath ( _T( "services\\DecalNet.NetService\\" ) );
   strPath += strProgID;

   return m_pDecal->get_Object ( strPath, IID_IDispatch, reinterpret_cast< LPVOID * > ( pVal ) );
}

STDMETHODIMP cNetService::get_Hooks(IACHooks **pVal)
{
	return m_pHooks->QueryInterface( pVal );
}

void __stdcall OnReceiveBlob (DWORD *pStructure)
{
	LPBYTE pPacket = (LPBYTE) (pStructure [g_dwNetBlobObjectDataOffset]);
	DWORD dwSize = pStructure [g_dwNetBlobObjectSizeOffset];
	if (cNetService::g_pService)
	{
		HookedMessage msg (pPacket, dwSize);
		cNetService::g_pService->onMessage (msg);
	}
}

void __declspec (naked) ReceiveBlob()
{
	_asm
	{
		push	ecx

		push	ecx
		call	OnReceiveBlob

		pop		ecx

		jmp		g_dwReceiveBlobReadyProc
	}
}

// Returns the function which previously was being called, after replacing it with the new call.
DWORD HookCall (DWORD dwCallAddress, DWORD dwReplacement)
{
	DWORD* pTemp = (DWORD *) (dwCallAddress + 1);

	HANDLE hProcess = OpenProcess
		(
		PROCESS_VM_WRITE | PROCESS_VM_OPERATION,
		FALSE,
		GetCurrentProcessId ()
		);

	DWORD dwOriginal = 0;

	if (hProcess)
	{
		dwOriginal = (*pTemp) + dwCallAddress + 5;

		DWORD dwTemp = dwReplacement - (dwCallAddress + 5);
		if (WriteProcessMemory
			(
			hProcess,
			pTemp,
			&dwTemp,
			sizeof (DWORD),
			NULL
			))
		{
		}
		else
		{
			dwOriginal = 0;
		}

		CloseHandle (hProcess);
	}

	return dwOriginal ;
}

