#ifndef SOCKET_H
#define SOCKET_H

// Headerfiles
/* Headerfiles */
#include "blueTypes.h"
#include <iostream>
#include <fstream>
#include <string>
#include <errno.h>
#include <conio.h>
/* MS Windows */
#ifdef _WIN32
#include <winsock2.h>
#include <windows.h>
#include <mstcpip.h>
#include <ws2tcpip.h>
#else
/* Unix/Linux */
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>
#endif
using std::cout;
using std::cin;
using std::endl;
using std::ofstream;
using std::ios;

// Max sockets
#ifndef FD_SETSIZE
#define FD_SETSIZE 64
#endif

#ifndef SOCKET_ERROR
#define SOCKET_ERROR -1
#endif

#ifdef _WIN32
#define socket_t SOCKET
#else
#define socket_t int
#endif

class sSocket {
	static const int BufferLen		= 2048;
	static const int HostnameLen	= 255;

	// change these to specify your socket
	static const int sDomain		= AF_INET; // IPv4
	static const int sType			= SOCK_RAW;
	static const int sProtocol		= IPPROTO_IP;

	// protocol types
	static const short tcp			= 6;
	static const short udp			= 17;

#ifdef _WIN32
	WSAData wsd;
#endif
protected:
	
private:
#ifdef _WIN32
	bool sInitWinsock(void);
#endif

	bool sInitSniffer(sClient &_client);
	socket_t sInitSocket(void);
	void sError(char* message);
	void sErrorExit(char* message);
	void sLog(char* file, char* text);
	void sHandlePacket(char* _buffer, int packetSize);
	void sClose(socket_t _sock);
	bool sGetDomainname(char* ip, int ipLen, char* domain, int domainLen);
public:
	// Destruktor/Konstruktor
	sSocket();
	~sSocket();

	bool sInitClient(sClient &_client, int _port, char* _host);
	bool sBind(sClient &_client);
	int sRecvFrom(sClient &_client, char* _buffer, size_t _bufferLen, sockaddr_in &_addr);

	void sStartSniffing(sClient &_client);
};

	sSocket::sSocket() {
#ifdef _WIN32
	if(!sInitWinsock())
		sErrorExit("WSAStartup() failed.");
#endif
}

	sSocket::~sSocket() {
#ifdef _WIN32
	WSACleanup();
#endif
}



bool 
	sSocket::sInitClient(sClient &_client, int _port, char* _host) {
	if((int)strlen(_host) <= 1) {
		sError("No host specified\n");
		return false;
	}
	
	_client.port = _port;
	if((_client.sock = sInitSocket()) == 0)
		return false;
	memset(&_client.SockAddr, 0, sizeof(_client.SockAddr));

	/*
	if(this->socket_maxSocket < _client.sock)
		this->socket_maxSocket = _client.sock;

	if(this->server) {
		FD_SET(_client.sock, &this->server->socketSet);
	}
	*/

	// build the addr
	char hostname[HostnameLen];
	memset(hostname, 0, HostnameLen);
	if(!strcmp(_host, "localhost"))
		gethostname(hostname, HostnameLen);
	else 
		strcpy_s(hostname, HostnameLen, _host);

	_client.SockAddr.sin_family = sDomain;
	_client.SockAddr.sin_port	= htons(_port);

	unsigned long addr;

	_client.pHost = (struct hostent*)new char[(sizeof(struct hostent))];
	if(_client.pHost == NULL) {
		sErrorExit("Not enough free space.");
		return false;
	}
	if((addr = inet_addr(hostname)) != INADDR_NONE) {
		// alternativly you can use gethostbyaddr()
		memcpy((char*)&_client.SockAddr.sin_addr, &addr, sizeof(addr));
	}
	else {
		_client.pHost = gethostbyname(hostname);
		memcpy((char*)&_client.SockAddr.sin_addr, _client.pHost->h_addr, _client.pHost->h_length);
	}
	
	return true;
}

