#include "stdafx.h"
#include "Pub.h"
#include "Wininet.h"
#include "ExDispID.h"
#include "MsHtmcid.h"

#include "PubDomExtender.h"
#include "DocEventHandler.h"
#include "WindowEventHandler.h"

#define ELEMTYPE_FLASH		0x0001
#define ELEMTYPE_IMAGES		0x0002

extern PUB::CAtlFileMapping<TCHAR> g_filemapBlockedList;

extern CString g_sCoRegKey;
extern CString g_sPubRegKey;

HMENU	g_hPubMenu = NULL;
WNDPROC	g_lpPrevWndProc = NULL;
BOOL	g_bResizePosChanged = FALSE;
BOOL	g_bResizeSysCmd = FALSE;
BOOL	g_bInSizeMove = FALSE;
BOOL	g_bCaptureChanged = FALSE;

#define MAX_PUB_INSTANCES	100

typedef struct tagPUB_INSTANCE
{
	HWND		hWnd;
	CPub*		pPub;
	DWORD		dwThreadId;
	DWORD		dwProcessId;
}
PUB_INSTANCE;

// Shared data among all instances
#pragma data_seg(".SHARED")
// Begin shared data
BOOL	g_bRegistrySettingsLoaded = FALSE;
DWORD	g_dwBlockedTotal = 0;
DWORD	g_dwBlockedSession = 0;
BOOL	g_bEnabled = TRUE;
BOOL	g_bBlockFlash = FALSE;

BOOL	g_bStatusMsg = TRUE;



BOOL	g_bBlockUncachedImages = FALSE;
int		g_nPubInstanceCount = 0;
DWORD	g_dwHideAnimationsHotkey = 0;
PUB_INSTANCE	g_aPubInstances[MAX_PUB_INSTANCES]= {0};
// End shared data
#pragma data_seg()
#pragma comment(linker, "/SECTION:.SHARED,RWS")


// CPub

//
// The functions below are used to manage the array of PUB_INSTANCE
// structures.
//
// The technique for handling PUB_INSTANCE comes from:
// WEB HIGHLIGHTER Version 1.1
// Copyright (c) 2001 Ziff Davis Media, Inc.
// All rights reserved.
// First Published in PC Magazine, US Edition, April 9, 2001.
// Programmer: Patrick Philippot
//
BOOL AddInstance(HWND hWnd, CPub* pPub, DWORD dwThreadId, DWORD dwProcessId)
{
	for (int i = 0; i < g_nPubInstanceCount; i++)
	{
		if ((g_aPubInstances[i].pPub == pPub) &&
			(g_aPubInstances[i].hWnd == hWnd) &&
			(g_aPubInstances[i].dwProcessId == dwProcessId))
		{
			// Already registered
			return TRUE;	
		}
	}

	if (g_nPubInstanceCount < MAX_PUB_INSTANCES)
	{
		g_aPubInstances[g_nPubInstanceCount].hWnd = hWnd;
		g_aPubInstances[g_nPubInstanceCount].pPub = pPub;
		g_aPubInstances[g_nPubInstanceCount].dwThreadId = dwThreadId;
		g_aPubInstances[g_nPubInstanceCount].dwProcessId = dwProcessId;

		g_nPubInstanceCount++;
		return TRUE;
	}
	return FALSE;
}

BOOL RemoveInstance(CPub* pPub, HWND hWnd, DWORD dwProcessId)
{
	BOOL	bOK = FALSE;

	for (int i = 0; i < g_nPubInstanceCount; i++)
	{
		if ((g_aPubInstances[i].pPub == pPub) &&
			(g_aPubInstances[i].hWnd == hWnd) &&
			(g_aPubInstances[i].dwProcessId == dwProcessId))
		{
			if (i < MAX_PUB_INSTANCES - 1)
			{
				memmove(&g_aPubInstances[i], &g_aPubInstances[i+1], 
					sizeof(PUB_INSTANCE) * (g_nPubInstanceCount - i - 1));
			}
			g_nPubInstanceCount--;
			bOK = TRUE;
			break;
		}
	}

	if (g_nPubInstanceCount == 0)	// Last instance
	{
	}

	return bOK;
}

CPub* LookupInstance(HWND hWnd)
{
	int i;

	for (i = 0; i < g_nPubInstanceCount; i++)
	{
		if (g_aPubInstances[i].hWnd == hWnd)
			return g_aPubInstances[i].pPub;
	}

	return NULL;
}

CPub* LookupInstance(DWORD dwThreadId)
{
	int i;
	for (i = 0; i < g_nPubInstanceCount; i++)
	{
		if (g_aPubInstances[i].dwThreadId == dwThreadId)
			return g_aPubInstances[i].pPub;
	}

	return NULL;
}

