Cookie Spy

/*
 Copyright (c) 2001
 Author: Konstantin Boukreev 
 E-mail: konstantin@mail.primorye.ru 

 Created: 16.02.2001 17:10:00
 Version: 1.0.0

*/
#include "stdafx.h"
#include "IeXtensions.h"
#include "CookieSpy.h"
#include "EditCookieDlg.h"

HRESULT CookieSpy::FinalConstruct()
{
	return CoCreateFreeThreadedMarshaler(GetControllingUnknown(), &m_pUnkMarshaler.p);
}

void CookieSpy::FinalRelease()
{
	m_pUnkMarshaler.Release();
}

/////////////////////////// 
// message handlers

LRESULT CookieSpy::OnCreate(UINT /*uMsg*/, WPARAM /*wParam*/, LPARAM /*lParam*/, BOOL& /*bHandled*/)
{
	CAxWindow wnd(m_hWnd);
	HRESULT hr = wnd.CreateControl(IDR_HTML_MAIN);
	if (SUCCEEDED(hr))
	{
		CComQIPtr<IDispatch> spDisp(GetUnknown());
		hr = wnd.SetExternalDispatch(spDisp);
	}
//	if (SUCCEEDED(hr))
//		hr = wnd.QueryControl(IID_IWebBrowser2, (void**)&m_spBrowser);
	return SUCCEEDED(hr) ? 0 : -1;
}

///////////////////////////
// IObjectWithSite

STDMETHODIMP CookieSpy::SetSite(IUnknown *pUnkSite)
{
	HRESULT hr;
	hr = DeskBandImpl<CookieSpy>::SetSite(pUnkSite);

	if (SUCCEEDED(hr))
	{
		if (m_spUnkSite)
		{				
			CComQIPtr<IServiceProvider> spSP(m_spUnkSite);
			if(spSP)
			{
				spSP->QueryService(__uuidof(SHDocVw::IWebBrowserApp), __uuidof(SHDocVw::IWebBrowser2), (void**)&m_spWebBrowser);
				theDispEvent::DispEventAdvise(m_spWebBrowser);
			}
		}
		else
		{
			if (m_spWebBrowser)
			{
				theDispEvent::DispEventUnadvise(m_spWebBrowser);	
				m_spWebBrowser.Release();
			}
		}
	}
	return hr;
};

///////////////////////////
// ICookieSpy

STDMETHODIMP CookieSpy::put_Default(VARIANT* pVar)
{
	ATLTRACE(_T("CookieSpy::put_Default\n"));	

	if (!pVar) return E_POINTER;
	if (pVar->vt == VT_DISPATCH)
	{
		m_spElem = pVar->pdispVal; 

		CComPtr<MSHTML::IHTMLElement> spElem; 
		m_spElem->get_parentElement(&spElem);
		if (spElem)
		{
			CComQIPtr<MSHTML::IHTMLBodyElement> spBody(spElem);
			if (spBody)
			{
				TCHAR path[1024] = {0};
				GetModuleFileName(_Module.m_hInstResource, path, 1024);  			
				CComBSTR url(L"res://");
				url += path;
				url += L"/RT_JPEG/#201";		

				HRESULT hr;
				hr = spBody->put_background(url); 
			}
		}
	}
	return S_OK;
}

STDMETHODIMP CookieSpy::get_Default(VARIANT* pVar)
{
	ATLTRACE(_T("CookieSpy::get_Default\n"));

	if (!pVar) return E_POINTER;
	VariantClear(pVar);
	pVar->vt = VT_DISPATCH;
	return m_spElem->QueryInterface(&pVar->pdispVal);  	
}

STDMETHODIMP CookieSpy::Event(VARIANT* vHint, VARIANT* vEvent)
{
	USES_CONVERSION;

	vHint;	// didn't use so far

	if (!vEvent || 
		vEvent->vt != VT_DISPATCH ||
		!vEvent->pdispVal)
		return E_INVALIDARG;
	
	CComQIPtr<MSHTML::IHTMLEventObj> spEvent(vEvent->pdispVal);
	if (!spEvent) return E_INVALIDARG;
	
	CComPtr<MSHTML::IHTMLElement> spElem;
	spEvent->get_srcElement(&spElem);
	if (!spElem) return S_FALSE;
	
	#ifdef _DEBUG
	CComBSTR sType;
	spEvent->get_type(&sType);
	ATLASSERT(wcscmp(sType, L"click") == 0);
	#endif //_DEBUG

	CComBSTR sClass;
	spElem->get_className(&sClass);

	if (!sClass.Length() ||
		wcscmp(sClass, L"label"))
		return S_OK;
	
	CComBSTR  sID;
	spElem->get_id(&sID); 
	ATLASSERT(sID.Length());

	return OnClickLabel(sID);	
}

