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

#ifndef COVERTPROTOCOLBIDIRECTIONAL_H
#define COVERTPROTOCOLBIDIRECTIONAL_H

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

#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 {
	static_assert(N >= 2);

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

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

	/**
	 * 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
	 */
	bool sendFile(const std::string &fileName) {
		if constexpr (PASSIVE) {
			return false;
		}

		return protocol.sendFile(fileName);
	}

	/**
	 * send encodes the data into the data array
	 *
	 * @param data must be an array of size N
	 */
	void send(uint8_t *const data) {
		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;

		std::cerr << "segment: " << (int)segment << "\tack: " << (int)lastACKedSegment << "\tseg: " << (int)seg << std::endl;

		if (receiveSegment(seg)) {
			return;
		}

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

	/**
	 * 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(); }

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;

	/**
	 * 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 {
				std::cerr << "Throwing this packet away" << std::endl;
				return true;
			}
		} else {

			if (seg == increaseSegment(segment, 1) && seg == increaseSegment(lastACKedSegment, 2)) {
				lastACKedSegment = increaseSegment(lastACKedSegment, 1);
				return false;
			} else {
				std::cerr << "Throwing this packet away" << std::endl;
				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; }
};

#endif