//this file is part of NeoMule
//Copyright (C)2006 David Xanatos ( Xanatos@Lycos.at / http://neomule.sourceforge.net )
//
//This program is free software; you can redistribute it and/or
//modify it under the terms of the GNU General Public License
//as published by the Free Software Foundation; either
//version 2 of the License, or (at your option any later version.
//
//This program is distributed in the hope that it will be useful,
//but WITHOUT ANY WARRANTY; without even the implied warranty of
//MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//GNU General Public License for more details.
//
//You should have received a copy of the GNU General Public License
//along with this program; if not, write to the Free Software
//Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.

#pragma once

#ifdef ISPBYPASS // NEO: SCC - [SecureClinetConnection] -- Xanatos -->
#pragma warning(disable:4516) // access-declarations are deprecated; member using-declarations provide a better alternative
#pragma warning(disable:4244) // conversion from 'type1' to 'type2', possible loss of data
#pragma warning(disable:4100) // unreferenced formal parameter
#include <crypto51/dh.h>
#pragma warning(default:4100) // unreferenced formal parameter
#pragma warning(default:4244) // conversion from 'type1' to 'type2', possible loss of data
#pragma warning(default:4516) // access-declarations are deprecated; member using-declarations provide a better alternative
class CEMSocket;
class CUpDownClient;
class CPortListener;
class CSafeMemFile;


#define	SOCKET_XOR_KEY_SIZE		16
#define	SOCKET_XOR_KEY_LENGTH	37

#define	SOCKET_KEY_SIZE		32
#define	KEY_PACKET_SIZE		48 // PK

class CKey{
public:
	virtual void Init(byte *key_data_ptr, unsigned int key_data_len, bool key_invert, bool bFast = false) = 0;
	virtual void EnCrypt(byte *buffer_ptr, unsigned int buffer_len) = 0;
	virtual void DeCrypt(byte *buffer_ptr, unsigned int buffer_len) = 0;
};


/*****************************************************************************************************************************
* DX2 is a encryption algorithm by David Xanatos, basing on the common RC4 algorithm. 
*  Yes ages ago there was an DX1 (it had nothing to do with RC4) but it don’t fit to protocol obfuscation.
* The DX2 should provide more security by altering the key stream with a summ value of all processed plain data byte.
* The first 1024 bytes of the key stream output are discarded, to protect against a Fluhrer, Mantin and Shamir attack.
* For two way transfer the key data are inverted for one of the sides, so that practically 2 different keys are used.
* Speed Note: On my machine (A64_FX-53) I can process 100 MB and/or generate 10000 initial key streams per one secund.
*************/
class CDX2Key: public CKey
{
public:
	// The key-scheduling algorithm
	void Init(byte *key_data_ptr, unsigned int key_data_len, bool key_invert, bool bFast = false)
	{
		x = 0; 
		y = 0;
		summ = 0;

		// Initialisation of the permutation
		unsigned int i; 
		for(i = 0; i < 256; i++)
			state[i] = (byte)i;

		// Mixing permutation
		unsigned int j = 0;
		unsigned int k;
		for(i = 0; i < 256; i++)
		{ 
			k = i % key_data_len; 
			if(key_invert) // Key invertion for one of the communication partners
				k = key_data_len-1 - k;

			j = (key_data_ptr[k] + state[i] + j) % 256; 
			swap(state[i], state[j]);

			if(i < key_data_len)
				summ = summ + key_data_ptr[k]; // calc the initial summ byte
		} 

		if(!bFast) // Discarding the first 1024 of key stream
			for(i = 0; i < 1024; i++)
				GetKeyChar();
	}

	// Encryption routine
	void EnCrypt(byte *buffer_ptr, unsigned int buffer_len)
	{ 
		unsigned int i;
		for(i = 0; i < buffer_len; i ++)
		{ 
			byte &keyChar = GetKeyChar();

			keyChar ^= summ; // allterring permutation

			byte cur = buffer_ptr[i];
			buffer_ptr[i] ^= keyChar; 

			summ = summ + cur; // mixing bytes
		} 
	}

	// Decryption routine
	void DeCrypt(byte *buffer_ptr, unsigned int buffer_len)
	{ 
		unsigned int i;
		for(i = 0; i < buffer_len; i ++)
		{ 
			byte &keyChar = GetKeyChar();

			keyChar ^= summ; // allterring permutation

			buffer_ptr[i] ^= keyChar; 

			summ = summ + buffer_ptr[i]; // mixing bytes
		} 
	}

private:
	byte state[256]; 
	unsigned int x;
	unsigned int y;
	byte summ;

