#include "../include/fileman.h"
#include "../include/base64.h"

#include <cstdio>
#include <cstring>
#include <iostream>

using std::string;
using std::vector;

FileMan::FileMan() {
	isputting = false;
	islisting = false;
	keyenabled = false;
	cryptoreadye = false;
	cryptoreadyd = false;
	ERR_load_crypto_strings();
}

FileMan::~FileMan() {
	cancelGet();
	cancelPut();
	cancelList();
	if (keyenabled)
		closeKey();
	ERR_free_strings();
}

bool FileMan::isGetting() { return getfile.is_open(); }

bool FileMan::isPutting() { return isputting; }

bool FileMan::isListing() { return islisting; }
bool FileMan::isListingSimple() { return islisting && !extendedListing; }
bool FileMan::isListingExtended() { return islisting && extendedListing; }

bool FileMan::isEncrypted() { return keyenabled; }

bool FileMan::openPut(const string &path) {
	std::fstream keyfile;
	putpath = path;
	putname = pathToFilename(path);
	putfile.open(path, std::ios::ate | std::ios::binary);
	if (putfile.is_open()) {
		size_t size = putfile.tellg();
		//~ putsize = size;
		// increment by one for CCATs signature in its own chunk
		//~ putchunks = 1 + (size / max_read_len + ((size % max_read_len) ? 1 : 0));
		if (keyenabled) {
			// fill IV for file
			keyfile.open("/dev/urandom", std::ios::binary | std::ios::in);
			keyfile.read((char *)iv, sizeof(iv));
			keyfile.close();
		}
		putfile.seekg(0);

		// if crypto enabled, encrypt now
		if (keyenabled) {
			cipherpath = path + ".tmp";

			if (!!std::ifstream(cipherpath)) {
				// close put if tmp file already exists
				closePut();
				return false;
			}

			std::ofstream cipherFile(cipherpath, std::ios::binary);
			if (!cipherFile) {
				closePut();
				return false;
			}

			// resize buffer to also fit IV and tag
			size_t additionalsize = sizeof(iv) + sizeof(tag);

			// skip first 32 bytes. Those are needed for the header
			cipherFile.seekp(additionalsize);

			if (!initCryptoE()) {
				// failed to init crypto, do not continue
				return false;
			}

			unsigned char *plain = new unsigned char[max_read_len];
			unsigned char *cipher = new unsigned char[max_read_len];
			int cipherlen;

			// prepend signature
			memcpy(plain, signature, sizeof(signature));

			if (!EVP_EncryptUpdate(cryptctxe, cipher, &cipherlen, plain, sizeof(signature))) {
				setOpensslError();
				std::cerr << __PRETTY_FUNCTION__ << " failed to update " << getOpensslError() << std::endl;
			}

			cipherFile.write((char *)cipher, cipherlen);

			while ((std::size_t)(putfile.tellg()) + max_read_len < size) {
				putfile.read((char *)plain, max_read_len);

				if (!EVP_EncryptUpdate(cryptctxe, cipher, &cipherlen, plain, max_read_len)) {
					setOpensslError();
					std::cerr << __PRETTY_FUNCTION__ << " failed to update " << getOpensslError() << std::endl;
				}

				cipherFile.write((char *)cipher, cipherlen);
			}

			if (putfile.tellg() < size - 1) {
				int rest = size - putfile.tellg();
				putfile.read((char *)plain, rest);

				if (!EVP_EncryptUpdate(cryptctxe, cipher, &cipherlen, plain, rest)) {
					setOpensslError();
					std::cerr << __PRETTY_FUNCTION__ << " failed to update " << getOpensslError() << std::endl;
				}

				cipherFile.write((char *)cipher, cipherlen);
			}

			if (!EVP_EncryptFinal_ex(cryptctxe, cipher, &cipherlen)) {
				setOpensslError();
				std::cerr << __PRETTY_FUNCTION__ << " failed to finalize " << getOpensslError() << std::endl;
			}
			// obtain tag
			if (!EVP_CIPHER_CTX_ctrl(cryptctxe, EVP_CTRL_GCM_GET_TAG, 16, tag)) {
				setOpensslError();
				std::cerr << __PRETTY_FUNCTION__ << " failed to get tag " << getOpensslError() << std::endl;
			}

			// prepend IV and tag
			cipherFile.seekp(0);
			cipherFile.write((char *)iv, sizeof(iv));
			cipherFile.write((char *)tag, sizeof(tag));

			cipherFile.close();
			putfile.close();

			delete[] cipher;
			delete[] plain;

			// now open cipher file as put file
			putfile.open(cipherpath, std::ios::ate | std::ios::binary);
			if (putfile.is_open()) {
				size = putfile.tellg();
				putfile.seekg(0);
			} else {
				closePut();
				return false;
			}

			deinitCryptoE();
		} else {
			size += sizeof(signature);
		}

		// calculate chunks
		putsize = size;
		putchunksRemaining = putchunks = size / max_read_len + (size % max_read_len > 0 ? 1 : 0);

		isputting = true;

		return true;
	}
	return false;
}

