/**
 * 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.
 **/
#if ! defined(LIBSECRECY_RAWKEY_HPP)
#define LIBSECRECY_RAWKEY_HPP

#include <libsecrecy/GPGMEContext.hpp>
#include <libsecrecy/Yarrow.hpp>
#include <libsecrecy/GCMStreamBase.hpp>
#include <nettle/aes.h>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <mutex>
#include <map>

// #include <iostream> // remove

namespace libsecrecy
{

	struct RawKeyBase : public NumToHex
	{
		virtual ~RawKeyBase() {}
		virtual std::string hashToHex() const = 0;

		static constexpr std::size_t getDigestSize()
		{
			return SHA256_DIGEST_SIZE;
		}

		static std::string computeHash(char const * a, char const * e)
		{
			sha256_ctx shactx;
			sha256_init(&shactx);
			sha256_update(&shactx,e-a,reinterpret_cast<uint8_t const *>(a));
			std::shared_ptr<uint8_t[]> H(new uint8_t[getDigestSize()]);
			sha256_digest(&shactx,getDigestSize(),H.get());
			return std::string(H.get(),H.get()+getDigestSize());
		}

		static std::string computeHash(uint8_t const * a, uint8_t const * e)
		{
			return computeHash(
				reinterpret_cast<char const *>(a),
				reinterpret_cast<char const *>(e)
			);
		}

		static std::string computeHash(std::string const & data)
		{
			char const * a = data.c_str();
			char const * e = a + data.size();
			return computeHash(a,e);
		}

		static std::string hashToHex(uint8_t const * ua, uint8_t const * ue)
		{
			std::ostringstream ostr;
			while ( ua != ue )
			{
				uint8_t const v = *(ua++);
				ostr.put(numToHex((v >> 4) & 0xF));
				ostr.put(numToHex((v >> 0) & 0xF));
			}
			return ostr.str();
		}

		static std::string getKeyDirectory()
		{
			static char const * libsecrecy_keydir_envvar = "LIBSECRECY_KEYDIR";
			static char const * HOME_envvar = "HOME";
			static char const * libsecrecy_default_dir = ".libsecrecy";

			char const * keydir = getenv(libsecrecy_keydir_envvar);
			std::string skeydir;

			if ( keydir )
			{
				skeydir = keydir;
			}
			else
			{
				char const * home = getenv(HOME_envvar);

				if ( ! home )
					throw std::runtime_error("libsecrecy::Environment variable HOME is not set");

				std::string const shome(home);
				std::filesystem::path pathhome(shome);
				std::filesystem::path pathdefdir(libsecrecy_default_dir);
				pathhome /= pathdefdir;
				skeydir = static_cast<std::string>(pathhome);
			}

			return skeydir;
		}

		static std::string getKeyHashDirectory()
		{
			std::string const secdir = getKeyDirectory();
			std::string const hashext = "hash";
			std::filesystem::path secpath(secdir);
			std::filesystem::path hashpath(hashext);
			secpath /= hashpath;
			return static_cast<std::string>(secpath);
		}

		static std::string getKeyNameDirectoryName()
		{
			return "name";
		}

		static std::string getKeyNameDirectory()
		{
			std::string const secdir = getKeyDirectory();
			std::string const hashext = getKeyNameDirectoryName();
			std::filesystem::path secpath(secdir);
			std::filesystem::path hashpath(hashext);
			secpath /= hashpath;
			return static_cast<std::string>(secpath);
		}