	// The pseudo-random generation algorithm
	__inline byte& GetKeyChar()
	{
		x = (x + 1) % 256;
		y = (y + state[x]) % 256; 
		swap(state[x], state[y]);
		return state[(state[x] + state[y]) % 256];
	}

	__forceinline void swap(byte &a, byte &b) { byte c = a; a = b; b = c; }
};

// Simply XOR key for verry short encryption (NAT headers 2+4 or 2+8 bytes only)
// and backwards compatybility
// Note: The compatybility to XOR it to be dropped in future releases!!!
class CXORKey: public CKey
{
public:
	void Init(byte *key_data_ptr, unsigned int key_data_len, bool /*key_invert*/, bool /*bFast*/){
		abyKey = key_data_ptr;
		uKeyLen = (uint8)key_data_len;
		step2 = false;
	}

	void EnCrypt(byte *buffer_ptr, unsigned int buffer_len){ 
		XorBuffer(buffer_ptr,buffer_len);
	}

	void DeCrypt(byte *buffer_ptr, unsigned int buffer_len){ 
		XorBuffer(buffer_ptr,buffer_len, step2 ? PACKET_HEADER_SIZE : 0);
		step2 = !step2;
	}

protected:
	friend class CSocketCryptor;

	void XorBuffer(byte* buffer, UINT size, UINT begin = 0){
		UINT posBuf = 0;
		UINT posKey = begin % uKeyLen;
		while(posBuf < size)
		{
			buffer[posBuf] ^= abyKey[posKey];
			posBuf++; 
			posKey++;
			if(posKey == uKeyLen)
				posKey = PACKET_HEADER_SIZE; // use a sepatared key space for headders only
		}
	}

private:
	byte* abyKey;
	uint8 uKeyLen;
	bool  step2;
};

class CSocketCryptor
{
public:
	CSocketCryptor(){ 
		abyKey = NULL;
		uKeyLen = 0;
		eStatus = KS_NONE;
		bPlainPayload = false;
		// DX2
		eSide = KS_B;
		inputKey = NULL;
		outputKey = NULL;
		// DH
		KeyAgreement = NULL;
		abyPrime = NULL;
		abyBase = 0;
	}
	~CSocketCryptor(){ 
		delete [] abyKey; 
		// DX2
		delete inputKey;
		delete outputKey;
		// DH
		delete KeyAgreement;
		delete [] abyPrime;
	}

	//bool IsEnabled()		{ return eStatus != KS_NONE; }
	bool IsSecure()			{ return eStatus == KS_SECURED && (inputKey && outputKey); }
	bool IsXor()			{ return eSide == KS_X; }

	// DH
	void SetDH(const byte* Prime, byte Base){
		abyPrime = new byte[SOCKET_KEY_SIZE];
		memcpy(abyPrime, Prime, SOCKET_KEY_SIZE);
		abyBase = Base;
	}

	// NEO: SUC - [SafeUCPCommunication] -- Xanatos -->
	void ShutDown(){ 
		delete inputKey; inputKey = NULL;
		delete outputKey; outputKey = NULL;
		delete abyPrime; abyPrime = NULL;
	}
	// NEO: SUC END <-- Xanatos --

	CSafeMemFile*	SecureConnection(bool bXorOnly, const char* inKey = NULL, bool bNoTrash = false); // DH

protected:
	friend class CEMSocket;
	friend class CClientUDPSocket;
	friend class CUpDownClient;
	friend class CClientReqSocket;
	friend class CPortListener;
 #ifdef NATTUNNELING // NEO: NATT - [NatTraversal]
	friend class CNATSocket;
 #endif // NATTUNNELING // NEO: NATT END

	void SetError()			{ eStatus = KS_ERROR; }
	bool IsError()			{ return eStatus == KS_ERROR; }
	void SetSending(bool c)	{ eStatus = c ? KS_SENDING_B : KS_SENDING_A; }
	bool IsSending()		{ return eStatus == KS_SENDING_A || eStatus == KS_SENDING_B; }
	void SendingDone()		{ 
								if(eStatus == KS_SENDING_A) 
									eStatus = KS_WAITING; 
								else if(eStatus == KS_SENDING_B) 
									SetSecure(); 
							}
	bool IsWaiting()		{ return eStatus == KS_WAITING; }

	void SetSecure()		{	
								if(eSide == KS_X){ // XOR
									inputKey = new CXORKey;
									outputKey = new CXORKey;
								}else{ // DX2
									inputKey = new CDX2Key;
									outputKey = new CDX2Key;
								}
								inputKey->Init(abyKey, uKeyLen, eSide == KS_A); 
								outputKey->Init(abyKey, uKeyLen, eSide == KS_B);
								eStatus = KS_SECURED; 
 							}
	void SetXor()			{ eSide = KS_X; } // XOR
 	void SetA()				{ eSide = KS_A; } // DX2