bool FileMan::openGet(const string &path) {
	getpath = path;
	getname = pathToFilename(path);
	getchunks = 0;
	getchunksRemaining = 0;
	getfile.open(path, std::ios::app | std::ios::binary | std::ios::out);
	if (getfile.tellp() != std::ios::beg) {
		closeGet();
		return false;
	}
	return true;
}

bool FileMan::openList(bool extended) {
	if (isListing()) {
		return false;
	}
	listdata = vector<Json::Value>();
	extendedListing = extended;
	islisting = true;
	return true;
}

bool FileMan::openKey(const string &path) {
	std::ifstream keyfile;
	//~ std::cerr << __PRETTY_FUNCTION__ << " begin" << std::endl;
	if (isPutting() || isGetting())
		return false; // do not enable key mid-operation
	if (keyenabled)
		closeKey();
	keyfile.open(path, std::ios::ate | std::ios::binary);
	//~ std::cerr << __PRETTY_FUNCTION__ << " open keyfile" << std::endl;
	if (keyfile.is_open()) {
		//~ std::cerr << __PRETTY_FUNCTION__ << " keyfile open " << keyfile.tellg() << " " << sizeof(key) << std::endl;
		if (keyfile.tellg() == sizeof(key)) {
			keyfile.seekg(0);
			//~ std::cerr << "keyfile is at " << keyfile.tellg() << std::endl;
			keyfile.read((char *)key, sizeof(key));
			keyfile.close();
			keyenabled = true;
			return true;
		}
	}
	return false;
}

void FileMan::closePut() {
	putfile.close();
	putname = "";
	putpath = "";
	putchunks = 0;
	putchunksRemaining = 0;
	isputting = false;
	memset(iv, 0, sizeof(iv));
	memset(tag, 0, sizeof(tag));

	if (keyenabled && !!std::ifstream(cipherpath)) {
		// delete tmp cipher file if it exists
		std::remove(cipherpath.c_str());
	}
}

void FileMan::closeGet() {
	getfile.close();
	getname = "";
	getpath = "";
	getchunks = 0;
	getchunksRemaining = 0;
	memset(iv, 0, sizeof(iv));
	memset(tag, 0, sizeof(tag));
}

void FileMan::closeList() {
	listdata = vector<Json::Value>();
	islisting = false;
	listchunks = 0;
	listchunksRemaining = 0;
}

bool FileMan::closeKey() {
	if (isPutting() || isGetting())
		return false;
	if (keyenabled) {
		memset(key, 0, sizeof(key));
		keyenabled = false;
	}
	return true;
}

void FileMan::cancelPut() {
	if (isPutting()) {
		closePut();
	}
}

void FileMan::cancelGet() {
	if (isGetting()) {
		closeGet();
		std::remove(getname.c_str());
	}
}

void FileMan::cancelList() {
	if (isListing()) {
		closeList();
	}
}

void FileMan::writeGet(const vector<char> data) {
	if (getchunksRemaining == getchunks - 1 && !keyenabled) {
		// check if signature matches
		if (memcmp(data.data(), signature, 4)) {
			// mismatch, encrypted file without enabled key, write as is
			getfile.write(data.data(), data.size());
		} else {
			// skip signature for unencrypted files
			getfile.write(data.data() + 4, data.size() - 4);
		}
	} else {
		getfile.write(data.data(), data.size());
	}
	getchunksRemaining--;
}

void FileMan::writeBase64(string data) {
	vector<char> decode = base64::decodeVector(data);
	keyenabled ? writeEnc(decode) : writeGet(decode);
}