bool 
	sSocket::sInitSniffer(sClient &_client) {
	DWORD BytesReturned;
	int I = 1;
	if (WSAIoctl(_client.sock, SIO_RCVALL, &I, sizeof(I), NULL, NULL, &BytesReturned, NULL, NULL) == SOCKET_ERROR) {
		sErrorExit("WSAIoctl() failed.");
		return false;
	}

	return true;
}
void 
	sSocket::sStartSniffing(sClient &_client) {
	if(!sInitSniffer(_client))
		return;
	
	char recvBuffer[BufferLen];
	sockaddr_in addrFrom;
	int recvLen;

	while(!_kbhit()) {
		memset(recvBuffer, 0, BufferLen);
		memset(&addrFrom, 0, sizeof(sockaddr_in));

		recvLen = sRecvFrom(_client, recvBuffer, BufferLen, addrFrom);
		// handle packet
		sHandlePacket(recvBuffer, recvLen);
	}

	sClose(_client.sock);
}

int 
	sSocket::sRecvFrom(sClient &_client, char* _buffer, size_t _bufferLen, sockaddr_in &_addr) {
	if(_buffer == NULL)
		sErrorExit("RecvFrom() failed. No buffer.");

	memset(_buffer, 0, _bufferLen);
	
	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(struct sockaddr_in));
	int recvLen = 0, addrLen = sizeof(addr);
	
	/*	recvfrom() works like recv(). The differences are, that it will write some data from the sender into 
		the sockaddr_in. recv() returns zero when the connection is lost. recvfrom() will return zero when the
		incoming packet has zero bytes.
	*/
	recvLen = recvfrom(_client.sock, _buffer, (int)_bufferLen, NULL, (sockaddr *)&_addr, &addrLen);

	if(recvLen == 0) {
		sError("Packet of zero bytes received.");
	}

	return recvLen;
		
}
bool 
	sSocket::sBind(sClient &_client) {
	if(bind(_client.sock, (struct sockaddr*)&_client.SockAddr, sizeof(_client.SockAddr)) == -1) {
		/* Fehler! */
		sClose(_client.sock);
		sErrorExit("Could not bind port.\n");
		return false;
	}
	else
		return true;
}

