#include "stdafx.h"

#include "cipher.h"

#ifdef _IMPROVED_PACKET_ENCRYPTION_

//#pragma warning(push)
//#pragma warning(disable: 4100 4127 4189 4231 4512 4706)

#include <cryptopp/modes.h>
#include <cryptopp/nbtheory.h>
#include <cryptopp/osrng.h>

// Diffie-Hellman key agreement
#include <cryptopp/dh.h>
#include <cryptopp/dh2.h>

// AES winner and candidates
#include <cryptopp/aes.h>
#include <cryptopp/cast.h>
#include <cryptopp/rc6.h>
#include <cryptopp/mars.h>
#include <cryptopp/serpent.h>
#include <cryptopp/twofish.h>
// Other block ciphers
#include <cryptopp/blowfish.h>
#include <cryptopp/camellia.h>
#include <cryptopp/des.h>
#include <cryptopp/idea.h>
#include <cryptopp/rc5.h>
#include <cryptopp/seed.h>
#include <cryptopp/shacal2.h>
#include <cryptopp/skipjack.h>
#include <cryptopp/tea.h>

#ifdef __THEMIDA__
#include <ThemidaSDK.h>
#endif

#include "Debug.h"


using namespace CryptoPP;

// Block cipher algorithm selector abstract base class.
struct BlockCipherAlgorithm {
	enum {
		kDefault, // to give more chances to default algorithm
		// AES winner and candidates
//		kAES, // Rijndael
		kRC6,
		kMARS,
		kTwofish,
		kSerpent,
		kCAST256,
		// Other block ciphers
		kIDEA,
		k3DES, // DES-EDE2
		kCamellia,
		kSEED,
		kRC5,
		kBlowfish,
		kTEA,
		//kSKIPJACK,
		kSHACAL2,
		// End sentinel
		kMaxAlgorithms
	};

	BlockCipherAlgorithm() {}
	virtual ~BlockCipherAlgorithm() {}

	static BlockCipherAlgorithm* Pick(int hint);

	virtual int GetBlockSize() const = 0;
	virtual int GetDefaultKeyLength() const = 0;
	virtual int GetIVLength() const = 0;

	virtual SymmetricCipher* CreateEncoder(const CryptoPP::byte* key, size_t keylen,
		const CryptoPP::byte* iv) const = 0;
	virtual SymmetricCipher* CreateDecoder(const CryptoPP::byte* key, size_t keylen,
		const CryptoPP::byte* iv) const = 0;
};

// Block cipher (with CTR mode) algorithm selector template class.
template<class T>
struct BlockCipherDetail : public BlockCipherAlgorithm {
	BlockCipherDetail() {}
	virtual ~BlockCipherDetail() {}

	virtual int GetBlockSize() const { return T::BLOCKSIZE; }
	virtual int GetDefaultKeyLength() const { return T::DEFAULT_KEYLENGTH; }
	virtual int GetIVLength() const { return T::IV_LENGTH; }

	virtual SymmetricCipher* CreateEncoder(const CryptoPP::byte* key, size_t keylen,
		const CryptoPP::byte* iv) const 
	{
		return new typename CTR_Mode<T>::Encryption(key, keylen, iv);
	}
	virtual SymmetricCipher* CreateDecoder(const CryptoPP::byte* key, size_t keylen,
		const CryptoPP::byte* iv) const 
	{
		return new typename CTR_Mode<T>::Decryption(key, keylen, iv);
	}
};

// Key agreement scheme abstract class.
class KeyAgreement {
public:
	KeyAgreement() {}
	virtual ~KeyAgreement() {}

	virtual size_t Prepare(void* buffer, size_t* length) = 0;
	virtual bool Agree(size_t agreed_length, const void* buffer, size_t length) = 0;

	const SecByteBlock& shared() const { return shared_; }

protected:
	SecByteBlock shared_;
};

// Crypto++ Unified Diffie-Hellman key agreement scheme implementation.
class DH2KeyAgreement : public KeyAgreement {
public:
	DH2KeyAgreement();
	virtual ~DH2KeyAgreement();

	virtual size_t Prepare(void* buffer, size_t* length);
	virtual bool Agree(size_t agreed_length, const void* buffer, size_t length);

private:
	DH dh_;
	DH2 dh2_;
	SecByteBlock spriv_key_;
	SecByteBlock epriv_key_;
};

Cipher::Cipher()
: activated_(false), encoder_(NULL), decoder_(NULL), key_agreement_(NULL) {
}

Cipher::~Cipher() {
	if (activated_) {
		CleanUp();
	}
}

void Cipher::CleanUp() {
	if (encoder_ != NULL) {
		delete encoder_;
		encoder_ = NULL;
	}
	if (decoder_ != NULL) {
		delete decoder_;
		decoder_ = NULL;
	}
	if (key_agreement_ != NULL) {
		delete key_agreement_;
		key_agreement_ = NULL;
	}
	activated_ = false;
}

