#include <stdexcept>
#include <cassert>
#include "aes.h"
#include "envelope.h"
#include "rng.h"

#undef DEBUG

#ifdef DEBUG
#include "base64.h"
#endif

using namespace std;

namespace Envelope{

	bool safeKey(const string &key){
		//checks against keys 0 and 1, because they don't undergo change in RSA
		int i;
		for(i=0;i<(int)key.size();i++)if(key[i]!=0)break;
		if(i==(int)key.size())return false; //key is 0
		if(i==(int)key.size()-1&&key[i]==1)return false; //key is 1
		return true; //fine
	}

	string encrypt(const string &data,const RSA::Key &pubkey){
		const int keylen=8; //256-bit
		CryptoRng crng;
		string aeskey(4*keylen,'\0');
		do {
			for(int i=0;i<keylen;i++)*(uint32_t*)&aeskey[4*i]=crng.get();
		} while(!safeKey(aeskey));
#ifdef DEBUG
		cerr<<"decrkey="<<Base64::encode(aeskey)<<endl;
#endif
		
		string payload(AES::encrypt(data,aeskey,AES::AES_256_CBC));
#ifdef DEBUG
		cerr<<"payload="<<Base64::encode(payload)<<endl;
#endif
		Bigint rsadata;
		for(int i=0;i<(int)aeskey.size();i++){
			if(i!=0)rsadata<<=8;
			rsadata+=(uint8_t)aeskey[i];
		}
#ifdef DEBUG
		cerr<<"rsadata="<<rsadata<<endl;
		cerr<<"pubkey = {"<<pubkey.mod<<" , "<<pubkey.exp<<'}'<<endl;
#endif
		Bigint res(RSA::encrypt(rsadata,pubkey));
#ifdef DEBUG
		cerr<<"rsaencr="<<res<<endl;
#endif

		vector<uint8_t> bytes; //bytes in little-endian order (so, reverse!)
		while(res!=0){
			bytes.push_back(res.lowdigits()&0xff);
			res>>=8;
		}
#ifdef DEBUG
		cerr<<"encrkey="<<Base64::encode(string(bytes.rbegin(),bytes.rend()))<<endl;
#endif

		payload.reserve(payload.size()+bytes.size()+2);
		for(int i=bytes.size()-1;i>=0;i--)payload.push_back(bytes[i]); //append in big-endian order again
		payload.push_back(bytes.size()>>8);
		payload.push_back((uint8_t)bytes.size()&0xff);

		return payload;
	}

	string decrypt(const string &data,const RSA::Key &privkey){
#ifdef DEBUG
		cerr<<"=== DECRYPT ==="<<endl;
#endif
		if(data.size()<2)throw invalid_argument("Envelope data length invalid");
		int encrkeylen=((uint16_t)(uint8_t)data[data.size()-2]<<8)+(uint8_t)data.back();
		assert(encrkeylen<(1<<16));
#ifdef DEBUG
		cerr<<"encrkeylen="<<encrkeylen<<endl;
#endif
		if((int)data.size()<encrkeylen+2)throw invalid_argument("Envelope key format invalid");

		string encrkey(encrkeylen,'\0'); //in big-endian
		for(int i=0;i<encrkeylen;i++){
			encrkey[i]=data[data.size()-2-encrkeylen+i];
		}
#ifdef DEBUG
		cerr<<"encrkey="<<Base64::encode(encrkey)<<endl;
#endif

		Bigint rsadata;
		for(int i=0;i<encrkeylen;i++){
			if(i!=0)rsadata<<=8;
			rsadata+=(uint8_t)encrkey[i];
		}
#ifdef DEBUG
		cerr<<"rsaencr="<<rsadata<<endl;
		cerr<<"privkey = {"<<privkey.mod<<" , "<<privkey.exp<<'}'<<endl;
#endif
		Bigint res(RSA::decrypt(rsadata,privkey));
#ifdef DEBUG
		cerr<<"rsadata="<<res<<endl;
#endif

		vector<uint8_t> bytes; //bytes in little-endian order
		for(int i=0;i<32;i++){ //32 is the number of bytes in a 256-bit AES key (256/8=32)
			bytes.push_back(res.lowdigits()&0xff);
			res>>=8;
		}
		if(res!=0){
			throw invalid_argument("Envelope RSA private key incorrect");
		}
		string decrkey(bytes.size(),'\0');
		for(int i=0;i<(int)bytes.size();i++)decrkey[bytes.size()-1-i]=bytes[i];
#ifdef DEBUG
		cerr<<"decrkey="<<Base64::encode(decrkey)<<endl;
#endif

#ifdef DEBUG
		cerr<<"payload="<<Base64::encode(data.substr(0,data.size()-2-encrkeylen))<<endl;
#endif
		try {
			return AES::decrypt(data.substr(0,data.size()-2-encrkeylen),decrkey,AES::AES_256_CBC);
		} catch(invalid_argument){
			throw invalid_argument("Envelope RSA private key incorrect");
		}
	}

}