///////////////////////////
// DWebBrowserEvents2

HRESULT __stdcall CookieSpy::BeforeNavigate2(IDispatch *pDisp, VARIANT *URL, VARIANT *Flags, VARIANT *TargetFrameName, VARIANT *PostData, VARIANT *Headers, VARIANT_BOOL *Cancel)
{
	Flags;
	TargetFrameName;
	PostData;
	Headers;
	Cancel;

	ATLTRACE(_T("CookieSpy::BeforeNavigate2\n"));

	UpdateCookies(pDisp, L"BeforeNavigate");
	return S_OK;
}

HRESULT __stdcall CookieSpy::DocumentComplete(IDispatch* pDisp, VARIANT* URL)
{	
	ATLTRACE(_T("CookieSpy::DocumentComplete\n"));	

	UpdateCookies(pDisp, L"DocumentComplete");
	return S_OK;
}

HRESULT __stdcall CookieSpy::NavigateComplete2(IDispatch* pDisp, VARIANT* URL)
{
	ATLTRACE(_T("CookieSpy::NavigateComplete2\n"));

	UpdateCookies(pDisp, L"NavigateComplete");
	return S_OK;
}

////////////////////////////////////
//

void CookieSpy::UpdateCookies(IDispatch* pDisp, const WCHAR* sEvent)
{
	CComQIPtr<SHDocVw::IWebBrowser2> spWB(pDisp);
	if (!spWB) return;
	
	CComPtr<IDispatch> spDocDisp;
	spWB->get_Document(&spDocDisp);
	m_spDoc = spDocDisp;

	if (m_spDoc)
	{			
		BSTR sCookie = 0;		
		m_spDoc->get_cookie(&sCookie);		
		RefreshCookieTable(sCookie, sEvent);
	}
}

void CookieSpy::RefreshCookieTable(BSTR sCookie, const WCHAR* sEvent)
{
	if (!m_spElem) return;
	if (!m_spWebBrowser) return;

	CComBSTR url, lname;

	m_spWebBrowser->get_LocationURL(&url);
	m_spWebBrowser->get_LocationName(&lname);
	
	bool bNewLocation;

	if (m_url == url.m_str)
	{
		bNewLocation = false;
	}
	else
	{
		bNewLocation = true;
		m_url = url;
		m_CookieEntries.Clear();
	}
	
	typedef std::basic_ostringstream<wchar_t> ostringstream;
	typedef std::basic_string<wchar_t> string;

	ostringstream os;

	// table's header
	os << L"<div class=\"status\" style=\"text-align:center\"><a href=\"" << url.m_str << L"\">" << lname.m_str << L"</a>";
	os << L"<span> - " << sEvent << L"<span></div>";
	os << L"<hr><table border=\"0\" cellpadding=\"1\" cellspacing=\"2\" width=\"100%%\" class=\"text\">";
	os << L"<tr><th><span class=\"label\" id=\"e84289f8-316e-4105-984f-28031bb5e26f\">+</span></th><th>name</th><th>value</th></tr>";
	
	// table's data
	if (SysStringLen(sCookie))
	{
		CookieEntry* pEntry = 0;
		wchar_t tokens[]	= L";=";
		wchar_t* token		= wcstok(sCookie, tokens);
		
		while (token)
		{			
			if (!pEntry)
			{					
				while(iswspace(*token)) token++;

				pEntry = m_CookieEntries.Find(token);
				if (pEntry) 
				{
					// update date
					pEntry->m_sName  = token;
					pEntry->m_nParam = 1;					
				}
				else
				{
					pEntry = m_CookieEntries.Add(token);
				}

				ATLASSERT(wcslen(token));
												
				os << L"<tr><td><span class=\"label\" id=\"" << token << L"\">x</span></td><td";
				if (!bNewLocation && !pEntry->m_nParam)
					os << L" style=\"color:red\"";
				os << L">" << token << L"</td>";
				
			}
			else
			{	
				ATLASSERT(wcslen(token));

				os << L"<td";
				if (!bNewLocation && (!pEntry->m_sValue || wcscmp(pEntry->m_sValue, token)))
					os << L" style=\"color:red\"";				
				os <<L">" << token << L"</td></tr>";

				pEntry->m_sValue = token;
				pEntry = 0;
			}
						
			token = wcstok(0, tokens);
		}
	}
		
	// table's end
	os << L"</table><hr>";
		
	m_cookies.Empty(); 
	m_cookies.Attach(sCookie);	
	m_CookieEntries.Reorganize();

	string s = os.str();

	BSTR bstr = SysAllocStringLen(s.c_str(), s.size());
	HRESULT hr = m_spElem->put_innerHTML(bstr);
	SysFreeString(bstr);
}