// 
// Prevent scripts from resizing the browser window.  Basically, this
// works by assuming that if the window is being resized and
//     1) it is not due to a system command (maximize, minimize or restore)
//     2) it is not when IE is first being launched in this process
//     3) it is not because the user double-clicked the title bar 
//     4) the IE has not captured the mouse, i.e., the user isn't 
//        dragging the frame, and
//     5) the user is not using just his eyes to resize the frame ;)
// then it must be due to script and therefore should be blocked.
//
// Credit and kudos go to Ferdinand Oeinck and Gerk Huisma!
//
LRESULT CALLBACK CPub::ResizeWndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
	CPub* pPub;

	// We first try to find a matching window
	do
	{
		pPub = LookupInstance(hWnd);
		if (pPub == NULL)
			hWnd = GetParent(hWnd);
	}
	while ((pPub == NULL) && (hWnd != NULL));

	// If the previous attempt failed, this may be because the IE window that
	// has the input focus is not a child of the main browser window. A typical
	// case is the Favorites | Add Favorites | New Folder command. In this case,
	// we try to lookup using the current thread id
	if (pPub == NULL)
		pPub = LookupInstance(GetCurrentThreadId());

	ATLASSERT(pPub != NULL);
	if (pPub == NULL)	// Should not happen
		return 1;		// No other choice

	if (msg == WM_CAPTURECHANGED)
	{
		g_bCaptureChanged = TRUE;
	}
	else if (msg == WM_ENTERSIZEMOVE)
	{
		g_bInSizeMove = TRUE;
	}
	else if (msg == WM_EXITSIZEMOVE)
	{
		g_bInSizeMove = FALSE;
		g_bCaptureChanged = FALSE;
	}
	else if (msg == WM_SYSCOMMAND)
	{
		// Set resize exception:
		if (wParam == SC_MAXIMIZE || 
			wParam == SC_RESTORE || 
			wParam == SC_MINIMIZE)
		{
			ATLTRACE(_T("(%ld) WM_SYSCOMMAND\n"), ::GetCurrentThreadId());
			g_bResizeSysCmd = TRUE;
		}
	}
	
	else if (msg == WM_WINDOWPOSCHANGED)
	{
		ATLTRACE(_T("(%ld) WM_WINDOWPOSCHANGED\n"), ::GetCurrentThreadId());

		// g_bResizePosChanged is FALSE until after window position
		// change so that IE can restore the window position when initially
		// launched.  Remember that g_bResizePosChanged is global and therefore
		// shared by all BHOs in this process.
		g_bResizePosChanged = TRUE;

		// Reset
		g_bResizeSysCmd = FALSE;

		g_bCaptureChanged = FALSE;
	}
	return CallWindowProcW((WNDPROC)pPub->m_lpResizeOldWndProc, hWnd, msg, wParam, lParam);
}

LRESULT CALLBACK CPub::IEKeyboardProc(int nCode, WPARAM wParam, LPARAM lParam)
{
	// We need a pointer to an instance of CPub to work with. This value
	// has been stored in the g_aPubInstances collection when CPub::SetSite has
	// been called.
	HWND hWnd = GetFocus();
	CPub* pPub;

	// We first try to find a matching window
	do
	{
		pPub = LookupInstance(hWnd);
		if (pPub == NULL)
			hWnd = GetParent(hWnd);
	}
	while ((pPub == NULL) && (hWnd != NULL));

	// If the previous attempt failed, this may be because the IE window that
	// has the input focus is not a child of the main browser window. A typical
	// case is the Favorites | Add Favorites | New Folder command. In this case,
	// we try to lookup using the current thread id
	if (pPub == NULL)
		pPub = LookupInstance(GetCurrentThreadId());

	ATLASSERT(pPub != NULL);
	if (pPub == NULL)	// Should not happen
		return 1;		// No other choice -- will lock keyboard

	if (!pPub->m_bHotkeysEnabled || 
		(nCode < 0) || 
		(lParam & 0xC0000000))	// bit 30 or bit 31 is set
	{
		// Process the key only once.
		// Disable hotkeys if still navigating.
		return CallNextHookEx(pPub->m_hHook, nCode, wParam, lParam);
	}

	UINT nState = pPub->GetKeyState();

	if (g_dwHideAnimationsHotkey != 0 &&
		pPub->KeyMatch(nState, wParam, g_dwHideAnimationsHotkey))
	{
		pPub->BlockWebPageElements(ELEMTYPE_FLASH | ELEMTYPE_IMAGES);
		return 1;
	}

	// Note that we do not let the keystroke go through.
	// Which means that if the keystroke is used by MSIE
	// the corresponding function will not be triggered.
	// Use CallNextHook instead if you want that the original
	// function be triggered as well.
	return CallNextHookEx(pPub->m_hHook, nCode, wParam, lParam);
}

UINT CPub::GetKeyState()
{
	UINT	nState = 0;

	// Make KeyState
	if (::GetKeyState(VK_SHIFT) < 0)
		nState += HOTKEYF_SHIFT;
	if (::GetKeyState(VK_MENU) < 0)
		nState += HOTKEYF_ALT;
	if (::GetKeyState(VK_CONTROL) < 0)
		nState += HOTKEYF_CONTROL;

	return nState;
}

