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

#ifndef COVERTPROTOCOLBIDIRECTIONAL_H
#define COVERTPROTOCOLBIDIRECTIONAL_H

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

#include "../ChannelControls.h"
#include "CovertProtocol.hpp"

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

public:
	/**
	 * CovertProtocolBidirectional constructor
	 */
	CovertProtocolBidirectional() : segment(PASSIVE ? 1 : 0) {
		lastData = new uint8_t[N]();
		protocol.reset();
	}

	/**
	 * CovertProtocol destructor
	 */
	~CovertProtocolBidirectional() { delete[](lastData); }

	/**
	 * send encodes the data into the data array
	 *
	 * @param data must be an array of size N
	 */
	void send(uint8_t *const data) {
		if (sendResetFlag) {
			sendResetFlag = false;
			data[0] = (segment << 6) | 1;
			std::memset(data + 1, 0, N - 1);
			return;
		}

		if (sendSegment()) {
			std::memcpy(data, lastData, N);
			return;
		}

		if constexpr (!PASSIVE) {
			protocol.send(data + 1);
		}

		data[0] = (segment << 6);

		// save data to lastData
		std::memcpy(lastData, data, N);
	}

	/**
	 * receive extracts data from the data array
	 *
	 * @param data must be an array of size N
	 */
	void receive(const uint8_t *const data) {
		uint8_t seg = data[0] >> 6;
		bool reset = data[0] & 1;

		if (reset) {
			resetState();
			return;
		}

		if (receiveSegment(seg)) {
			return;
		}

		if constexpr (PASSIVE) {
			protocol.receive(data + 1);
		}
	}

	/* 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) { return protocol.sendFile(fileName); }

	/**
	 * Get the progress
	 *
	 * @return progress counters
	 */
	virtual std::pair<uint32_t, uint32_t> getProgress() { return protocol.getProgress(); }

	/**
	 * Test if a transfer is running
	 *
	 * @return true - a transfer runs | false - no transfer runs
	 */
	virtual bool isTransferRunning() { return protocol.isTransferRunning(); }

	/**
	 * Resets state and sets reset flag so a reset signal is sent in the next packet
	 */
	virtual void reset() {
		sendResetFlag = true;
		resetState();
	}

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

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

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

private:
	/**
	 * current segment counter
	 */
	uint8_t segment;

	/**
	 * last acknowledged segment counter
	 */
	uint8_t lastACKedSegment = 0;

	/**
	 * lastData sent
	 */
	uint8_t *lastData;

	/**
	 * Simple protocol
	 */
	CovertProtocol<N - 1, PASSIVE> protocol;

	/**
	 * Sends a reset signal on next send call
	 */
	bool sendResetFlag = true;

	/**
	 * Evaluates received segment number and increases own segment number if necessary
	 *
	 * @param seg received segment number
	 * @return true if segment should discarded
	 */
	inline bool receiveSegment(uint8_t seg) {
		if constexpr (PASSIVE) {

			if (seg == segment && seg == increaseSegment(lastACKedSegment, 1)) {
				lastACKedSegment = seg;
				return false;
			} else {
				return true;
			}
		} else {

			if (seg == increaseSegment(segment, 1) && seg == increaseSegment(lastACKedSegment, 2)) {
				lastACKedSegment = increaseSegment(lastACKedSegment, 1);
				return false;
			} else {
				return true;
			}
		}
	}

	/**
	 * Returns true if the segment should be repeated.
	 *
	 * @return true if segment should be repeated
	 */
	inline bool sendSegment() {
		if (segment == lastACKedSegment) {
			segment = increaseSegment(segment, 1);
			return false;
		} else {
			return true;
		}
	}

	/**
	 * Calculates seg + inc
	 * @return seg + inc
	 */
	inline uint8_t increaseSegment(uint8_t seg, uint8_t inc) { return seg + inc & 0x3; }

	/**
	 * Resets the state of the protocol.
	 * This method can / should only be called from this class.
	 */
	void resetState() {
		protocol.reset();
		segment = (PASSIVE ? 1 : 0);
		lastACKedSegment = 0;
		std::memset(lastData, 0, N);
	}
};

#endif