		static void createDirectory(std::string const & skeydir)
		{
			std::filesystem::path keypath(skeydir);
			bool running = true;

			while ( running )
			{
				if ( ! std::filesystem::status_known(std::filesystem::status(keypath)) )
				{
					std::ostringstream ostr;
					ostr << "libsecrecy::RawKey::createKeyDirectory: Unable to access key director " << skeydir;
					throw std::runtime_error(ostr.str());
				}

				if ( std::filesystem::exists(keypath) )
				{
					if ( std::filesystem::is_directory(keypath) )
						running = false;
					else if ( std::filesystem::is_symlink(keypath) )
						keypath = std::filesystem::read_symlink(keypath);
					else
					{
						std::ostringstream ostr;
						ostr << "libsecrecy::RawKey::createKeyDirectory: " << static_cast<std::string>(keypath)
							<< " is not a directory or symlink";
						throw std::runtime_error(ostr.str());
					}
				}
				else
				{
					bool const ok = std::filesystem::create_directories(keypath);

					if ( ok )
					{
						std::filesystem::permissions(keypath,std::filesystem::perms::owner_all);
						running = false;
					}
					else
					{
						std::ostringstream ostr;
						ostr << "libsecrecy::RawKey::createKeyDirectory: unable to create directory " << static_cast<std::string>(keypath);
						throw std::runtime_error(ostr.str());
					}
				}
			}
		}

		static void createKeyDirectory()
		{
			std::string const skeydir(getKeyDirectory());
			createDirectory(skeydir);
			createDirectory(getKeyHashDirectory());
			createDirectory(getKeyNameDirectory());
		}

		static std::string getKeyFileNameCryptDigest(std::string const & ext)
		{
			std::string const skeydir = getKeyHashDirectory();
			std::filesystem::path keypath(skeydir);
			std::filesystem::path keyfn(ext);
			keypath /= keyfn;
			std::string const fullkeyfn(static_cast<std::string>(keypath));
			return fullkeyfn;
		}

		static std::string getNameFileName(std::string const & keyname)
		{
			std::filesystem::path p(RawKeyBase::getKeyNameDirectory());
			p /= keyname;
			return static_cast<std::string>(p);
		}

		static bool keyNameExists(std::string const & keyname)
		{
			return std::filesystem::exists(getNameFileName(keyname));
		}

		static std::string getDefaultKeyName()
		{
			return "default_key";
		}

		static std::string getDefaultKeyFileName()
		{
			std::filesystem::path p(RawKeyBase::getKeyDirectory());
			p /= getDefaultKeyName();
			return static_cast<std::string>(p);
		}

		static void printKeyNames(std::ostream & out)
		{
			std::string const keynamedir = getKeyNameDirectory();

			for ( auto const & p : std::filesystem::directory_iterator(keynamedir) )
			{
				std::filesystem::path const keypath = p.path();
				std::string hash;

				if ( std::filesystem::exists(keypath) && std::filesystem::is_symlink(keypath) )
				{
					hash = static_cast<std::string>(std::filesystem::path(followSymLink(keypath)).filename());
				}

				out << static_cast<std::string>(keypath.filename()) << "\t" << hash << "\n";
			}
		}

		static std::string followSymLink(std::string namefn, std::size_t const maxfollow = 128)
		{
			std::filesystem::path curpath(namefn);

			std::size_t followed = 0;
			while ( followed++ < maxfollow && std::filesystem::exists(curpath) && std::filesystem::is_symlink(curpath) )
			{
				std::filesystem::path npath = std::filesystem::read_symlink(curpath);

				if ( npath.is_absolute() )
				{
					curpath = npath;
				}
				else
				{
					std::filesystem::path xpath = curpath;
					xpath.remove_filename();
					xpath /= npath;
					xpath = std::filesystem::absolute(xpath).lexically_normal();
					curpath = xpath;
				}
			}

			if ( !std::filesystem::exists(curpath) )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::followSymLink: followed symlink to nonexistent " << static_cast<std::string>(curpath);
				throw std::runtime_error(ostr.str());
			}
			else if ( std::filesystem::is_symlink(curpath) )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::followSymLink: too many levels of symlinks from " << namefn;
				throw std::runtime_error(ostr.str());
			}

