/**
 * Copyright 2020 German Tischler-Höhle
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **/
#include <libsecrecy/GCMInputStream.hpp>
#include <libsecrecy/GCMOutputStream.hpp>

#include <iostream>

#define DEBUG

/**
 * generate a key and store it
 **/
template<typename key_type>
void writeKey(std::string const & gpgrecipient, std::string const & keyname)
{
	libsecrecy::GPGMEContext gpgme;
	gpgme.loadKey(gpgrecipient.c_str(),1 /* secret */);

	libsecrecy::Yarrow Y;
	key_type K(Y);
	K.writeKey(gpgme,keyname);

	std::string const hexhash = K.hashToHex();
	std::cout << hexhash << std::endl;

	#if defined(DEBUG)
	// read back key and check for equality
	std::string const keyfiledata = libsecrecy::RawKeyBase::readKeyData(hexhash,false /* do not try to resolve name */);
	std::istringstream keydataistr(keyfiledata);
	libsecrecy::RawKeyBase::expectMagic(keydataistr);
	std::string const keycipher = libsecrecy::GCMStreamBase::readString(keydataistr);
	std::string const rkeyname = libsecrecy::GCMStreamBase::readString(keydataistr);

	if ( keycipher != K.getPrefix() )
		throw std::runtime_error("Decoded key does not match encoded key (wrong cipher)");

	key_type KB(gpgme,keydataistr);

	bool const eq = (K == KB);

	if ( ! eq )
	{
		throw std::runtime_error("Decoded key does not match encoded key");
	}
	#endif
}


