CCats/daemon/include/CovertChannel/Protocols/CovertProtocol.hpp

#ifndef COVERTPROTOCOL_H
#define COVERTPROTOCOL_H

#include <fstream>
#include <iostream>
#include <string>
#include <type_traits>

#include <boost/filesystem.hpp>

#include "../../Config.h"
#include "../../Notifications.h"
#include "../../Queue.h"
#include "../ChannelControls.h"

/**
 * @class CovertProtocol
 *
 * An unidirectional Covert Channel protocol with a variable size template.
 *
 * The protocol works unidirectional so the active side uses send to encode data to be sent and the passive side uses receive to extract data.
 *
 * @param N number of bytes which can be used to transmit data
 * @param PASSIVE passive mode
 */
template <int N, bool PASSIVE> class CovertProtocol : public ChannelControls {
	static_assert(N >= 1);

public:
	/**
	 * CovertProtocol constructor
	 */
	CovertProtocol() : fileDirectory(Config::getValue("filedirectory")) {}

	/**
	 * CovertProtocol destructor
	 *
	 * Closes the file
	 */
	~CovertProtocol() { file.close(); }

	/**
	 * send encodes the data into the data array
	 *
	 * 1. Send file name size (length)
	 * 2. Send file name
	 * 3. Send data size (length)
	 * 4. Send data
	 *
	 * @param data must be an array of size N
	 */
	void send(uint8_t *const data) {
		static_assert(!PASSIVE);

		switch (state) {
		case ProtocolState::idle:
			data[0] = 0;
			return;

		case ProtocolState::fileNameSize:
			fileNameSize = fileName.size();
			data[0] = fileNameSize;
			fileNamePosition = 0;
			state = ProtocolState::fileName;
			break;

		case ProtocolState::fileName:
			if (fileNameSize - fileNamePosition >= N) {
				for (int i = 0; i < N; i++) {
					data[i] = fileName.at(fileNamePosition++);
				}
			} else {
				int diff = fileNameSize - fileNamePosition;
				for (int i = 0; i < diff; i++) {
					data[i] = fileName.at(fileNamePosition++);
				}
			}

			if (fileNamePosition == fileNameSize) {
				file.open(fileDirectory + fileName, std::ios::in | std::ios::binary | std::ios::ate);
				if (!file.is_open()) {
					file.close();
					std::cerr << "File \"" << fileName << "\" does exists. Resetting!!!" << std::endl;
					Notifications::newNotification("File \"" + fileName + "\" does not exists. Skipping.");
					// state = ProtocolState::error;
					Queue::channel->reset();
					return;
				}

				dataSize = file.tellg();
				dataPosition = 0;
				file.seekg(std::ios::beg);

				state = ProtocolState::dataSize;
			}

			break;

		case ProtocolState::dataSize:
			if constexpr (N >= 4) {
				for (; dataPosition < 4; dataPosition++) {
					data[dataPosition] = dataSize >> ((3 - dataPosition) * 8);
				}
			} else {
				uint32_t oldDataPosition = dataPosition;
				uint32_t limit = dataPosition + N;
				if (limit > 4) {
					limit = 4;
				}

				for (; dataPosition < limit; dataPosition++) {
					data[dataPosition - oldDataPosition] = dataSize >> ((3 - dataPosition) * 8);
				}
			}

			if (dataPosition == 4) {
				dataPosition = 0;
				state = ProtocolState::data;
			}

			break;

		case ProtocolState::data:
			if (dataSize - dataPosition >= N) {
				file.read((char *)data, N);
				dataPosition += N;
			} else {
				file.read((char *)data, dataSize - dataPosition);
				dataPosition = dataSize;
			}

			std::cout << "sent " << dataPosition << "/" << dataSize << std::endl;

			if (dataPosition == dataSize) {
				file.close();
				state = ProtocolState::idle;
				std::cout << "finished sending file \"" << fileName << "\"" << std::endl;
				Notifications::newNotification("Finished sending file \"" + fileName + "\".");

				// schedule next file transfer
				Queue::schedule();
			}

			break;

		case ProtocolState::error:
			Queue::channel->reset();
			break;
		}

		return;
	}

	/**
	 * receive extracts data from the data array
	 *
	 * 1. Receive file name size (length)
	 * 2. Receive file name
	 * 3. Receive data size (length)
	 * 4. Receive data
	 *
	 * @param data must be an array of size N
	 */
	void receive(const uint8_t *const data) {
		static_assert(PASSIVE);

		switch (state) {
		case ProtocolState::idle:
			if (data[0] == 0) {
				return;
			}

			// no break because the first data received is the filename size
		case ProtocolState::fileNameSize:
			Notifications::newNotification("Incoming transmission.");
			std::cout << "incoming file transmission" << std::endl;

			fileNameSize = data[0];
			fileNamePosition = 0;
			fileName = "";
			transferStart = std::time(nullptr);
			state = ProtocolState::fileName;
			break;

		case ProtocolState::fileName:
			if (fileNameSize - fileNamePosition >= N) {
				uint8_t oldFileNamePosition = fileNamePosition;
				for (; fileNamePosition < oldFileNamePosition + N; fileNamePosition++) {
					fileName += data[fileNamePosition - oldFileNamePosition];
				}
			} else {
				uint8_t oldFileNamePosition = fileNamePosition;
				for (; fileNamePosition < fileNameSize; fileNamePosition++) {
					fileName += data[fileNamePosition - oldFileNamePosition];
				}
			}

			if (fileNamePosition == fileNameSize) {
				file.open(fileDirectory + fileName, std::ios::in);
				if (file.is_open()) {
					file.close();
					std::cerr << "File \"" << fileName << "\" already exists. Resetting!!!" << std::endl;
					Notifications::newNotification("File \"" + fileName + "\" already exists. Skipping.");
					// state = ProtocolState::error;
					Queue::channel->reset();
					return;
				}

				file.close();

				file.open(fileDirectory + fileName, std::ios::out | std::ios::binary | std::ios::app);
				if (!file.is_open()) {
					std::cerr << "File \"" << fileName << "\" could not be opened! Resetting!!!" << std::endl;
					Notifications::newNotification("File \"" + fileName + "\" could not be written. Skipping.");
					// state = ProtocolState::error;
					Queue::channel->reset();
					return;
				}

				dataSize = 0;
				dataPosition = 0;
				state = ProtocolState::dataSize;

				std::cout << "starting receiving file \"" << fileName << "\"" << std::endl;
				Notifications::newNotification("Receiving file \"" + fileName + "\".");
			}

			break;

		case ProtocolState::dataSize:
			if constexpr (N >= 4) {
				for (; dataPosition < 4; dataPosition++) {
					dataSize = dataSize | data[dataPosition] << ((3 - dataPosition) * 8);
				}
			} else {
				uint32_t oldDataPosition = dataPosition;
				uint32_t limit = dataPosition + N;
				if (limit > 4) {
					limit = 4;
				}

				for (; dataPosition < limit; dataPosition++) {
					dataSize = dataSize | data[dataPosition - oldDataPosition] << ((3 - dataPosition) * 8);
				}
			}

			if (dataPosition == 4) {
				dataPosition = 0;
				state = ProtocolState::data;
			}

			break;
		case ProtocolState::data:

			if (dataSize - dataPosition >= N) {
				file.write((char *)data, N);
				dataPosition += N;
			} else {
				file.write((char *)data, dataSize - dataPosition);
				dataPosition = dataSize;
			}

			std::cout << "received " << dataPosition << "/" << dataSize << std::endl;

			if (dataPosition == dataSize) {
				file.close();
				state = ProtocolState::idle;

				std::cout << "finished receiving file \"" << fileName << "\"" << std::endl;
				Notifications::newNotification("Finished receiving file \"" + fileName + "\".");
			}

			break;
		case ProtocolState::error:
			//
			break;
		}
	}

	/* ChannelControls */

	/**
	 * Starts sending a file.
	 *
	 * Starts sending a file if no transmission is running and the file exists.
	 *
	 * @param fileName name of the file in the file directory
	 * @return true - file will be sent | false - file was not accepted
	 */
	virtual bool sendFile(const std::string &fileName) {
		if constexpr (PASSIVE) {
			return false;
		}

		if (state != ProtocolState::idle || file.is_open()) {
			return false;
		}

		file.open(fileDirectory + fileName, std::ios::in);
		if (file.is_open()) {
			file.close();
			this->fileName = fileName;
			transferStart = std::time(nullptr);
			state = ProtocolState::fileNameSize;
			Notifications::newNotification("Start sending file \"" + fileName + "\".");
			std::cout << "starting sending file \"" << fileName << "\"" << std::endl;
			return true;
		} else {
			return false;
		}
	}

	/**
	 * Get the progress
	 *
	 * @return progress counters
	 */
	virtual std::pair<uint32_t, uint32_t> getProgress() {
		if (state == ProtocolState::data) {
			return std::pair<uint32_t, uint32_t>(dataPosition, dataSize);
		} else {
			return std::pair<uint32_t, uint32_t>(0, 0);
		}
	}

	/**
	 * Test if a transfer is running
	 *
	 * @return true - a transfer runs | false - no transfer runs
	 */
	virtual bool isTransferRunning() { return state != ProtocolState::idle; }

	/**
	 * Resets protocol state
	 */
	virtual void reset() {
		file.close();

		if (PASSIVE && state == ProtocolState::data) {
			// delete file if in passive mode and data has been received
			deleteFile(fileName);
			Notifications::newNotification("Transfer of file \"" + fileName + "\" was aborted.");
		}

		state = ProtocolState::idle;
		dataSize = 0;
		dataPosition = 0;
		fileNameSize = 0;
		fileNamePosition = 0;
		fileName = "";

		Queue::schedule();
	}

	/**
	 * Test if a transfer is running
	 *
	 * @return true - a transfer runs | false - no transfer runs
	 */
	virtual std::time_t getTransferStart() { return transferStart; }

	/**
	 * Get file name of the file which is currently be sent or received.
	 *
	 * @return file name
	 */
	virtual std::string getFileName() { return fileName; };

	/* =============== */

private:
	/**
	 * folder of the files
	 */
	const std::string fileDirectory;

	/**
	 * States of the data transmission
	 *
	 * @warning error state cannot be recovered!!!
	 */
	enum struct ProtocolState { idle, fileNameSize, fileName, dataSize, data, error };

	/**
	 * State of the data transmission
	 */
	ProtocolState state = ProtocolState::idle;

	/**
	 * size of the file to be sent/received
	 */
	uint32_t dataSize;

	/**
	 * position in the file to be sent/received
	 */
	uint32_t dataPosition;

	/**
	 * file to be sent/received
	 */
	std::fstream file;

	/**
	 * size of the file name
	 */
	uint8_t fileNameSize;

	/**
	 * position of the file name so far sent/received
	 */
	uint8_t fileNamePosition;

	/**
	 * Deletes a file
	 */
	void deleteFile(const std::string &fileName) {
		std::string fname = this->fileDirectory;
		fname.append(fileName);

		if (boost::filesystem::exists(fname)) {
			boost::filesystem::remove(fname);
		}
	}
};

#endif