			return static_cast<std::string>(curpath);
		}

		static std::string readKeyData(std::string hexhash, bool tryname)
		{
			if ( !hexhash.size() && tryname )
			{
				if ( std::filesystem::exists(getDefaultKeyFileName()) )
				{
					if ( std::filesystem::is_symlink(getDefaultKeyFileName()) )
					{
						std::filesystem::path curpath(followSymLink(getDefaultKeyFileName()));
						curpath = std::filesystem::absolute(curpath);
						hexhash = static_cast<std::string>(curpath.filename());
					}
					else
					{
						std::ostringstream ostr;
						ostr << "libsecrecy::RawKeyBase::construct: default key file name " << getDefaultKeyFileName() << " is not a symbolic link";
						throw std::runtime_error(ostr.str());
					}
				}
				else
				{
					std::ostringstream ostr;
					ostr << "libsecrecy::RawKeyBase::construct: key specification is empty but there is no default key";
					throw std::runtime_error(ostr.str());
				}
			}

			std::string fullkeyname = RawKeyBase::getKeyFileNameCryptDigest(hexhash);

			if ( ! std::filesystem::exists(fullkeyname) && tryname )
			{
				std::string const namefn = getNameFileName(hexhash);

				if ( std::filesystem::exists(namefn) && std::filesystem::is_symlink(namefn) )
				{
					fullkeyname = followSymLink(namefn);
				}
				else
				{
					std::ostringstream ostr;
					ostr << "libsecrecy::RawKeyBase::construct: no key found for " << hexhash;
					throw std::runtime_error(ostr.str());
				}
			}

			if ( ! std::filesystem::exists(fullkeyname) )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::construct: no key found for " << hexhash << " as " << fullkeyname << " does not exist";
				throw std::runtime_error(ostr.str());
			}

			if ( std::filesystem::is_directory(fullkeyname) )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::construct: no key found for \"" << hexhash << "\" as " << fullkeyname << " is a directory";
				throw std::runtime_error(ostr.str());
			}

			std::ifstream istr(fullkeyname,std::ios::binary);

			if ( ! istr.is_open() )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::construct: unable to open key file " << fullkeyname << " for " << hexhash;
				throw std::runtime_error(ostr.str());
			}

			std::ostringstream keyfiledatastr;
			while ( istr && istr.peek() != std::istream::traits_type::eof() )
			{
				int const c = istr.get();
				assert ( c != std::istream::traits_type::eof() );
				keyfiledatastr.put(c);
			}

			std::string const keyfiledata = keyfiledatastr.str();

			if ( keyfiledata.size() < RawKeyBase::getDigestSize() )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::construct: key file data is too short (shorter than digest size)";
				throw std::runtime_error(ostr.str());
			}

			std::string const shadigest = keyfiledata.substr(keyfiledata.size() - RawKeyBase::getDigestSize());
			std::string const compdigest = RawKeyBase::computeHash(keyfiledata.substr(0,keyfiledata.size() - RawKeyBase::getDigestSize()));

			if ( shadigest != compdigest )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKeyBase::construct: key file checksum error";
				throw std::runtime_error(ostr.str());
			}

			return keyfiledata;
		}

		static std::string getMagic()
		{
			return "LIBSECRECY_KEY";
		}

		static void expectMagic(std::istream & in)
		{
			std::string const magic = GCMStreamBase::readString(in);

			if ( magic != getMagic() )
			{
				std::ostringstream estr;
				estr << "libsecrecy::RawKey::expectMagic: mismatch in file identifier";
				throw std::runtime_error(estr.str());
			}
		}
	};

	struct RawKeyObject
	{
		std::shared_ptr<RawKeyBase> ptr;
		std::string cipher;
		std::string name;
		std::string hash;

		RawKeyObject()
		{
		}
		RawKeyObject(
			std::shared_ptr<RawKeyBase> rptr,
			std::string rcipher,
			std::string rname,
			std::string rhash
		) : ptr(rptr), cipher(rcipher), name(rname), hash(rhash)
		{
		}

		bool isNull() const
		{
			return !ptr;
		}

		std::string toString() const
		{
			std::ostringstream ostr;
			ostr << "RawKeyObject("
				<< "key=" << ptr->hashToHex() << ","
				<< "cipher=" << cipher << ","
				<< "name=" << name << ","
				<< "hash=" << hash << ")";
			return ostr.str();
		}
	};

	struct RawKeyBaseCacheInterface
	{
		virtual ~RawKeyBaseCacheInterface() {}
		virtual void addKey(RawKeyObject ptr) = 0;
		virtual RawKeyObject getKey(std::string const & id) = 0;
	};

	struct RawKeyBaseCache : public RawKeyBaseCacheInterface
	{
		std::mutex lock;
		std::map<std::string,RawKeyObject> M;

		RawKeyBaseCache() : lock(), M()
		{

		}
		virtual ~RawKeyBaseCache() {}
		virtual void addKey(RawKeyObject ptr)
		{
			std::lock_guard<std::mutex> slock(lock);

			M[ptr.name] = ptr;

			if ( ptr.hash != ptr.name )
				M[ptr.hash] = ptr;
		}
		virtual RawKeyObject getKey(std::string const & id)
		{
			RawKeyObject ptr;

			{
				std::lock_guard<std::mutex> slock(lock);
				auto const it = M.find(id);

				if ( it != M.end() )
					ptr = it->second;
			}

			return ptr;
		}
	};

	template< ::std::size_t _keysize>
	struct RawKey : public RawKeyBase
	{
		static ::std::size_t const keysize = _keysize;
		typedef RawKey<keysize> this_type;
		static ::std::size_t const digestsize = SHA256_DIGEST_SIZE;
		static ::std::size_t const hashblocksize = SHA256_BLOCK_SIZE;

		private:
		LockedMemory<uint8_t> L;
		std::shared_ptr<uint8_t[]> H;

		public:
		RawKey() : L(keysize), H(new uint8_t[digestsize])
		{
			std::fill(L.begin(),L.end(),0);
		}
		RawKey(uint8_t const * la, uint8_t const * le, uint8_t const * ha, uint8_t const * he) : L(keysize), H(new uint8_t[digestsize])
		{
			if ( le - la != static_cast< std::ptrdiff_t >(keysize) )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKey(): key passed has wrong size";
				throw std::runtime_error(ostr.str());
			}

			if ( he - ha != static_cast< std::ptrdiff_t >(digestsize) )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKey(): digest passed has wrong size";
				throw std::runtime_error(ostr.str());
			}

			std::copy(la,le,L.begin());
			std::copy(ha,he,H.get());
		}
		RawKey(Yarrow & Y)
		: L(keysize), H(new uint8_t[digestsize])
		{
			Y(L.get(),keysize);

			// compute hash for key by
			// - getting more random data
			// - putting it through sha256
			// - storing it in H
			LockedMemory<uint8_t> LH(hashblocksize);
			Y(LH.get(),hashblocksize);
			sha256_ctx shactx;
			sha256_init(&shactx);
			sha256_update(&shactx,hashblocksize,LH.get());
			sha256_digest(&shactx,digestsize,H.get());
			std::fill(LH.begin(),LH.end(),0);
		}
		RawKey(RawKey<keysize> const & K)
		: L(keysize), H(new uint8_t[hashblocksize])
		{
			std::copy(K.L.begin(),K.L.end(),L.begin());
			std::copy(K.H.get(),K.H.get()+digestsize,H.get());
		}
		RawKey<keysize> operator=(RawKey<keysize> const & K)
		{
			if ( this != &K )
			{
				std::copy(K.L.begin(),K.L.end(),L.begin());
				std::copy(K.H.get(),K.H.get()+digestsize,H.get());
			}
			return *this;
		}

		bool operator==(RawKey<keysize> const & K) const
		{
			if ( ! std::equal(L.begin(),L.end(),K.L.begin()) )
				return false;
			if ( ! std::equal(H.get(),H.get()+digestsize,K.H.get()) )
				return false;
			return true;
		}

		bool operator!=(RawKey<keysize> const & K) const
		{
			return ! operator==(K);
		}

		std::shared_ptr< LockedMemory<char> > toHex() const
		{
			return L.toHex();
		}

		uint8_t const * get() const
		{
			return L.get();
		}

		std::size_t size() const
		{
			return keysize;
		}

		std::string hashToHex() const
		{
			return RawKeyBase::hashToHex(H.get(),H.get() + digestsize);
		}

		std::string getHash() const
		{
			return std::string(H.get(),H.get() + digestsize);
		}


		std::string getKeyHashFileName() const
		{
			return getKeyFileNameCryptDigest(hashToHex());
		}


		void writeKey(libsecrecy::GPGMEContext & context, std::string const & cipher, std::string const & keyname)
		{
			std::string const keyfn = getKeyHashFileName();
			std::string const namefn = getNameFileName(keyname);

			if ( std::filesystem::exists(keyfn) )
			{
				std::ostringstream estr;
				estr << "libsecrecy::RawKey::writeKey: file " << keyfn << " already exists.";
				throw std::runtime_error(estr.str());
			}

			if ( std::filesystem::exists(namefn) )
			{
				std::ostringstream estr;
				estr << "libsecrecy::RawKey::writeKey: file " << namefn << " for name " << keyname << " already exists.";
				throw std::runtime_error(estr.str());
			}

			std::ostringstream sstr;

			GCMStreamBase::writeString(sstr,getMagic());
			// write prefix / cipher
			GCMStreamBase::writeString(sstr,cipher);
			// write name value
			GCMStreamBase::writeString(sstr,keyname);
			// write hash value
			GCMStreamBase::writeString(sstr,hashToHex());
			// write encrypted key
			context.encrypt(
				reinterpret_cast<char const *>(L.begin()),
				reinterpret_cast<char const *>(L.end()),
				sstr
			);

			std::string const data = sstr.str();

			// compute checksum
			std::string const chksum = computeHash(data);

			// create the key directory if necessary
			createKeyDirectory();
			std::ofstream ostr(keyfn,std::ios::binary);
			if ( ! ostr.is_open() )
			{
				std::ostringstream estr;
				estr << "libsecrecy::RawKey::writeKey: failed to open file " << keyfn;
				throw std::runtime_error(estr.str());
			}

			ostr.write(data.c_str(),data.size());
			ostr.write(chksum.c_str(),chksum.size());
			ostr.flush();

			if ( ! ostr )
			{
				std::ostringstream estr;
				estr << "libsecrecy::RawKey::writeKey: failed to write key data to " << keyfn;
				throw std::runtime_error(estr.str());
			}

			ostr.close();

			if ( !std::filesystem::path(keyname).parent_path().empty() )
			{
				std::ostringstream estr;
				estr << "libsecrecy::RawKey::writeKey: key name " << keyname << " can not contain a directory path";
				throw std::runtime_error(estr.str());
			}

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

			std::filesystem::current_path(RawKeyBase::getKeyNameDirectory());
			std::string const shex = hashToHex();

			std::filesystem::path ptarget("../hash");
			ptarget /= shex;

			std::filesystem::create_symlink(ptarget,std::filesystem::path(keyname));

			std::filesystem::current_path(lcwd);
		}

		void readKey(libsecrecy::GPGMEContext & context, std::istream & in)
		{
			std::string const shexhash = GCMStreamBase::readString(in);

			if ( shexhash.size() != 2*digestsize )
			{
				std::ostringstream ostr;
				ostr << "libsecrecy::RawKey::readKey: hash part of key name is not compatible with key type";
				throw std::runtime_error(ostr.str());
			}

			std::string::const_iterator it = shexhash.begin();
			uint8_t * h = H.get();

			while ( it != shexhash.end() )
			{
				uint8_t const high = hexToNum(*it++) << 4;
				assert ( it != shexhash.end() );
				uint8_t const low  = hexToNum(*it++) << 0;

				*(h++) = high | low;
			}

			assert ( h == H.get() + digestsize );

			context.decrypt(in,reinterpret_cast<char *>(L.begin()),reinterpret_cast<char *>(L.end()));
		}

	};

	typedef RawKey<AES128_KEY_SIZE> RawAES128Key;
	typedef RawKey<AES192_KEY_SIZE> RawAES192Key;
	typedef RawKey<AES256_KEY_SIZE> RawAES256Key;
};
#endif