size_t Cipher::Prepare(void* buffer, size_t* length) {
#ifdef __THEMIDA__
	VM_START
#endif

	assert(key_agreement_ == NULL);
	key_agreement_ = new DH2KeyAgreement();
	assert(key_agreement_ != NULL);
	size_t agreed_length = key_agreement_->Prepare(buffer, length);
	if (agreed_length == 0) {
		delete key_agreement_;
		key_agreement_ = NULL;
	}
#ifdef __THEMIDA__
	VM_END
#endif

	return agreed_length;
}

bool Cipher::Activate(bool polarity, size_t agreed_length,
					  const void* buffer, size_t length) {
#ifdef __THEMIDA__
	VM_START
#endif

	assert(activated_ == false);
	assert(key_agreement_ != NULL);
	bool result = false;
	if (key_agreement_->Agree(agreed_length, buffer, length)) {
		result = SetUp(polarity);
	}
	delete key_agreement_;
	key_agreement_ = NULL;
#ifdef __THEMIDA__
	VM_END
#endif

	return result;
}

bool Cipher::SetUp(bool polarity) {
#ifdef __THEMIDA__
	VM_START
#endif

	assert(key_agreement_ != NULL);
	const SecByteBlock& shared = key_agreement_->shared();

	// Pick a block cipher algorithm

	if (shared.size() < 2) {
		return false;
	}

	int hint_0 = shared.BytePtr()[*(shared.BytePtr()) % shared.size()];
	int hint_1 = shared.BytePtr()[*(shared.BytePtr() + 1) % shared.size()];
	BlockCipherAlgorithm* detail_0 = BlockCipherAlgorithm::Pick(hint_0);
	BlockCipherAlgorithm* detail_1 = BlockCipherAlgorithm::Pick(hint_1);
	assert(detail_0 != NULL);
	assert(detail_1 != NULL);
	std::unique_ptr<BlockCipherAlgorithm> algorithm_0(detail_0);
	std::unique_ptr<BlockCipherAlgorithm> algorithm_1(detail_1);

	const size_t key_length_0 = algorithm_0->GetDefaultKeyLength();
	const size_t iv_length_0 = algorithm_0->GetBlockSize();

	if (shared.size() < key_length_0 || shared.size() < iv_length_0) {
		return false;
	}

	const size_t key_length_1 = algorithm_1->GetDefaultKeyLength();
	const size_t iv_length_1 = algorithm_1->GetBlockSize();

	if (shared.size() < key_length_1 || shared.size() < iv_length_1) {
		return false;
	}

	// Pick encryption keys and initial vectors

	SecByteBlock key_0(key_length_0), iv_0(iv_length_0);
	SecByteBlock key_1(key_length_1), iv_1(iv_length_1);

	size_t offset;

	key_0.Assign(shared, key_length_0);
	offset = key_length_0;
	offset = min(key_length_0, shared.size() - key_length_1);
	key_1.Assign(shared.BytePtr() + offset, key_length_1);

	offset = shared.size() - iv_length_0;
	iv_0.Assign(shared.BytePtr() + offset, iv_length_0);
	offset = (offset < iv_length_1 ? 0 : offset - iv_length_1);
	iv_1.Assign(shared.BytePtr() + offset, iv_length_1);

	// Create encryption/decryption objects

	if (polarity) {
		encoder_ = algorithm_1->CreateEncoder(key_1, key_1.size(), iv_1);
		decoder_ = algorithm_0->CreateDecoder(key_0, key_0.size(), iv_0);
	} else {
		encoder_ = algorithm_0->CreateEncoder(key_0, key_0.size(), iv_0);
		decoder_ = algorithm_1->CreateDecoder(key_1, key_1.size(), iv_1);
	}

	assert(encoder_ != NULL);
	assert(decoder_ != NULL);
#ifdef __THEMIDA__
	VM_END
#endif

	return true;
}

BlockCipherAlgorithm* BlockCipherAlgorithm::Pick(int hint) {
	BlockCipherAlgorithm* detail;
	int selector = hint % kMaxAlgorithms;
	switch (selector) {
	//case kAES:
	//	detail = new BlockCipherDetail<AES>();
		break;
	case kRC6:
		detail = new BlockCipherDetail<RC6>();
		break;
	case kMARS:
		detail = new BlockCipherDetail<MARS>();
		break;
	case kTwofish:
		detail = new BlockCipherDetail<Twofish>();
		break;
	case kSerpent:
		detail = new BlockCipherDetail<Serpent>();
		break;
	case kCAST256:
		detail = new BlockCipherDetail<CAST256>();
		break;
	case kIDEA:
		detail = new BlockCipherDetail<IDEA>();
		break;
	case k3DES:
		detail = new BlockCipherDetail<DES_EDE2>();
		break;
	case kCamellia:
		detail = new BlockCipherDetail<Camellia>();
		break;
	case kSEED:
		detail = new BlockCipherDetail<SEED>();
		break;
	case kRC5:
		detail = new BlockCipherDetail<RC5>();
		break;
	case kBlowfish:
		detail = new BlockCipherDetail<Blowfish>();
		break;
	case kTEA:
		detail = new BlockCipherDetail<TEA>();
		break;
//	case kSKIPJACK:
//		detail = new BlockCipherDetail<SKIPJACK>();
//		break;
	case kSHACAL2:
		detail = new BlockCipherDetail<SHACAL2>();
		break;
	case kDefault:
	default:
		detail = new BlockCipherDetail<Twofish>(); // default algorithm
		break;
	}

	return detail;
}