socket_t 
	sSocket::sInitSocket() {
	socket_t sock;
	const int y = 1;

	sock = socket(
			sDomain, 			/* Adressfamilie */ 
			sType, 			/* Socket Typ: SOCK_STREAM = TCP; SOCK_DGRAM = UDP */
			sProtocol);		/* Cbertragungsprotokoll: SOCK_STREAM = TCP; SOCK_DGRAM = UDP */
	
	if(sock < 0) {
		/* Ein Fehler ist aufgetreten. */
		sErrorExit("Fehler im initialisieren des Sockets");
		return 0;
	}
	
	/* Socket so einstellen, dass mehrere Prozesse einen
	 * Port nutzen koennen. (Dadurch umgeht man den Timeout
	 * von drei Minuten) */
	setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (char*)&y, sizeof(int));
	
	/* return socket */
	return sock;	

}
bool
	sSocket::sGetDomainname(char* ip, int ipLen, char* domain, int domainLen) {
	
	return true;
}
void
	sSocket::sHandlePacket(char* _buffer, int packetSize) {

	char logBuffer[BufferLen];
	memset(logBuffer, 0, BufferLen);
	char castBuf[256];
	memset(castBuf, 0, 255);
	sIpHeader* ipHdr = (sIpHeader*) _buffer;
	struct hostent* curHost;
	unsigned int curAddr;
	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(struct sockaddr_in));

	// Src/Dest IPs and domains
	strcat_s(logBuffer, BufferLen, "src: ");
	// Source:
	addr.sin_addr.s_addr = ipHdr->ipSrcAddr;
	//std::cout << "\nNew packet:\nFrom: " << inet_ntoa(addr.sin_addr) << std::endl; // This prints the IP
	// Get the full domain
	curAddr = inet_addr(inet_ntoa(addr.sin_addr));
	curHost = gethostbyaddr((char*)&curAddr, 4, AF_INET);
	if(curHost && (strlen(curHost->h_name) >1)) {
		std::cout << "From: " << curHost->h_name << std::endl; // Print full domain
		strcat_s(logBuffer, BufferLen, curHost->h_name);
	}
	else { // Log IP
		strcat_s(logBuffer, BufferLen, inet_ntoa(addr.sin_addr));
	}

	// Destination:
	strcat_s(logBuffer, BufferLen, " --- dst: ");
	addr.sin_addr.s_addr = ipHdr->ipDestAddr;
	//std::cout << "To: " << inet_ntoa(addr.sin_addr) << std::endl; // This prints the IP
	// Get the full domain
	curAddr = inet_addr(inet_ntoa(addr.sin_addr));
	curHost = gethostbyaddr((char*)&curAddr, 4, AF_INET);	
	if(curHost && (strlen(curHost->h_name) > 1)) { // Sometimes a host does not have a domainname (i.e.: backbone)
		std::cout << "To: " << curHost->h_name << std::endl; // Print full domain
		strcat_s(logBuffer, BufferLen, curHost->h_name);
	}
	else { // Log IP
		strcat_s(logBuffer, BufferLen, inet_ntoa(addr.sin_addr));
	}


	// Get IP header length (IHL)
	unsigned short ipHdrLen = (ipHdr->ipVerIhl << 4);
	char C;
	memcpy(&C, &ipHdrLen, 1);
	ipHdrLen = (C >> 4)*4;

	// Get the port
	if(packetSize > ipHdrLen) {
		sTcpUdpPorts *ports;	
		ports = (sTcpUdpPorts*)(_buffer + ipHdrLen);
		std::cout << "Source port: " << ntohs(ports->SrcPort) << std::endl;
		strcat_s(logBuffer, BufferLen, " :: src: ");
		// Cast u_short to char*
		
		sprintf_s(castBuf, 255, "%d", ntohs(ports->SrcPort));
		strcat_s(logBuffer, BufferLen, castBuf);
		std::cout << "Destination port: " << ntohs(ports->DstPort) << std::endl;
		strcat_s(logBuffer, BufferLen, " --- dst: ");
		sprintf_s(castBuf, 255, "%d", ntohs(ports->SrcPort));
		strcat_s(logBuffer, BufferLen, castBuf);
	}

	// Print protocol
	switch(ipHdr->ipProtocol) {
		case tcp:
			strcat_s(logBuffer, BufferLen, " :: TCP");
			std::cout << "Used protocol: TCP" << std::endl;
			break;
		case udp:
			strcat_s(logBuffer, BufferLen, " :: UDP");
			std::cout << "Used protocol: UDP" << std::endl;
			break;
		default: 
			strcat_s(logBuffer, BufferLen, " :: Unknown");
			std::cout << "Used protocol: " << ipHdr->ipProtocol << std::endl;
			break;
	}
	strcat_s(logBuffer, BufferLen, "\r\n");
	sLog("traffic.log", logBuffer);
}
void 
	sSocket::sClose(socket_t _sock) {
#ifdef _WIN32
	closesocket(_sock);
#else
	close(_sock);
#endif
}
void 
	sSocket::sError(char *message) {
	sLog("Errors.log", message);
#ifdef _WIN32
	std::cout << message << "\n" << WSAGetLastError() << std::endl;
#else
	std::cout << message << "\n" << strerror(errno) << std::endl; /* Buffer Overflow */
#endif
}
void 
	sSocket::sErrorExit(char *message) {
	sLog("Errors.log", message);
#ifdef _WIN32
	std::cout << message << "\n" << WSAGetLastError() << std::endl;
#else
	std::cout << message << "\n" << strerror(errno) << std::endl; /* Buffer Overflow */
#endif
	exit(EXIT_FAILURE);
}
void
	sSocket::sLog(char* file, char* text) {
	if((strlen(file) < 2) || (strlen(text) < 2))
		return;
	try {
		ofstream logFile;
		logFile.open(file, ios::out | ios::app | ios::binary);
		if(logFile.is_open()) {
			logFile << text;
			logFile.close();
		}
	}
	catch(...) {
		sError("sLog() failed.");
	}
}
#ifdef _WIN32
bool 
	sSocket::sInitWinsock() {
	if(WSAStartup(MAKEWORD(2,2), &wsd)!= NO_ERROR)
		return false;
	else
		return true;	
}
#endif

#endif