#include "stdafx.h"
#include "eMule.h"
#include "otherfunctions.h"
//#include "emuleDlg.h"
//#include "KademliaWnd.h"
//#include "KadContactListCtrl.h"
#include "./kademlia/kademlia/Kademlia.h"
#include "./kademlia/routing/Contact.h"
#include "./kademlia/routing/RoutingZone.h"
#include "KadHelper.h"
// BEGIN netfinity: Safe KAD - Ensure we are actually searching for these nodes
#include "./kademlia/utils/UInt128.h"
#include "./kademlia/kademlia/SearchManager.h"
// END netfinity: Safe KAD
#include "./kademlia/net/KademliaUDPListener.h"
#include "./kademlia/net/PacketTracking.h"
#include "opcodes.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#endif

using namespace Kademlia;

//>>> WiZaRd::Kad Security
// the current implementation can and should be enhanced with more checks for 
// packets size to track down "bad" users/clients and shut them off the network
// these checks are a first step...
void _CheckPacketLength(LPCSTR func, const UINT uLenPacket, const UINT desired, const uint8 type) 
{
	CString strError;
	switch(type)
	{
		case _KAD_CHECK_EQUAL:
			if(uLenPacket != desired)
				strError.Format(L"***NOTE: Received wrong size (%u != %u) packet in %hs", uLenPacket, desired, func);
			break; 
		case _KAD_CHECK_BIGGER:
			if(uLenPacket < desired)
				strError.Format(L"***NOTE: Received wrong size (%u < %u) packet in %hs", uLenPacket, desired, func);
			break; 
		case _KAD_CHECK_SMALLER:
			if(uLenPacket > desired)
				strError.Format(L"***NOTE: Received wrong size (%u > %u) packet in %hs", uLenPacket, desired, func);
			break;
		default:
			ASSERT(0);
			break;
	}
	if (!strError.IsEmpty())
		throw strError;
}

void _CheckMessageBounce(LPCSTR func, const UINT uLenPacket, const UINT uIP, const uint16 /*uUDPPort*/, const uint8 byOpcode, const bool bDontRemove)
{	
	if (Kademlia::CKademlia::GetUDPListener()->IsOnOutTrackList(uIP, byOpcode, bDontRemove))
	{ 
		CString strError;
		strError.Format(L"***NOTE: Possible message bounce detected in %hs! Ignoring packet - size (%u)", func, uLenPacket);
		throw strError;
	}
}

void _CheckPacketTracking(LPCSTR func, const UINT uLenPacket, const UINT uIP, const uint16 /*uUDPPort*/, const uint8 byOpcode, const bool bDontRemove)
{
	if (!Kademlia::CKademlia::GetUDPListener()->IsOnOutTrackList(uIP, byOpcode, bDontRemove))
	{ 
		CString strError;
		strError.Format(L"***NOTE: Received unrequested packet, size (%u) in %hs", uLenPacket, func);
		throw strError;
	}
}

//check whether additional bytes are in a packet
void _CheckPacketPos(LPCSTR func, const uint64& pos, const uint64& uLenPacket)
{
	if(pos != uLenPacket)
	{
		CString strError;
		strError.Format(L"***NOTE: There were %I64u additional bytes in %hs", uLenPacket-pos, func);
		throw strError;
	}
}

void _CheckPacketPos(LPCSTR func, const uint64& add)
{
	if(add != 0)
	{
		CString strError;
		strError.Format(L"***NOTE: There were %I64u additional bytes in %hs", add, func);
		throw strError;
	}
}
//<<< WiZaRd::Kad Security
// BEGIN netfinity: Safe KAD - Ensure we are actually searching for these nodes
void _CheckForValidTarget(LPCSTR func, const Kademlia::CUInt128& uTarget)
{
	if (!CSearchManager::AlreadySearchingFor(uTarget))
	{
		CString strError;
		strError.Format(L"***NOTE: Received nodes we didn't search for in %hs", func);
		throw strError;
	}
}
// END netfinity: Safe KAD

// BEGIN netfinity: Safe KAD
void SetGoodRandom(Kademlia::CUInt128& ui)
{
	while (!IsGoodRandom(ui))
		ui.SetValueRandom();
}

//Check if a number appears to be constructed rather than randomly generated
bool IsGoodRandom(const Kademlia::CUInt128& ui)
{
	int	iByteCnt[256];
	memset(iByteCnt, 0, sizeof(iByteCnt));
	const ULONG* m_uData = ui.GetULONGPointer();
	for(uint8 i = 0; i < 4; ++i)
	{
		++iByteCnt[(m_uData[i]>>24)&0xFF];
		++iByteCnt[(m_uData[i]>>16)&0xFF];
		++iByteCnt[(m_uData[i]>>8)&0xFF];
		++iByteCnt[(m_uData[i])&0xFF];
	}
	for(int i = 0; i < 256; ++i)
		if(iByteCnt[i] > 6) 
			return false; // It is highly unlikely that more than 6 out of 16 bytes would roll the same number on a 256 sided dice
	return true;
}
// END netfinity: Safe KAD