void FileMan::putListData(vector<Json::Value> files) {
	if (!extendedListing) {
		listdata.insert(listdata.end(), files.begin(), files.end());
	} else {
		for (Json::Value file : files) {

			Json::Value fileInfo;
			fileInfo["name"] = file["name"];
			fileInfo["size"] = file["size"];

			string headString = file["head"].asString();
			decryptability decr = isDecryptable(base64::decodeVector(headString));
			if (decr == FileMan::unknown) {
				fileInfo["encrypted"] = "unknown";
			} else if (decr == FileMan::undecryptable) {
				fileInfo["encrypted"] = "undecryptable";
			} else if (decr == FileMan::decryptable) {
				fileInfo["encrypted"] = "decryptable";
			} else {
				fileInfo["encrypted"] = "unencrypted";
			}

			listdata.push_back(fileInfo);
		}
	}
	listchunksRemaining--;
}

vector<Json::Value> FileMan::getListData() { return listdata; }

int FileMan::getGetChunks() { return getchunks; }

int FileMan::getGetRemainingChunks() { return getchunksRemaining; }

void FileMan::setGetChunks(int chunks) {
	getchunks = chunks;
	getchunksRemaining = chunks - 1;
}

void FileMan::setListChunks(int chunks) {
	listchunks = chunks;
	listchunksRemaining = chunks - 1;
}

vector<char> FileMan::readPut() {
	int chunk = putchunks - putchunksRemaining;
	vector<char> data;

	int toread = max_read_len;
	if (putchunksRemaining == 1) {
		toread = putsize - (putchunks - 1) * max_read_len;
	}

	data.resize(toread);

	if (chunk == 0 && !keyenabled) {
		memcpy(&data.front(), signature, sizeof(signature));
		putfile.read(&data.front() + sizeof(signature), toread - sizeof(signature));
	} else {
		putfile.read(&data.front(), toread);
	}

	putchunksRemaining--;
	return data;
}

string FileMan::readBase64() { return base64::encodeVector(readPut()); }

string FileMan::getPutName() { return putname; }

string FileMan::getGetName() { return getname; }

int FileMan::getPutChunks() { return putchunks; }

int FileMan::getPutRemainingChunks() { return putchunksRemaining; }

int FileMan::getPutSize() { return putsize; }

int FileMan::getListRemainingChunks() { return listchunksRemaining; }

int FileMan::getListChunks() { return listchunks; }

string FileMan::pathToFilename(string path) {

	int lastFoundIndex = -1;

	for (int currentIndex = path.find("/"); currentIndex != string::npos; currentIndex = path.find("/", currentIndex + 1)) {

		// check if the "/" was escaped
		if (currentIndex > 0 && path[currentIndex - 1] == '\\')
			continue;

		/*  // check unnecessary, because error occurs when trying to open file
		   anyways

		    // check if the "/" is at the end of path name
		    if (currentIndex + 1 == path.length()) {
		      // ERROR: INVALID FILENAME, BECAUSE ENDS WITH "/"
		    }
		*/
		// otherwise we found a valid "/"
		lastFoundIndex = currentIndex;
	}

	return path.substr(lastFoundIndex + 1);
}

string FileMan::getOpensslError() {
	pendingerr = false;
	return opensslerr;
}

void FileMan::setOpensslError() {
	opensslerr = ERR_error_string(ERR_get_error(), NULL);
	pendingerr = true;
}

bool FileMan::initCryptoE() {
	// try to initialize crypto context
	if (!cryptoreadye) {
		//~ std::cerr << __PRETTY_FUNCTION__ << " init crypto" << std::endl;
		if (!(cryptctxe = EVP_CIPHER_CTX_new())) {
			setOpensslError();
			return false;
		}
		if (!EVP_EncryptInit_ex(cryptctxe, EVP_aes_256_gcm(), NULL, key, iv)) {
			setOpensslError();
			return false;
		}
		cryptoreadye = true;
	}
	return true;
}

bool FileMan::initCryptoD() {
	// try to initialize crypto context
	if (!cryptoreadyd) {
		//~ std::cerr << __PRETTY_FUNCTION__ << " init crypto" << std::endl;
		if (!(cryptctxd = EVP_CIPHER_CTX_new())) {
			setOpensslError();
			return false;
		}
		if (!EVP_DecryptInit_ex(cryptctxd, EVP_aes_256_gcm(), NULL, key, iv)) {
			setOpensslError();
			return false;
		}
		cryptoreadyd = true;
	}
	return true;
}