bool CookieSpy::GetCookieDataFromCookieFolder(WCHAR* sCookie, WCHAR* sValue, BSTR* psUrl, SYSTEMTIME& ExpireDate)
{	
	USES_CONVERSION;

	CComPtr<MSHTML::IHTMLLocation> spLoc;
	m_spDoc->get_location(&spLoc);
	if (!spLoc) return false;
						
	CComBSTR sHostName;
	spLoc->get_hostname(&sHostName); 
	if (!sHostName.Length()) return false;
	
	WCHAR* sHost = 0;

	sHost = wcschr(sHostName, L'.');
	if (!sHost) return false;	
	
	WCHAR* p = wcsrchr(sHost, L'.');
	if (!p) return false;	
		
	*p = 0;	
	sHost++;
	if (!wcslen(sHost)) return false;

	TCHAR sUser[UNLEN + 1];
	DWORD dwSize = UNLEN + 1;
	if (!GetUserName(sUser, &dwSize))
		return false;

	TCHAR szPath[MAX_PATH] = {0};

	/*
	PFNSHGETFOLDERPATH pfnSHGetFolderPath = 0;
	HMODULE hMod = LoadLibrary(_T("shfolder.dll"));
	#ifdef _UNICODE
	pfnSHGetFolderPath = (PFNSHGETFOLDERPATHW)GetProcAddress(hMod, "SHGetFolderPathW");	
	#else
	pfnSHGetFolderPath = (PFNSHGETFOLDERPATHA)GetProcAddress(hMod, "SHGetFolderPathA");	
	#endif
	
	HRESULT hr = pfnSHGetFolderPath(0, CSIDL_COOKIES, 0, 0, szPath);
	FreeLibrary(hMod);
	if (FAILED(hr))	return false;
	*/

	if (!SHGetSpecialFolderPath(0, szPath, CSIDL_COOKIES, FALSE))
		return false;
	
	PathAddBackslash(szPath);
	
	TCHAR sSearch[MAX_PATH];
	wsprintf(sSearch, _T("%s%s@%s????.txt"), szPath, sUser, W2T(sHost));

	HANDLE f;
	WIN32_FIND_DATA fd;

	f = FindFirstFile(sSearch, &fd);	
	if (f == INVALID_HANDLE_VALUE)
		return false;

	bool r = false;

	do 
	{		
		TCHAR cookieFile[MAX_PATH];
		lstrcpy(cookieFile, szPath);		
		lstrcat(cookieFile, fd.cFileName);

		r = GetCookieDataFromCookieFile(cookieFile, sCookie, sValue, psUrl, ExpireDate);
		if (r) break;
		
	} while(FindNextFile(f, &fd));

	FindClose(f);

	return r;
}