int createKey(char const * progname, int argc, char * argv[])
{
	if ( argc < 3 )
	{
		std::cerr << "usage: " << progname << " createKey <cipher> <gpgid> <keyname>" << std::endl;
		std::cerr << std::endl;
		std::cerr << "ciphers available: AES128, AES192, AES256\n";
		return EXIT_FAILURE;
	}

	std::string const scipher(argv[0]);
	std::string const sgpgid(argv[1]);
	std::string const skeyname(argv[2]);

	if ( scipher == "AES256" )
	{
		writeKey<libsecrecy::AES256EncryptKey>(sgpgid,skeyname);
	}
	else if ( scipher == "AES192" )
	{
		writeKey<libsecrecy::AES192EncryptKey>(sgpgid,skeyname);
	}
	else if ( scipher == "AES128" )
	{
		writeKey<libsecrecy::AES128EncryptKey>(sgpgid,skeyname);
	}
	else
	{
		std::cerr << "Unsupported cipher " << scipher << std::endl;
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}

static std::string getExportedKeyMagic()
{
	return "LIBSECRECY_EXPORTED_KEY";
}

template<typename key_type>
void exportKey(std::string const & sgpgid, std::istream & keydataistr, std::string const & keyname)
{
	libsecrecy::GPGMEContext gpgme;
	key_type KB(gpgme,keydataistr);
	std::shared_ptr< libsecrecy::LockedMemory<char> > const keyhex = KB.keyToHex();

	std::ostringstream ostr;
	libsecrecy::GCMStreamBase::writeString(ostr,getExportedKeyMagic());
	libsecrecy::GCMStreamBase::writeString(ostr,KB.getPrefix());
	libsecrecy::GCMStreamBase::writeString(ostr,keyname);
	libsecrecy::GCMStreamBase::writeString(ostr,KB.hashToHex());

	std::size_t const keysize = keyhex->size();
	std::size_t const digestsize = libsecrecy::RawKeyBase::getDigestSize();

	libsecrecy::GCMStreamBase::writeNumber(ostr,keysize);
	for ( std::size_t i = 0; i < keysize; ++i )
		ostr.put(0);
	for ( std::size_t i = 0; i < digestsize; ++i )
		ostr.put(0);

	std::string const predata = ostr.str();

	libsecrecy::LockedMemory<uint8_t> L(predata.size());
	std::copy(predata.begin(),predata.end(),L.begin());

	std::copy(keyhex->begin(),keyhex->end(),L.end() - (keysize+digestsize));

	std::string const hash = libsecrecy::RawKeyBase::computeHash(L.begin(),L.end() - digestsize);

	assert ( hash.size() == digestsize );

	std::copy(hash.begin(),hash.end(),L.end() - digestsize);

	gpgme.loadKey(sgpgid.c_str(),1 /* secret */);

	std::ostringstream cryptostr;
	gpgme.encrypt(L.begin(),L.end(),cryptostr);

	std::string const cryptdata = cryptostr.str();

	std::ostringstream cryptasciistr;
	cryptasciistr << getExportedKeyMagic() << "\n";

	for ( std::size_t i = 0; i < cryptdata.size(); ++i )
	{
		cryptasciistr.put(libsecrecy::NumToHex::numToHex((cryptdata[i]>>4)&0xF));
		cryptasciistr.put(libsecrecy::NumToHex::numToHex((cryptdata[i]>>0)&0xF));
	}

	cryptasciistr << "\n";

	std::cout << cryptasciistr.str() << std::flush;
}

int exportKey(char const * progname, int argc, char * argv[])
{
	if ( argc < 2 )
	{
		std::cerr << "usage: " << progname << " createKey <keyname> <gpgid>" << std::endl;
		std::cerr << std::endl;
		return EXIT_FAILURE;
	}

	std::string const skeyname(argv[0]);
	std::string const sgpgid(argv[1]);

	// read back key and check for equality
	std::string const keyfiledata = libsecrecy::RawKeyBase::readKeyData(skeyname,true /* try to resolve name */);
	std::istringstream keydataistr(keyfiledata);
	libsecrecy::RawKeyBase::expectMagic(keydataistr);
	std::string const keycipher = libsecrecy::GCMStreamBase::readString(keydataistr);
	std::string const rkeyname = libsecrecy::GCMStreamBase::readString(keydataistr);

	if ( keycipher == "AES256" )
	{
		exportKey<libsecrecy::AES256EncryptKey>(sgpgid,keydataistr,rkeyname);
	}
	else if ( keycipher == "AES192" )
	{
		exportKey<libsecrecy::AES192EncryptKey>(sgpgid,keydataistr,rkeyname);
	}
	else if ( keycipher == "AES128" )
	{
		exportKey<libsecrecy::AES128EncryptKey>(sgpgid,keydataistr,rkeyname);
	}
	else
	{
		std::cerr << "Unsupported cipher " << keycipher << std::endl;
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}

template<typename key_type>
void importKey(
	libsecrecy::LockedMemory<uint8_t> const & L,
	std::shared_ptr<uint8_t[]> const & H,
	std::string const & name,
	std::string const & gpgrecipient
)
{
	std::size_t const digestsize = libsecrecy::RawKeyBase::getDigestSize();
	typedef typename key_type::raw_key_type raw_key_type;
	std::shared_ptr<raw_key_type> kptr(new raw_key_type(L.begin(),L.end(),H.get(),H.get()+digestsize));
	key_type K(kptr);
	libsecrecy::GPGMEContext context;
	context.loadKey(gpgrecipient.c_str(),1 /* secret */);
	K.writeKey(context,name);
}

int importKey(char const * progname, int argc, char * argv[])
{
	if ( argc < 1 )
	{
		std::cerr << "usage: " << progname << " importKey <gpgid> [<name>]" << std::endl;
		std::cerr << std::endl;
		return EXIT_FAILURE;
	}

	std::string const gpgrecipient = argv[0];
	std::vector<std::string> Vlines;
	std::string line;
	while ( std::getline(std::cin,line) && Vlines.size() < 3 )
		Vlines.push_back(line);

	if ( Vlines.size() != 2 )
	{
		std::cerr << "Number of lines is wrong in input" << std::endl;
		return EXIT_FAILURE;
	}

	if ( Vlines[0] != getExportedKeyMagic() )
	{
		std::cerr << "File identifier is wrong in input" << std::endl;
		return EXIT_FAILURE;
	}

	std::ostringstream bindatastr;
	std::string const & hexinput = Vlines[1];

	if ( hexinput.size() % 2 != 0 )
	{
		std::cerr << "Data line has uneven length" << std::endl;
		return EXIT_FAILURE;
	}

	for ( std::size_t i = 0; i < hexinput.size(); )
	{
		uint8_t const uhigh = libsecrecy::NumToHex::hexToNum(hexinput[i++]);
		uint8_t const ulow  = libsecrecy::NumToHex::hexToNum(hexinput[i++]);
		bindatastr.put((uhigh<<4)|ulow);
	}

	std::string const bindata = bindatastr.str();
	std::size_t const digestsize = libsecrecy::RawKeyBase::getDigestSize();

	libsecrecy::GPGMEContext gpgme;
	std::istringstream bindataistr(bindata);
	std::pair< std::shared_ptr<libsecrecy::LockedMemory<char> >, std::size_t> decdataP = gpgme.decrypt(bindataistr);

	if ( decdataP.second < digestsize )
	{
		std::cerr << "Data is too short" << std::endl;
		return EXIT_FAILURE;
	}

	char const * c_bindata_a = decdataP.first->get();
	char const * c_bindata_e = c_bindata_a + decdataP.second - digestsize;
	char const * c_bindata_z = c_bindata_e + digestsize;
	std::string const hash = libsecrecy::RawKeyBase::computeHash(c_bindata_a,c_bindata_e);
	std::string const binhash(c_bindata_e,c_bindata_z);

	if ( binhash != hash )
	{
		std::cerr << "Checksum mismatch: " << binhash << " != " << hash << std::endl;
		return EXIT_FAILURE;
	}

	std::string const smagic = libsecrecy::GCMStreamBase::readString(c_bindata_a,c_bindata_e);

	if ( smagic != getExportedKeyMagic() )
	{
		std::cerr << "File identifier is wrong in decrypted data" << std::endl;
		return EXIT_FAILURE;
	}

	std::string const scipher = libsecrecy::GCMStreamBase::readString(c_bindata_a,c_bindata_e);
	std::string const sname = libsecrecy::GCMStreamBase::readString(c_bindata_a,c_bindata_e);
	std::string const shash = libsecrecy::GCMStreamBase::readString(c_bindata_a,c_bindata_e);

	if ( shash.size() % 2 != 0 )
	{
		std::cerr << "Key hash in decrypted data has uneven size" << std::endl;
		return EXIT_FAILURE;
	}

	std::size_t const hashsize = shash.size()/2;
	std::shared_ptr<uint8_t[]> H(new uint8_t[hashsize]);
	// std::copy(shash.begin(),shash.end(),H.get());

	for ( std::size_t i = 0; i < hashsize; ++i )
	{
		uint8_t const uhigh = libsecrecy::NumToHex::hexToNum(shash[2*i+0]);
		uint8_t const ulow = libsecrecy::NumToHex::hexToNum(shash[2*i+1]);
		H[i] = (uhigh<<4)|ulow;
	}

	std::size_t const lkeysize = libsecrecy::GCMStreamBase::readNumber(c_bindata_a,c_bindata_e);

	if ( lkeysize % 2 )
	{
		std::cerr << "Key size stored is not an even number" << std::endl;
		return EXIT_FAILURE;
	}

	std::size_t const keylen = lkeysize/2;
	std::size_t const keyav = c_bindata_e - c_bindata_a;

	if ( keyav < lkeysize )
	{
		std::cerr << "Key data is truncated" << std::endl;
		return EXIT_FAILURE;
	}

	libsecrecy::LockedMemory<uint8_t> L(keylen);
	for ( std::size_t i = 0; i < keylen; ++i )
	{
		uint8_t const uhigh = libsecrecy::NumToHex::hexToNum(c_bindata_a[2*i+0]);
		uint8_t const ulow  = libsecrecy::NumToHex::hexToNum(c_bindata_a[2*i+1]);
		L[i] = (uhigh<<4)|ulow;
	}

	if ( scipher == "AES128" )
	{
		importKey<libsecrecy::AES128EncryptKey>(L,H,sname,gpgrecipient);
	}
	else if ( scipher == "AES192" )
	{
		importKey<libsecrecy::AES192EncryptKey>(L,H,sname,gpgrecipient);

	}
	else if ( scipher == "AES256" )
	{
		importKey<libsecrecy::AES256EncryptKey>(L,H,sname,gpgrecipient);
	}
	else
	{
		std::cerr << "Unsupported cipher " << scipher << std::endl;
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}

int decrypt(char const *, int, char *[])
{
	std::size_t const putbacksize = 16;
	libsecrecy::RawKeyBaseCache cache;
	libsecrecy::GCMInputStream GIS(std::cin,putbacksize,&cache);
	std::size_t const readblocksize = 64*1024;
	std::shared_ptr<char[]> B(new char[readblocksize]);
	std::size_t s = 0;

	std::chrono::time_point<std::chrono::system_clock> const before = std::chrono::system_clock::now();

	while ( GIS )
	{
		GIS.read(B.get(),readblocksize);

		if ( GIS.bad() )
		{
			std::cerr << "Failed to read encrypted stream" << std::endl;
			return EXIT_FAILURE;
		}

		std::cout.write(B.get(),GIS.gcount());
		s += GIS.gcount();

		if ( ! std::cout )
		{
			std::cerr << "Failed to write output" << std::endl;
			return EXIT_FAILURE;
		}
	}

	std::cout.flush();

	if ( ! std::cout )
	{
		std::cerr << "Failed to write output" << std::endl;
		return EXIT_FAILURE;
	}

	std::chrono::time_point<std::chrono::system_clock> const after = std::chrono::system_clock::now();
	std::chrono::duration<double> const dif = after-before;
	double const secs = dif.count();

	double const bytespersec = s / secs;
	double const mbpersec = bytespersec / (1024.0*1024.0);

	std::cerr << "Decoded " << s << " bytes in time " << secs << "s (" << mbpersec << " mb/s)" << std::endl;

	return EXIT_SUCCESS;
}

int setDefaultKey(char const * progname, int argc, char * argv[])
{
	if ( !(0 < argc) )
	{
		std::cerr << "usage: " << progname << " setDefaultKey <keyname>" << std::endl;
		return EXIT_FAILURE;
	}

	std::string const skeyname = argv[0];

	if ( ! libsecrecy::RawKeyBase::keyNameExists(skeyname) )
	{
		std::cerr << "Key name " << skeyname << " is not available" << std::endl;
		return EXIT_FAILURE;
	}

	std::string const keydir = libsecrecy::RawKeyBase::getKeyDirectory();

	std::filesystem::path const lcwd = std::filesystem::current_path();

	std::filesystem::current_path(keydir);

	if ( std::filesystem::exists(libsecrecy::RawKeyBase::getDefaultKeyName()) )
		std::filesystem::remove(libsecrecy::RawKeyBase::getDefaultKeyName());

	std::filesystem::path target(libsecrecy::RawKeyBase::getKeyNameDirectoryName());
	target /= skeyname;

	std::filesystem::create_symlink(target,libsecrecy::RawKeyBase::getDefaultKeyName());

	std::filesystem::current_path(lcwd);

	return EXIT_SUCCESS;
}

int encrypt(char const * progname, int argc, char * argv[])
{
	if ( !(0 < argc) )
	{
		std::cerr << "usage: " << progname << " encrypt <keyhash>" << std::endl;
		return EXIT_FAILURE;
	}

	std::string const skeyhash(argv[0]);

	libsecrecy::RawKeyBaseCache cache;

	libsecrecy::GCMOutputStream gcmout(
		std::cout,skeyhash,
		libsecrecy::GCMOutputStream::getDefaultBlockSize(),
		libsecrecy::GCMOutputStream::getDefaultAuthDataSize(),
		&cache
	);

	std::size_t const readblocksize = 64*1024;
	std::shared_ptr<char[]> B(new char[readblocksize]);

	while ( std::cin )
	{
		std::cin.read(B.get(),readblocksize);

		if ( std::cin.bad() )
		{
			std::cerr << "Failed to read plain stream" << std::endl;
			return EXIT_FAILURE;
		}

		std::size_t const r = std::cin.gcount();
		gcmout.write(B.get(),r);

		if ( ! gcmout )
		{
			std::cerr << "Failed to write encrypted stream" << std::endl;
			return EXIT_FAILURE;
		}
	}

	gcmout.flush();

	if ( ! gcmout )
	{
		std::cerr << "Failed to write to encrypted stream" << std::endl;
		return EXIT_FAILURE;
	}

	return EXIT_SUCCESS;
}

int listKeys()
{
	libsecrecy::RawKeyBase::printKeyNames(std::cout);
	return EXIT_SUCCESS;
}

int main(int argc, char * argv[])
{
	try
	{
		if ( argc < 2 )
		{
			std::cerr << "usage: " << argv[0] << " <command>" << std::endl;
			std::cerr << std::endl;
			std::cerr << "commands available: createKey, encrypt, decrypt, exportKey, importKey, listKeys, setDefaultKey\n";
			return EXIT_FAILURE;
		}

		std::string const scommand(argv[1]);

		if ( scommand == "createKey" )
			return createKey(argv[0],argc-2,argv+2);
		else if ( scommand == "decrypt" )
			return decrypt(argv[0],argc-2,argv+2);
		else if ( scommand == "encrypt" )
			return encrypt(argv[0],argc-2,argv+2);
		else if ( scommand == "exportKey" )
			return exportKey(argv[0],argc-2,argv+2);
		else if ( scommand == "importKey" )
			return importKey(argv[0],argc-2,argv+2);
		else if ( scommand == "setDefaultKey" )
			return setDefaultKey(argv[0],argc-2,argv+2);
		else if ( scommand == "listKeys" )
			return listKeys();
		else
		{
			std::cerr << "Unknown command " << scommand << std::endl;
			return EXIT_FAILURE;
		}
	}
	catch(std::exception const & ex)
	{
		std::cerr << ex.what() << std::endl;
		return EXIT_FAILURE;
	}
}
