#include "SocketC.h"

//初始化socket
bool SocketC::initSocket()
{
	//初始化套接字库
	WORD w_req = MAKEWORD(2, 2);//版本号
	WSADATA wsadata;
	int err;
	err = WSAStartup(w_req, &wsadata);
	if (err != 0)
	{
		return false;
	}
	//检测版本号
	if (LOBYTE(wsadata.wVersion) != 2 || HIBYTE(wsadata.wHighVersion) != 2)
	{
		WSACleanup();
		return false;
	}
	else
	{
		return true;
	}
}

//构造
SocketC::SocketC(int mode)
{
	m_mode = mode;
}

//析构
SocketC::~SocketC()
{
	if (state == 1)
	{
		Close();
	}
}

//连接到服务器
bool SocketC::Connect(string IP, UINT Prot)
{
	if (!initSocket())
	{//初始化失败
		return false;
	}
	struct addrinfo* answer, hint;
	char ipstr[16] = { 0 };
	ZeroMemory(&hint, sizeof(hint));
	hint.ai_family = AF_INET;
	hint.ai_socktype = SOCK_STREAM;
	SOCKADDR_IN server_addr;
	//解析域名（如果传入的IP是域名的情况）
	if (getaddrinfo(IP.c_str(), NULL, &hint, &answer) == 0)
	{
		server_addr.sin_family = AF_INET;
		server_addr.sin_addr.S_un.S_addr = ((struct sockaddr_in*)(answer->ai_addr))->sin_addr.S_un.S_addr;
		server_addr.sin_port = htons(Prot);
		((SOCKADDR_IN*)&hint.ai_addr)->sin_port = htons(Prot);

		//创建套接字
		sclient = socket(AF_INET, SOCK_STREAM, 0);
		//连接到服务器
		if (connect(sclient, (SOCKADDR*)&server_addr, sizeof(SOCKADDR)) == SOCKET_ERROR)
		{//连接失败
			WSACleanup();
			return false;
		}
		else
		{//连接成功创建线程接收数据
			state = 1;
			rec = std::thread((&SocketC::Receive), this);
			return true;
		}
	}
	return FALSE;
}

//接收数据函数
VOID SocketC::Receive()
{
	char* buff = NULL;
	while(state)
	{
		//开始接收数据
		char buf[1024] = { 0 };
		int len = recv(sclient, buf, 1020, 0);
		if (len ==SOCKET_ERROR || len == 0)
		{//服务器关闭了
			state = 0;
			if (Cfunc != NULL)
			{
				Cfunc();
			}
			//cout << "服务器断开连接";
			break;
		}
		if (m_mode == 0)
		{//普通模式直接处理
			if (Rfunc != NULL)
				Rfunc(buf, len);
		}
		else if (m_mode == 1)
		{//处理收到的数据避免出现粘包的情况
			int isStr;
			buff = RecvBuff(buf, len, isStr);
			while (buff)
			{
				if (Rfunc != NULL)
				{
					if (isStr == 1)
					{//如果是字符串先解密
						AES aes;
						string str = aes.DecryptionAES(buff, "qwertyuiopasdfgh");
						Rfunc((char*)str.c_str(), str.length());
					}
					else
						Rfunc(buff, len);
				}
				len = 0;
				ZeroMemory(buff, 10001);
				buff = RecvBuff(NULL, len, isStr);
			}
		}
	}
	return;
}

//处理收到的数据
char* SocketC::RecvBuff(char* buf, int& len, int& isStr)
{
	MsgHead* h;
	//lock_guard<std::mutex> guard(msgLock);
	//把收到的数据复制到缓冲区
	if (len != 0)
		memcpy_s(tmpBuf + bufLen, len, buf, len);
	bufLen += len;
	//判断缓冲区内的消息长度是否大于消息头
	if (bufLen > sizeof(MsgHead))
	{//检查消息头是由有问题
		h = (MsgHead*)tmpBuf;
		if (!CheckHead(h))
		{//消息头不合规关闭连接
			closesocket(sclient);
			return NULL;
		}
		//判断当前消息是否收完
		if (bufLen - sizeof(MsgHead) >= h->bufLen)
		{//当前消息收完了
			memcpy_s(buff, h->bufLen, tmpBuf + sizeof(MsgHead), h->bufLen);
			len = h->bufLen;
			isStr = h->isStr;
			bufLen = bufLen - len - sizeof(MsgHead);
			//缓冲区数据前移
			if (bufLen > 0)
				memcpy_s(tmpBuf, bufLen, tmpBuf + sizeof(MsgHead) + len, bufLen);
			return buff;
		}
	}
	return NULL;
}

//检查消息头
BOOL SocketC::CheckHead(MsgHead* h)
{
	//消息长度不能大于10000
	if (h->bufLen > 10000)
	{
		return FALSE;
	}
	//机密数据并验证发过来的Token是否正确
	string str, str1;
	str = to_string(h->tm);
	str += to_string(h->bufLen);
	md5 md;
	str1 = md.StringToMD5(str);
	str = h->token;
	if (str == str1)
	{//token正确验证发送消息的时间
		time_t t1;
		time(&t1);
		if (abs(t1 - h->tm) < 300)
		{
			return TRUE;
		}
	}
	return FALSE;
}

//发送数据包
int SocketC::SendData(const char* Data, int len, int isStr)
{
	if (Data == NULL)
		return -1;
	if (m_mode == 0)
	{//普通模式直接发送
		return send(sclient, Data, len, 0);
	}
	else if (m_mode == 1)
	{//生成token
		MsgHead h;
		h.isStr = isStr;
		h.bufLen = len;
		time(&h.tm);
		string str = to_string(h.tm);
		str += to_string(len);
		md5 md;
		str = md.StringToMD5(str);
		memcpy_s(h.token, str.length(), str.c_str(), str.length());
		char* buff = new char[sizeof(MsgHead) + len];
		memcpy_s(buff, sizeof(MsgHead), &h, sizeof(MsgHead));
		memcpy_s(buff + sizeof(MsgHead), len, Data, len);
		//发送数据
		int lenth = send(sclient, buff, len + sizeof(MsgHead), 0) - sizeof(MsgHead);
		delete[] buff;
		if (lenth < 0)
		{
			lenth = -1;
		}
		return lenth;
	}
}

//发送字符串
BOOL SocketC::SendStr(string str)
{
	if (str == "")
		return FALSE;
	if (m_mode == 0)
	{//普通模式直接发送
		if (send(sclient, str.c_str(), str.length(), 0) == -1)
		{
			return FALSE;
		}
	}
	else if (m_mode == 1)
	{//加密发送的字符串并发送
		AES aes;
		string sendStr = aes.EncryptionAES(str, "qwertyuiopasdfgh");
		if (-1 == SendData(sendStr.c_str(), sendStr.length(), 1))
		{
			return FALSE;
		}
	}
	return TRUE;
}

//设置接收数据的回调函数
VOID SocketC::SetRecvFunc(function<VOID(char*, int)> fun)
{
	Rfunc = fun;
}

//设置断开连接的回调函数
VOID SocketC::SetCloseFunc(function<VOID()> fun)
{
	Cfunc = fun;
}

//断开连接
void SocketC::Close()
{
	if (state != 0)
	{
		state = 0;
		//关闭套接字
		shutdown(sclient, SD_BOTH);
		closesocket(sclient);
		//释放DLL资源
		rec.join();
	}
	WSACleanup();
}