void FileMan::deinitCryptoE() {
	if (cryptoreadye) {
		EVP_CIPHER_CTX_free(cryptctxe);
		cryptctxe = NULL;
		cryptoreadye = false;
	}
}

void FileMan::deinitCryptoD() {
	if (cryptoreadyd) {
		EVP_CIPHER_CTX_free(cryptctxd);
		cryptctxd = NULL;
		cryptoreadyd = false;
	}
}

void FileMan::writeEnc(const vector<char> data) {
	//~ std::cerr << __PRETTY_FUNCTION__ << " begin" << std::endl;
	writeGet(data);
	if (getchunksRemaining < 0) {
		// loaded everything, try to decrypt
		unsigned char *cipher, *plain;
		int plainlen = 0, finallen = 0;
		getfile.close();
		getfile.open(getpath, std::ios::binary | std::ios::in | std::ios::ate);
		size_t size = getfile.tellg();
		if (size < (sizeof(iv) + sizeof(tag))) {
			// avoid underflow with files that are too small
			return;
		}
		size -= (sizeof(iv) + sizeof(tag));
		cipher = new unsigned char[size];
		plain = new unsigned char[size];
		getfile.seekg(0);
		getfile.read((char *)iv, sizeof(iv));
		getfile.read((char *)tag, sizeof(tag));
		getfile.read((char *)cipher, size);
		getfile.close();

		if (initCryptoD()) {
			if (!EVP_DecryptUpdate(cryptctxd, plain, &plainlen, cipher, size)) {
				setOpensslError();
				std::cerr << __PRETTY_FUNCTION__ << " failed to update " << getOpensslError() << std::endl;
			}
			//~ std::cerr << __PRETTY_FUNCTION__ << " decrypted" << std::endl;
			if (!EVP_CIPHER_CTX_ctrl(cryptctxd, EVP_CTRL_GCM_SET_TAG, 16, tag)) {
				setOpensslError();
				std::cerr << __PRETTY_FUNCTION__ << " failed to finalize " << getOpensslError() << std::endl;
			}
			//~ std::cerr << __PRETTY_FUNCTION__ << " set tag" << std::endl;
			if (0 < EVP_DecryptFinal_ex(cryptctxd, plain + plainlen, &finallen)) {
				plainlen += finallen;
				//~ std::cerr << __PRETTY_FUNCTION__ << " finalized with len " << plainlen << std::endl;
				getfile.close();
				// check signature
				if (memcmp(plain, signature, 4)) {
					std::cerr << __PRETTY_FUNCTION__ << " signature mismatch" << std::endl;
				} else {
					// signatur matches, skip it and dump the rest to disk
					getfile.open(getpath, std::ios::binary | std::ios::out | std::ios::trunc);
					getfile.write((char *)(plain + sizeof(signature)), plainlen - sizeof(signature));
					getfile.close();
				}
			} else {
				setOpensslError();
				std::cerr << __PRETTY_FUNCTION__ << " failed to finalize " << getOpensslError() << std::endl;
			}
			deinitCryptoD();
		}
		delete[] cipher;
		delete[] plain;
	}
}

FileMan::decryptability FileMan::isDecryptable(const vector<char> data) {
	unsigned char plain[sizeof(signature)] = {0}, *cipher;
	int plainlen;
	// check if signature matches in plaintext
	if (data.size() < sizeof(signature) || memcmp(data.data(), signature, sizeof(signature))) {
		// either size or signature dont match
		if (data.size() >= (sizeof(signature) + sizeof(iv) + sizeof(tag)) && keyenabled) {
			// enough data, key is enabled
			memcpy(iv, data.data(), sizeof(iv));
			memcpy(tag, data.data() + sizeof(iv), sizeof(tag));
			cipher = (unsigned char *)(data.data() + sizeof(iv) + sizeof(tag));

			if (initCryptoD()) {
				// crypto is ready
				if (!EVP_DecryptUpdate(cryptctxd, plain, &plainlen, cipher, sizeof(signature))) {
					setOpensslError();
					std::cerr << __PRETTY_FUNCTION__ << " failed to update " << getOpensslError() << std::endl;
				}

				deinitCryptoD();

				// check again if signature matches
				if (memcmp(plain, signature, sizeof(signature))) {
					// still no dice
					return undecryptable;
				}

				// yup, got good data
				return decryptable;
			}
		}
		// size mismatch or no key enabled, consider unknown
		return unknown;
	}
	// it does, assume plaintext
	return plaintext;
}