#include <iostream>
#include <string>
#include <fstream>
#include <sstream>
#include <map>

#define uc unsigned char


using namespace std;

namespace Utils {

	/**
	 * Функция читает 6 байт из файла и возвращает строку с MAC адресом
	 * @param file - файл
	 * @return строка с MAC адресом
	 */
	std::string readMAC(std::fstream &file) {
		char MAC[6];
		file.read(MAC, 6);

		std::ostringstream ss;
		for (int i = 0; i < 6; ++i) {
			if (i != 0) ss << ':';
			ss.width(2); //< Два символа для каждого байта
			ss.fill('0'); //< Дополняем нулем, если не хватает нужного количества символов
			ss << std::hex << (MAC[i] & 0xFF);
		}

		// TODO: FIX IT! Пустой MAC
		if (ss.str() == "00:00:00:00:00:00") return Utils::readMAC(file);
		else return (ss.str());
	}

	/**
	 * Функция читает 2 байта из файла и определяет тип кадра
	 * 0x0800 - IPv4, 0x8060 - ARP
	 * @param file - файл
	 * @return строка - тип
	 */
	std::tuple<std::string, std::string, int> getFrameType(std::fstream &file) {
		char bytes[2];
		file.read(bytes, 2);

		int packetSize = (unsigned char)bytes[0] << 8 | (unsigned char)bytes[1] << 0;

		std::string protocol;
		if (bytes[0] == 8 && bytes[1] == 0) protocol = "IPv4";
		else if (bytes[0] == 8 && bytes[1] == 6) protocol = "ARP";
		else protocol = "Unknown";

		std::string frameType;
		if (packetSize <= 0x05DC) {
			file.read(bytes, 2); // Cчитываем дополнительно два байта
			file.seekg(-2, std::ios_base::cur);
			if ((uc)bytes[0] == 0xFF && (uc)bytes[1] == 0xFF) frameType = "Raw 802.3";
			else if ((uc)bytes[0] == 0xAA && (uc)bytes[1] == 0xAA) frameType = "Ethernet SNAP";
			else    frameType = "802.3/LLC";
		}
		else       frameType = "Ethernet DIX";

		return std::make_tuple(frameType, protocol, packetSize);
	}

	std::string readIP(std::fstream &file) {
		char ip[4];
		file.read(ip, 4);

		std::ostringstream ss;
		for (int i = 0; i < 4; ++i) {
			if (i != 0) ss << '.';
			ss << std::dec << (ip[i] & 0xFF);
		}

		return ss.str();
	}

	/**
	 * Функция читает два байта и возвращает число
	 * @param file - файл
	 * @return целове число
	 */
	int readShort(std::fstream &file) {
		char bytes[2];
		file.read(bytes, 2);

		// return (unsigned char)bytes[1] << 0 | (unsigned char)bytes[0] << 8;
		return (unsigned char)bytes[0] << 8 | (unsigned char)bytes[1] << 0;
	}

	/**
	 * Функция читает байт и определяет по нему тип пакета
	 * 0x01 - ICMP, 0x06 - TCP, 0x17 - UDP
	 * @return строка - тип пакета
	 */
	std::string getPacketType(std::fstream &file) {
		char byte;
		file.read(&byte, 1);

		switch (byte) {
		case 0x01: return std::string("ICMP");
		case 0x06: return std::string("TCP");
		case 0x11: return std::string("UDP");
		default: return std::string("Unknown");
		}
	}

	void printMap(std::map<std::string, int> &map) {
		for (auto &elem : map) {
			std::cout << elem.first << ": " << elem.second << std::endl;
		}
	}

}

namespace Packets {

	/**
	 * Функция парсит IPv4 пакет и выводит его данные
	 * @link https://en.wikipedia.org/wiki/IPv4#Header
	 * @param file - файл
	 */
	void IPv4(std::fstream &file) {
		file.ignore(2); // Пропускаем 2 байта (Version, IHL, DSCP, ECN)

		int packetSize = Utils::readShort(file); // 2 байта - Размер IPv4 пакета
		cout << "IPv4 packet size: " << packetSize << endl;

		file.ignore(5); // Пропускаем 5 байт (Identification, Flags, Fragment Offset, TTL)

		cout << "Packet type: " << Utils::getPacketType(file) << endl; // 1 байт - Протокол

		file.ignore(2); // Пропускаем 2 байта (Header Checksum)

		cout << "Source IP4     : " << Utils::readIP(file) << endl; // 4 байта
		cout << "Destination IP4: " << Utils::readIP(file) << endl; // 4 байта

		file.ignore(packetSize - 20); // Пропускаем данные пакета (размер пакета - 20 байт на Header)
	}

	/**
	 * Размер ARP пакета меньше 46 байт
	 * Но т.к. минимальное значение пакета должно быть 46 байт
	 * То пропускаем 46 байт!
	 */
	void ARP(std::fstream &file) {
		file.ignore(46);
	}

}

using namespace std;

int main() {
	string filePath = "F:/files/ethers07.bin";
	std::fstream file(filePath, std::ios::in | std::ios::binary);

	// Счетчики
	std::map<std::string, int> frames;
	std::map<std::string, int> protocols;

	int frameNumber = 1;
	while (file.good() && !file.eof()) {
		// cout << "File pointer:" << file.tellg() << endl;
		cout << "Frame number: " << frameNumber++ << endl;
		// Читаем два MAC адреса из файла
		cout << "Destination MAC: " << Utils::readMAC(file) << endl;
		cout << "Source      MAC: " << Utils::readMAC(file) << endl;

		string frameType;
		string protocol;
		int packetSize;
		tie(frameType, protocol, packetSize) = Utils::getFrameType(file);

		if (frames.count(frameType) == 0) frames.insert({ frameType, 1 });
		else frames.at(frameType)++;

		if (protocols.count(protocol) == 0) protocols.insert({ protocol, 1 });
		else protocols.at(protocol)++;


		cout << "Frame type: " << frameType << endl;

		if (protocol == "IPv4") {
			Packets::IPv4(file);
		}
		else if (protocol == "ARP") {
			Packets::ARP(file);
		}
		else if (protocol == "Unknown") {
			cout << "Packet size: " << packetSize << endl;
			file.seekg(packetSize, std::ios_base::cur);
		}
		// TODO: Legacy code: eof ставится только после попытки чтения
		file.ignore(1);
		if (file.eof()) {
			cout << endl << "Frame types:" << endl;
			Utils::printMap(frames);
			cout << endl << "Protocol:" << endl;
			Utils::printMap(protocols);
			system("pause");
			return 0;
		}
		file.seekg(-1, std::ios_base::cur);

		cout << endl << endl;
	}
	return 0;
}