bool CookieSpy::GetCookieDataFromCookieFile(TCHAR* sFile, WCHAR* sCookie, WCHAR* sValue, BSTR* psUrl, SYSTEMTIME& ExpireDate)
{
	USES_CONVERSION;

	HANDLE hFile = CreateFile(sFile, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
	if (INVALID_HANDLE_VALUE == hFile)return false;

	DWORD cbSize = GetFileSize(hFile, 0);

	if (-1 == cbSize)
	{
		CloseHandle(hFile);
		return false;
	}
	
	char* pData = (char*)_alloca(cbSize);
	ReadFile(hFile, pData, cbSize, &cbSize, 0);
	CloseHandle(hFile);
	
	char* tokenUrl = 0;
	char* tokenLowDate = 0;
	char* tokenHighDate = 0;	
	bool r = false;
	int i = 0;

	char* CookieName  = W2A(sCookie);
	char* CookieValue = W2A(sValue);

	char tokens[] = "\n";	
	char* token = strtok(pData, tokens);

	while (token)
	{	
		switch(i)
		{
			case 0: //cookie name
				r = strcmp(token, CookieName) == 0;
				break;
			case 1: // cookie value
				if (r) r = strcmp(token, CookieValue) == 0;
				break;
			case 2: // cookie path
				if (r) tokenUrl = token;
				break;
			case 3: // secure ?
				break;
			case 4: // expire date
				if (r) tokenLowDate = token;
				break;
			case 5: // expire date
				if (r) tokenHighDate = token;
				break;
			case 6: // creation date ?
				break;
			case 7: // creation date ?	
				break;
			case 8: // * <- end of cookie record
				if (*token != '*') return false;				
				break;			
		}

		if (i == 8) 
		{
			if (r) break;
			i = 0;
		}
		else
		{
			i++;
		}

		token = strtok(0, tokens);
	}

	if (r)
	{
		FILETIME ft;
		ft.dwLowDateTime  = atol(tokenLowDate);
		ft.dwHighDateTime = atol(tokenHighDate);		
		FileTimeToSystemTime(&ft, &ExpireDate);

		char ValidUrl[MAX_PATH] = "http://";
		strcat(ValidUrl, tokenUrl);

		*psUrl = SysAllocString(A2W(ValidUrl));
	}

	return r;
}

HRESULT CookieSpy::OnClickLabel(WCHAR* sID)
{			
	CookieEntry* pEntry = 0;
	TCHAR* sTitle		= _T("");
	TCHAR* sHelp		= _T("");
	WCHAR* sValue		= L"";	
	bool r				= false;
	bool bSession		= false;

	CComBSTR sUrl;

	SYSTEMTIME st;
	GetSystemTime(&st);

	if (wcscmp(sID, L"e84289f8-316e-4105-984f-28031bb5e26f"))
	{
		pEntry = m_CookieEntries.Find(sID);

		if (pEntry)
		{
			sTitle = _T("Edit Cookie");

			sValue = pEntry->m_sValue;						
			r = GetCookieDataFromCookieFolder(sID, pEntry->m_sValue, &sUrl, st);
			if (!r)
			{
				bSession = true;
			//	sHelp = _T("");
			}			
		}
		else
		{
			MessageBox(_T("Error, Cannot find the cookie"), _T("Cookie Spy"), MB_ICONWARNING);
			return E_FAIL;
		}
	}
	else
	{
		sTitle = _T("Add new Cookie");
		sID = L"";
		st.wYear++;
	}
	
	if (!sUrl.Length())
	{		
		sUrl = m_url;
	}

	EditCookieDlg dlg(sUrl, sID, sValue, st, bSession, sTitle, sHelp);
	if (IDOK == dlg.DoModal())
	{
		CComBSTR& sCookie = dlg.m_sCookie;

		sCookie += L"=";
		sCookie += dlg.m_sValue;
		sCookie += L";";

		if (!dlg.m_bSession)
		{
			SYSTEMTIME& st = dlg.m_st; 

			static wchar_t* Days[]  = {L"Sun", L"Mon", L"Tue", L"Wed", L"Thu", L"Fri", L"Sat"};
			static wchar_t* Months[] = {L"Jan", L"Feb", L"Mar", L"Apr", L"May", L"Jun", L"Jul", L"Aug", L"Sen", L"Oct", L"Nov", L"Dec"};

			wchar_t expires[100];
			_snwprintf(expires, 100,
				L" expires=%s, %u %s %u %02u:%02u:%02u GMT", 
				Days[st.wDayOfWeek], 
				st.wDay, 
				Months[st.wMonth - 1],
				st.wYear,
				st.wHour,
				st.wMinute,
				st.wSecond);

			sCookie += expires;
		}

		BOOL r = InternetSetCookieW(dlg.m_sUrl, 0, sCookie);

		if (!r)
		{
			MessageBox(_T("Error, Cannot set the cookie"), _T("Cookie Spy"), MB_ICONWARNING);
		}
	}

	return S_OK;
}

////////////////////////////////////////////
// CookieSpy::CookieEntries

CookieSpy::CookieEntry* CookieSpy::CookieEntries::Add(WCHAR* sName)
{
	push_back(CookieEntry(sName));
	return &back();
}

CookieSpy::CookieEntry* CookieSpy::CookieEntries::Find(WCHAR* sName)
{
	ATLASSERT(sName && wcslen(sName));

	LONG l = wcslen(sName);
	
	for (iterator x = begin(); x != end(); x++)
	{
		ATLASSERT((*x).m_sName && wcslen((*x).m_sName));
	
		if ((*x).m_nHint == l &&
			wcscmp((*x).m_sName, sName) == 0)
		{
			return &(*x);
		}
	}
		
	return 0;
}

void CookieSpy::CookieEntries::Reorganize()
{
	iterator x = begin();

	while (x != end())
	{
		if ((*x).m_nParam == 0xffffffff)
		{
			iterator _x = x++;
			erase(_x);
			continue;
		}
		
		(*x).m_nParam = 0xffffffff;
		 x++;
	}
}

void CookieSpy::CookieEntries::Clear()
{
	clear();
}