/*****************************************************************//**
	 * @file   miragezip.h
	 * @brief  This file defines the interface for the MirageZip
	 *         back end, and provides functions for setting a password
	 *         and evaluating its strength, zipping and encrypting
	 *         imported files with libzip, and concatenating the
	 *         archive with an imported image file.
	 *
	 * @author Kevin Crepps
	 * @date   August 2024
	 *********************************************************************/
	
	#pragma once

	#include <sstream>
	#include <fstream>
	#include <filesystem>
	#include <memory>

	#include <Windows.h>

	#include "zip.h"

    /**
	* Custom return codes specifying that
	* a given operation:
	*
	* - Succeeded (0)
	* - Failed and aborting is recommended (1)
	* - Failed but program can safely continue (2)
	*/
	#define SUCCESS 0
	#define FAILURE_ABORT 1
	#define FAILURE_CONTINUE 2

	/**
	* Constants used in determining password strength.
	*/
	#define WEAK_SCORE 2
	#define MEDIUM_SCORE 3
	#define STRONG_SCORE 4
	#define MEDIUM_LENGTH 8
	#define STRONG_LENGTH 12
	#define NUM_SETS 4
	

	class MirageZip
	{
	public:
		/**
		* Used to specify path type with SetPath().
		*/
		enum PATH_TYPE
		{
			IMAGE,
			FILE,
			EXPORT
		};

		/**
		* @brief
		* Password containing null terminator is indicative of initial
		* state. Call private function Init() to create app data folder.
		*/
		MirageZip() : password("\0") { Init(); }

		/**
		* @brief
		* Stores error description for later retrieval using GetError().
		*
		* Overloaded for rvalues.
		*/
		inline void SetError(std::string&& arg) noexcept { error = std::move(arg); }

		/**
		* @brief
		* Stores error description for later retrieval using GetError().
		*
		* Overloaded for lvalues.
		*/
		inline void SetError(const std::string& arg) noexcept { error = arg; }

		/**
		* @brief
		* Retrieve exception info and other error messages set
		* internally with SetError().
		*/
		inline std::string GetError() const noexcept { return error; }

		/**
		* @brief
		* User sets the path of an image, file(s), and export location,
		* specified with PATH_TYPE enum defined above.
		*/
		void SetPath(PATH_TYPE, const std::string&) noexcept;

		/**
		* @brief
		* Use point system to evaluate password strength based on
		* character diversity and length.
		*
		* @return
		* Returns password strength score indicating:
		*   weak - fewer than three points
		*   medium - three points
		*   strong - four points
		*/
		unsigned int TestPassword(const std::string&) const noexcept;

		/**
		* @brief
		* Set the password after evaluating, used to encrypt file
		* archive in ZipFile().
		*/
		inline void SetPassword(const std::string& arg) noexcept { password = arg; }

		/**
		* @brief
		* Zip and password-protect files using libzip functions.
		*
		* @return
		* Custom return code defined at the top of miragezip.h.
		*/
		unsigned int ZipFile();

		/**
		* @brief
		* Combines image and archive using Windows copy command
		* with binary flag, issued through CreateProcessA().
		*/
		unsigned int Concatenate() const;

		/**
		* @brief
		* Get path of new zip file - used in HideFile().
		*/
		inline const char* GetArchivePath() const noexcept { return archivePath.c_str(); }

		~MirageZip() {}

	private:
		std::string error,      // Error info buffer, set in SetError()
			archivePath,        // Path to zip; this is set in CreateAppData()
			imagePath,          // Path to user-selected image file
			filePath,           // Path to user-selected file to be zipped
			exportPath,         // Destination of output file
			password;           // Password used to encrypt archive

		struct stat Statinfo;   /* Used for checking whether
								directory exists with stat() */

		/**
		* Character set types used in password evaluation.
		*/
		enum CHAR_TYPE
		{
			LOWER,
			UPPER,
			NUMBERS,
			CHARACTERS
		};

		/**
		* @brief
		* Called from the constructor. CreateAppData is run from here.
		* If unable to create directory, user is notified, program terminates.
		*/
		void Init();

		/**
		* @brief
		* Checks for existence of %LOCALAPPDATA%\MirageZip, attempts to
		* create it if it doesn't exist.
		*
		* @return
		* Custom return code defined at the top of miragezip.h.
		*/
		unsigned int CreateAppData();
	};

	/**
	* @brief
	* Non-member non-friend function attempts to zip file,
	* combine it with image, then purge app data files.
	*/
	extern unsigned int HideFile(MirageZip*);



	/*****************************************************************//**
	* @file   miragezip.cpp
	* @brief  MirageZip implementation.
	*
	* @author Kevin Crepps
	* @date   August 2024
	*********************************************************************/

	#include "miragezip.h"
	
	void MirageZip::Init()
	{
		/* If %LOCALAPPDATA%\MirageZip doesn't exist, and
            attempt to create it fails, notify user and abort */

		if (CreateAppData())
		{
			MessageBoxA(NULL, GetError().c_str(), "Failed to create app data folder.", MB_OK);
			abort();
		}
	}

	unsigned int MirageZip::CreateAppData()
	{
		try
		{
			// Set archivePath member, check whether folder already exists
			std::string path{ getenv("LOCALAPPDATA") };
			path += "\\MirageZip";
			archivePath = path + "\\archive.zip";
			stat(path.c_str(), &Statinfo);

			// If directory doesn't exist, create it
			if (!(Statinfo.st_mode & S_IFDIR))
			{
				std::error_code err;
				err.clear();

				if (!std::filesystem::create_directories(path, err))
				{
					std::stringstream msg("Error code: ");
					msg << err;
					SetError(std::move(msg.str()));
					return FAILURE_ABORT;
				}
			}

			return SUCCESS;
		}
		catch (...)
		{
			SetError("Unknown exception thrown.");
			return FAILURE_ABORT;
		}
	}
	
	void MirageZip::SetPath(PATH_TYPE type, const std::string& path) noexcept
	{
		switch (type)
		{
		case IMAGE:
			imagePath = path;
			break;

		case FILE:
			filePath = path;
			break;
		
		case EXPORT:
			exportPath = path;
		}
	}

	unsigned int MirageZip::TestPassword(const std::string& pw) const noexcept
	{
		bool found[NUM_SETS]{ false, false, false, false };
		char lowerSet[]{ "abcdefghijklmnopqrstuvwxyz" },
			upperSet[]{ "ABCDEFGHIJKLMNOPQRSTUVWXYZ" },
			numberSet[]{ "0123456789" },
			charSet[]{ "!@#$%^&*()-_=+,.<>/?|[]{}:;~\'\"\\" };
		char* sets[NUM_SETS]{ lowerSet, upperSet, numberSet, charSet };
		float points{ 0.0f };

		 /* Compare each character in password to each character set if not yet matched,
			add half a point for each character set with one or more instances
			of lowercase, uppercase, digits or special characters in password   */

		for (auto& c : pw)
		{
			for (unsigned char type = LOWER; type < NUM_SETS; ++type)
			{
				if (!found[type])
				{
					if (strchr(sets[type], c))
					{
						points += 0.5f;
						found[type] = true;
					}
				}
			}
		}

		// If pass is medium length (less than 12), add one point if it has at least three character types
		if (pw.length() >= MEDIUM_LENGTH && pw.length() < STRONG_LENGTH && points >= 1.5f)
			++points;

		// If strong length, add two points if it has at least two character types
		else if (pw.length() >= STRONG_LENGTH && points >= 1.0f)
			points += 2;

		/*     strong - four points
			   medium - three points
				weak - fewer than three points    */
		return (unsigned int)std::floor(points);
	}
	
	unsigned int MirageZip::ZipFile()
	{
	    /* This function will read file data into a buffer, zip the file using the buffer, and encrypt
            the data using 256-bit AES if a password has been set by the user. Error handling is done in
            HideFile()  */
	
		// Create archive and open
		std::remove(archivePath.c_str());
		int errCode = 0;
		zip* archive = zip_open(archivePath.c_str(), ZIP_CREATE, &errCode);
		if (!archive) return FAILURE_ABORT;

		// Open file, get size
		std::ifstream file(filePath, std::ios::binary);
		if (!file.is_open()) return FAILURE_ABORT;
		file.seekg(0, std::ios::end);
		int fileSize = file.tellg();
		file.seekg(0, std::ios::beg);

		// Store file data in resource management object, close file
		std::unique_ptr<char[]> fileData(new char[fileSize]);
		file.read(fileData.get(), fileSize);
		file.close();

		// Get filename from file path
		std::string fileName(filePath);
		size_t strPos = fileName.find_last_of("\\");
		fileName.erase(0, strPos + 1);

		// Zip file using buffer
		zip_source_t* source;
		source = zip_source_buffer(archive, fileData.get(), fileSize, 0);
		if (!source) return FAILURE_ABORT;
		zip_file_add(archive, fileName.c_str(), source, 0);

		// Encrypt archive if password was set
		if (password != "\0")
			zip_file_set_encryption(archive, 0, ZIP_EM_AES_256, password.c_str());
		
		// Close archive
		zip_close(archive);

		return SUCCESS;
	}

	unsigned int MirageZip::Concatenate() const
	{
		/* Execute Windows system command to combine files, e.g.:

            COPY /B <image path> + <archive path> <destination path>

            Error handling is done in HideFile()   */
		
		STARTUPINFOA info{ sizeof(info) };
		PROCESS_INFORMATION processInfo;
		std::stringstream cmd{ "" };
		cmd << "/c COPY /B \"" << imagePath << "\" + \"" << archivePath << "\" \"" << exportPath << "\"";

		if (CreateProcessA("C:\\Windows\\System32\\cmd.exe",
					cmd.str().data(), NULL, NULL, TRUE, 0, NULL, NULL, &info, &processInfo))
		{
			WaitForSingleObject(processInfo.hProcess, INFINITE);
			CloseHandle(processInfo.hProcess);
			CloseHandle(processInfo.hThread);
			return SUCCESS;
		}

		return FAILURE_ABORT;
	}
	
	// -- External --
	unsigned int HideFile(MirageZip* obj)
	{
		//  Attempt to zip selected file
		try
		{
			obj->ZipFile();
		}
		catch (const std::exception& e)
		{
			obj->SetError(e.what());
			return FAILURE_ABORT;
		}
		catch (...)
		{
			obj->SetError("Exception thrown while attempting to zip file.");
			return FAILURE_ABORT;
		}

		// Attempt to concatenate image and temp archive
		try
		{
			obj->Concatenate();
		}
		catch (const std::exception& e)
		{
			obj->SetError(e.what());
			return FAILURE_ABORT;
		}
		catch (...)
		{
			obj->SetError("Exception thrown while attempting to concatenate files.");
			return FAILURE_ABORT;
		}

		// Attempt to delete temp archive
		try
		{
			std::remove(obj->GetArchivePath());
		}
		catch (const std::exception& e)
		{
			obj->SetError(e.what());
			return FAILURE_CONTINUE;
		}
		catch (...)
		{
			obj->SetError("Exception thrown while attempting to remove temp archive.");
			return FAILURE_CONTINUE;
		}

		return SUCCESS;
	}