DH2KeyAgreement::DH2KeyAgreement() : dh_(), dh2_(dh_) {
}

DH2KeyAgreement::~DH2KeyAgreement() {
}

size_t DH2KeyAgreement::Prepare(void* buffer, size_t* length) {
#ifdef __THEMIDA__
	VM_START
#endif

	// RFC 5114, 1024-bit MODP Group with 160-bit Prime Order Subgroup
	// http://tools.ietf.org/html/rfc5114#section-2.1
	Integer p("0xB10B8F96A080E01DDE92DE5EAE5D54EC52C99FBCFB06A3C6"
		"9A6A9DCA52D23B616073E28675A23D189838EF1E2EE652C0"
		"13ECB4AEA906112324975C3CD49B83BFACCBDD7D90C4BD70"
		"98488E9C219A73724EFFD6FAE5644738FAA31A4FF55BCCC0"
		"A151AF5F0DC8B4BD45BF37DF365C1A65E68CFDA76D4DA708"
		"DF1FB2BC2E4A4371");

	Integer g("0xA4D1CBD5C3FD34126765A442EFB99905F8104DD258AC507F"
		"D6406CFF14266D31266FEA1E5C41564B777E690F5504F213"
		"160217B4B01B886A5E91547F9E2749F4D7FBD7D3B9A92EE1"
		"909D0D2263F80A76A6A24C087A091F531DBF0A0169B6A28A"
		"D662A4D18E73AFA32D779D5918D08BC8858F4DCEF97C2A24"
		"855E6EEB22B3B2E5");

	Integer q("0xF518AA8781A8DF278ABA4E7D64B7CB9D49462353");

	// Schnorr Group primes are of the form p = rq + 1, p and q prime. They
	// provide a subgroup order. In the case of 1024-bit MODP Group, the
	// security level is 80 bits (based on the 160-bit prime order subgroup).		

	// For a compare/contrast of using the maximum security level, see
	// dh-unified.zip. Also see http://www.cryptopp.com/wiki/Diffie-Hellman
	// and http://www.cryptopp.com/wiki/Security_level .

	AutoSeededRandomPool rnd;

	dh_.AccessGroupParameters().Initialize(p, q, g);

	if(!dh_.GetGroupParameters().ValidateGroup(rnd, 3)) {
		// Failed to validate prime and generator
		return 0;
	}

	size_t count = 0;

	p = dh_.GetGroupParameters().GetModulus();
	q = dh_.GetGroupParameters().GetSubgroupOrder();
	g = dh_.GetGroupParameters().GetGenerator();

	// http://groups.google.com/group/sci.crypt/browse_thread/thread/7dc7eeb04a09f0ce
	Integer v = ModularExponentiation(g, q, p);

	if(v != Integer::One()) {
		// Failed to verify order of the subgroup
		return 0;
	}

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

	spriv_key_.New(dh2_.StaticPrivateKeyLength());
	epriv_key_.New(dh2_.EphemeralPrivateKeyLength());
	SecByteBlock spub_key(dh2_.StaticPublicKeyLength());
	SecByteBlock epub_key(dh2_.EphemeralPublicKeyLength());

	dh2_.GenerateStaticKeyPair(rnd, spriv_key_, spub_key);
	dh2_.GenerateEphemeralKeyPair(rnd, epriv_key_, epub_key);

	// Prepare key agreement data
	const size_t spub_key_length = spub_key.size();
	const size_t epub_key_length = epub_key.size();
	const size_t data_length = spub_key_length + epub_key_length;

	if (*length < data_length) {
		// Not enough data buffer length b-l-a-c-k
		return 0;
	}

	*length = data_length;
	CryptoPP::byte* buf = (CryptoPP::byte*)buffer;
	memcpy(buf, spub_key.BytePtr(), spub_key_length);
	memcpy(buf + spub_key_length, epub_key.BytePtr(), epub_key_length);

#ifdef __THEMIDA__
	VM_END
#endif

	return dh2_.AgreedValueLength();
}

bool DH2KeyAgreement::Agree(size_t agreed_length, const void* buffer, size_t length) {
	if (agreed_length != dh2_.AgreedValueLength()) {
		// Shared secret size mismatch
		return false;
	}
	const size_t spub_key_length = dh2_.StaticPublicKeyLength();
	const size_t epub_key_length = dh2_.EphemeralPublicKeyLength();
	if (length != (spub_key_length + epub_key_length)) {
		// Wrong data length
		return false;
	}
	shared_.New(dh2_.AgreedValueLength());
	const CryptoPP::byte* buf = (const CryptoPP::byte*)buffer;
	if (!dh2_.Agree(shared_, spriv_key_, epriv_key_, buf, buf + spub_key_length)) {
		// Failed to reach shared secret
		return false;
	}
	return true;
}

#endif
