#include <iostream>
#include <stdexcept>
#include <cstdlib>
#include <cstring>
#include "../base64.h"
#include "../envelope.h"

using namespace std;

const char *progname="envelope";

void usage(){
	cerr<<"Usage: "                                                                         <<endl
	    <<"   "<<progname<<" -k <keylength>   Generate a key with the given number of bits" <<endl
	    <<"   "<<progname<<" -e <pubkey>      Encrypt text on standard input"               <<endl
	    <<"   "<<progname<<" -d <privkey>     Decrypt text on standard input"               <<endl
	                                                                                        <<endl
	    <<"The program first encrypts the plaintext using 256-bit AES CBC with a randomly"  <<endl
	    <<"generated key, and then encrypts the key using RSA, with the specified key."     <<endl
	    <<"All RSA keys are encoded in base64 in an internal format."                       <<endl
	                                                                                        <<endl
	    <<"The <keylength> used for RSA key generation is the number of bits in the"        <<endl
	    <<"modulus. Advised for security is at least 2048, but that might take a while to"  <<endl
	    <<"generate (be patient...). (Note: duration of key generation has *absolutely no*" <<endl
	    <<"meaning as to the strength of the key.) There is no restriction to powers of two"<<endl
	    <<"in the key length (or any other special class of numbers), but stay above 512,"  <<endl
	    <<"or the AES key might not fit."                                                   <<endl
	                                                                                        <<endl
	    <<"  -k  The public and private keys are printed on standard output, separated by"  <<endl
	    <<"      newlines. Comments are on standard error."                                 <<endl
	    <<"  -e  Text to encrypt is expected on standard input; the public key of the"      <<endl
	    <<"      recipient is a command-line argument. Encrypted text is in base64 on"      <<endl
	    <<"      standard output."                                                          <<endl
	    <<"  -d  Text to decrypt is expected on standard input, in the base64 form that is" <<endl
	    <<"      outputted by the -e function. your private key is a command-line argument."<<endl
	    <<"      Decrypted text is on standard output."                                     <<endl;
}

void mode_keygen(int keylength){
	if(keylength<512){
		cerr<<"Key length too small, please specify length >= 512"<<endl;
		exit(1);
	}
	if(keylength>8192){
		cerr<<"Key length too large, please specify length <= 8192"<<endl;
		exit(1);
	}
	pair<RSA::Key,RSA::Key> keyset=RSA::genkeys(keylength);
	cerr<<"Public key:"<<endl;
	cout<<RSA::exportKey(keyset.first)<<endl;
	cerr<<"Private key:"<<endl;
	cout<<RSA::exportKey(keyset.second)<<endl;
}

void mode_encrypt(const string &pubkeyrepr){
	RSA::Key key;
	try {
		key=RSA::importKey(pubkeyrepr);
	} catch(invalid_argument){
		cerr<<"The given public key is not a valid key!"<<endl;
		exit(1);
	}
	string data;
	char buf[1024];
	while(cin){
		cin.read(buf,sizeof(buf));
		int nread=cin.gcount();
		if(nread==0)continue;
		data.append(buf,nread);
	}
	cout<<Base64::encode(Envelope::encrypt(data,key))<<endl;
}

void mode_decrypt(const string &privkeyrepr){
	RSA::Key key;
	try {
		key=RSA::importKey(privkeyrepr);
	} catch(invalid_argument){
		cerr<<"The given private key is not a valid key!"<<endl;
		exit(1);
	}
	string data;
	char buf[1024];
	while(cin){
		cin.read(buf,sizeof(buf));
		int nread=cin.gcount();
		if(nread==0)continue;
		data.append(buf,nread);
	}
	try {
		cout<<Envelope::decrypt(Base64::decode(data),key)<<flush;
	} catch(invalid_argument){
		cerr<<"Private key doesn't match encrypted text!"<<endl;
		exit(1);
	}
}

int main(int argc,char **argv){
	if(argc==3&&strcmp(argv[1],"-k")==0){
		char *endp;
		int keylength=strtol(argv[2],&endp,10);
		if(!argv[2][0]||*endp){
			cerr<<"Invalid number '"<<argv[2]<<"'"<<endl;
			usage();
			return 1;
		}
		mode_keygen(keylength);
	} else if(argc==3&&strcmp(argv[1],"-e")==0){
		mode_encrypt(argv[2]);
	} else if(argc==3&&strcmp(argv[1],"-d")==0){
		mode_decrypt(argv[2]);
	} else {
		usage();
		return 1;
	}
}