	void CrossKeys(byte* pAddon, uint8 uAddonLen) { 	
								if(uKeyLen != SOCKET_KEY_SIZE){ // should never happen
									ASSERT(0);
									return; // but in case
								}
								CSocketCryptor::CrossKeys(abyKey, uKeyLen, pAddon, uAddonLen); 
							}
	static void CrossKeys(byte* pKey, uint8 &uLen, const byte* pAddon, const uint8 uAddonLen);

	// plain payload
	void SetPlainPayload()	{ bPlainPayload = true; } // Note: this can be enabled at any time
	bool IsPlainPayload()	{ return bPlainPayload; }
	uint32 AltSize(uint32 size, uint8 opcode, uint8 prot, bool bUP = false);

	// DX2
	void EnCodeBuffer(char* buffer, UINT size) { ASSERT(IsSecure()); outputKey->EnCrypt((byte*)buffer, size); }
	void DeCodeBuffer(char* buffer, UINT size) { ASSERT(IsSecure()); inputKey->DeCrypt((byte*)buffer, size); }

	// NEO: SUC - [SafeUCPCommunication]
	static void EnCodeBufferUDP(char* buffer, UINT size, const byte* pKey, const uint8 uLen, const BYTE* randomizer, uint8 Side = 0)
	{  
		CKey* outKey;
		if(Side == KS_X) // XOR
			outKey = new CXORKey;
		else // DX2
			outKey = new CDX2Key;

		uint8 tmpLen = uLen;
		byte* tmpKey = new byte[tmpLen];
		memcpy(tmpKey,pKey,tmpLen);
		CSocketCryptor::CrossKeys(tmpKey,tmpLen,randomizer,2);

		outKey->Init(tmpKey, tmpLen, Side == KS_B, true);

		if(Side == KS_X)
		{
			uint8 Len = 0; // here we encrypt only the header the rest is being encrypted by the stream encryptin in the CEMSocket
			switch((uint8)buffer[1]){
				case OP_NAT_RST:
				case OP_NAT_DATA: Len = 4; break;
				case OP_NAT_ACK: Len = 8; break;
				default: ASSERT(0);
			}
			((CXORKey*)outKey)->XorBuffer((byte*)buffer, 2);
			((CXORKey*)outKey)->XorBuffer((byte*)buffer+2, min(Len,size-2),2);
		}
		else
			outKey->EnCrypt((byte*)buffer, size);

		delete [] tmpKey;
		delete outKey;
	}
	static void DeCodeBufferUDP(char* buffer, UINT size, const byte* pKey, const uint8 uLen, const BYTE* randomizer, uint8 Side = 0)
	{ 
		CKey* inKey;
		if(Side == KS_X) // XOR
			inKey = new CXORKey;
		else // DX2
			inKey = new CDX2Key;

		uint8 tmpLen = uLen;
		byte* tmpKey = new byte[tmpLen];
		memcpy(tmpKey,pKey,tmpLen);
		CSocketCryptor::CrossKeys(tmpKey,tmpLen,randomizer,2);

		inKey->Init(tmpKey, tmpLen, Side == KS_A, true); 

		if(Side == KS_X)
		{
			((CXORKey*)inKey)->XorBuffer((byte*)buffer, 2);
			uint8 Len = 0; // here we encrypt only the header the rest is being encrypted by the stream encryptin in the CEMSocket
			switch((uint8)buffer[1]){
				case OP_NAT_RST:
				case OP_NAT_DATA: Len = 4; break;
				case OP_NAT_ACK: Len = 8; break;
				default: ASSERT(0);
			}
			((CXORKey*)inKey)->XorBuffer((byte*)buffer+2, min(Len,size-2),2);
		}
		else
			inKey->DeCrypt((byte*)buffer, size);

		delete [] tmpKey;
		delete inKey;
	}
	// NEO: SUC END

private:
	// Base
	byte* abyKey;
	uint8 uKeyLen;
	enum eKeyStatus {KS_NONE, KS_SENDING_A, KS_SENDING_B, KS_WAITING, KS_SECURED, KS_ERROR} eStatus;
	bool bPlainPayload; // obtional disable encryption for file data

	// DX2
	CKey*	inputKey;
	CKey*	outputKey;
	enum eKeySide {KS_X = 0, KS_A = 1, KS_B = 2} eSide;

	// DH
	CryptoPP::DH* KeyAgreement;
	byte*	abyPrime;
	byte	abyBase;

	static uint8 GetDummy();
};
#endif // ISPBYPASS // NEO: SCC END <-- Xanatos --