//
// Check whether a key and a key option match
//
BOOL CPub::KeyMatch(UINT nState, WPARAM wParam, int nKey)
{
	return ((wParam == LOBYTE(nKey)) && (nState == HIBYTE(nKey)));
}

//
// Funnel web browser events through this class
//
HRESULT CPub::ManageBrowserConnection(ConnectType eConnectType)
{
	if (eConnectType == ConnType_Unadvise && m_dwBrowserCookie == 0)
		return S_OK;	// not advised, nothing to do

	ATLASSERT(m_spWebBrowser2);
	if (!m_spWebBrowser2)
		return S_OK;

	HRESULT hr = E_FAIL;
	if (eConnectType == ConnType_Advise)
	{
		ATLASSERT(m_dwBrowserCookie == 0);
		hr = AtlAdvise(m_spWebBrowser2, (IDispatch*)this, 
			__uuidof(DWebBrowserEvents2), &m_dwBrowserCookie);
		if (SUCCEEDED(hr) && m_dwBrowserCookie > 0)
		{
			// subclass browser window
			SHANDLE_PTR hWnd;
			hr = m_spWebBrowser2->get_HWND(&hWnd);
			if (SUCCEEDED(hr) && hWnd && !m_lpResizeOldWndProc)
			{
#pragma warning( push )
#pragma warning( disable : 4311 )
#pragma warning( disable : 4312 )
				m_lpResizeOldWndProc = (WNDPROC)::SetWindowLong((HWND)hWnd, 
					GWLP_WNDPROC, (LONG)ResizeWndProc);
#pragma warning( pop ) 
			}

			// Initialize GDI+.
			Gdiplus::GdiplusStartup(&m_gdiplusToken, &m_gdiplusStartupInput, NULL);
		}
	}
	else
	{
		hr = AtlUnadvise(m_spWebBrowser2, 
			__uuidof(DWebBrowserEvents2), m_dwBrowserCookie);
		m_dwBrowserCookie = 0;

		// un-subclass browser window
		SHANDLE_PTR hWnd;
		hr = m_spWebBrowser2->get_HWND(&hWnd);
		if (SUCCEEDED(hr) && hWnd && m_lpResizeOldWndProc)
		{
#pragma warning( push )
#pragma warning( disable : 4311 )
			::SetWindowLong((HWND)hWnd, GWLP_WNDPROC, (LONG)m_lpResizeOldWndProc);
			m_lpResizeOldWndProc = NULL;
#pragma warning( pop ) 
		}

		// Shutdown GDI+.
		Gdiplus::GdiplusShutdown(m_gdiplusToken);
	}
	ATLASSERT(SUCCEEDED(hr));
	return hr;
}

//
// Receive IPropertyNotifySink notifications
//
HRESULT CPub::ManagePropertyNotifyConnection(ConnectType eConnectType)
{
	if (eConnectType == ConnType_Unadvise && m_dwPropNotifyCookie == 0)
		return S_OK;	// not advised, nothing to do

	HRESULT hr = S_OK;

	if (SUCCEEDED(hr))
	{
		if (m_spHTMLDocument)
		{
			if (eConnectType == ConnType_Advise)
			{
				ATLASSERT(m_dwPropNotifyCookie == 0);
				hr = AtlAdvise(m_spHTMLDocument, (IDispatch*)this, IID_IPropertyNotifySink, &m_dwPropNotifyCookie);
			}
			else
			{
				hr = AtlUnadvise(m_spHTMLDocument, IID_IPropertyNotifySink, m_dwPropNotifyCookie);
				m_dwPropNotifyCookie = 0;
			}
		}
		else
		{
			ATLASSERT(FALSE);
			hr = E_FAIL;
		}
	}

	ATLASSERT(SUCCEEDED(hr));
	return hr;
}

//
// Set event handlers for HTMLDocumentEvents2 and HTMLWindowEvents2
//
HRESULT CPub::ManageEventHandlers(CConnInfo* pConnInfo, ConnectType eConnectType)
{
	HRESULT hr = E_FAIL;

	CComQIPtr<IWebBrowser2, &IID_IWebBrowser2> spWB(pConnInfo->m_spWBDisp);
	if (spWB)
	{
		CComPtr<IDispatch> spDocDisp;
		hr = spWB->get_Document(&spDocDisp);
		if (SUCCEEDED(hr) && spDocDisp)
		{
			CComQIPtr<IHTMLDocument2, &IID_IHTMLDocument2> spDoc(spDocDisp);
			if (spDoc)
			{
				IHTMLWindow2* spWnd = NULL;
				hr = spDoc->get_parentWindow(&spWnd);
				if (SUCCEEDED(hr) && spWnd)
				{
					if (eConnectType == ConnType_Advise)
					{
						DWORD dwCookie = 0;

						CComObject<CDocEventHandler>* pDocHandler = NULL;
						hr = CComObject<CDocEventHandler>::CreateInstance (&pDocHandler);
						ATLASSERT(SUCCEEDED(hr));
						if (SUCCEEDED(hr) && pDocHandler)