/***************************************************************************
 *   Copyright (C) 2005 by Joris Guisson                                   *
 *   joris.guisson@gmail.com                                               *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.             *
 ***************************************************************************/
#include "encryptedpacketsocket.h"
#include <errno.h>
#include <util/sha1hash.h>
#include <util/log.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <QtGlobal>
#ifndef Q_WS_WIN
#include <netinet/in.h>
#include <netinet/in_systm.h>
#include <netinet/ip.h> 
#else
#define IPTOS_THROUGHPUT 0x08
#endif
#include <netinet/tcp.h>
#include <net/socketmonitor.h>
#include <util/functions.h>
#include "rc4encryptor.h"

using namespace bt;
using namespace net;

namespace mse
{
	
	Uint8 EncryptedPacketSocket::tos = IPTOS_THROUGHPUT;

	EncryptedPacketSocket::EncryptedPacketSocket(int ip_version) : 
		net::PacketSocket(true, ip_version),
		enc(0),
		monitored(false)
	{
		sock->setBlocking(false);
		sock->setTOS(tos);
		reinserted_data = 0;
		reinserted_data_size = 0;
		reinserted_data_read = 0;
	}

	EncryptedPacketSocket::EncryptedPacketSocket(int fd,int ip_version) : 
		net::PacketSocket(fd, ip_version),
		enc(0),
		monitored(false)
	{
		sock->setBlocking(false);
		sock->setTOS(tos);
		reinserted_data = 0;
		reinserted_data_size = 0;
		reinserted_data_read = 0;
	}

	EncryptedPacketSocket::EncryptedPacketSocket(net::SocketDevice* sd) : 
		net::PacketSocket(sd),
		enc(0),
		monitored(false)
	{
		sd->setBlocking(false);
		sd->setTOS(tos);
		reinserted_data = 0;
		reinserted_data_size = 0;
		reinserted_data_read = 0;
	}

	EncryptedPacketSocket::~EncryptedPacketSocket()
	{
		if (monitored)
			stopMonitoring();
		
		delete [] reinserted_data;
		delete enc;
	}
	
	void EncryptedPacketSocket::startMonitoring(net::SocketReader* rdr)
	{
		setReader(rdr);
		SocketMonitor::instance().add(this);
		monitored = true;
		if (reinserted_data)
		{
			if (enc)
				enc->decrypt(reinserted_data + reinserted_data_read,
							 reinserted_data_size - reinserted_data_read);
			
			rdr->onDataReady(reinserted_data + reinserted_data_read,
							 reinserted_data_size - reinserted_data_read);
			delete [] reinserted_data;
			reinserted_data = 0;
			reinserted_data_size = 0;
		}
	}
	
	void EncryptedPacketSocket::stopMonitoring()
	{
		SocketMonitor::instance().remove(this);
		monitored = false;
		rdr = 0;
	}
		
	Uint32 EncryptedPacketSocket::sendData(const Uint8* data,Uint32 len)
	{
		if (enc)
		{
			// we need to make sure all data is sent because of the encryption
			Uint32 ds = 0;
			const Uint8* ed = enc->encrypt(data,len);
			while (sock->ok() && ds < len)
			{
				Uint32 ret = sock->send(ed + ds,len - ds);
				ds += ret;
				if (ret == 0)
				{
					Out(SYS_CON|LOG_DEBUG) << "ret = 0" << endl;
				}
			}
			if (ds != len)
				Out(SYS_CON|LOG_DEBUG) << "ds != len" << endl;
			return ds;
		}
		else
		{
			Uint32 ret = sock->send(data,len);
			if (ret != len)
				Out(SYS_CON|LOG_DEBUG) << "ret != len" << endl;
			return ret;
		}
	}
		
	Uint32 EncryptedPacketSocket::readData(Uint8* buf,Uint32 len)
	{
		Uint32 ret2 = 0;
		if (reinserted_data)
		{
			Uint32 tr = reinserted_data_size - reinserted_data_read;
			if (tr < len)
			{
				memcpy(buf,reinserted_data + reinserted_data_read,tr);
				delete [] reinserted_data;
				reinserted_data = 0;
				reinserted_data_size = reinserted_data_read = 0;
				ret2 = tr;
				if (enc)
					enc->decrypt(buf,tr);
			}
			else
			{
				tr = len;
				memcpy(buf,reinserted_data + reinserted_data_read,tr);
				reinserted_data_read += tr;
				if (enc)
					enc->decrypt(buf,tr);
				return tr;
			}
		}
		
		if (len == ret2)
			return ret2;
		
		Uint32 ret = sock->recv(buf + ret2,len - ret2);
		if (ret + ret2 > 0 && enc)
			enc->decrypt(buf,ret + ret2);
		
		return ret;
	}
		
	Uint32 EncryptedPacketSocket::bytesAvailable() const
	{
		Uint32 ba = sock->bytesAvailable();
		if (reinserted_data_size - reinserted_data_read > 0)
			return  ba + (reinserted_data_size - reinserted_data_read);
		else
			return ba;
	}
	
	void EncryptedPacketSocket::close()
	{
		sock->close();
	}
	
	bool EncryptedPacketSocket::connectTo(const QString & ip,Uint16 port)
	{
		// do a safety check
		if (ip.isNull() || ip.length() == 0)
			return false;
		
		return connectTo(net::Address(ip,port));
	}
	
	bool EncryptedPacketSocket::connectTo(const net::Address & addr)
	{
		// we don't wanna block the current thread so set non blocking
		sock->setBlocking(false);
		sock->setTOS(tos);
		if (sock->connectTo(addr))
			return true;
		
		return false;
	}
			
	void EncryptedPacketSocket::initCrypt(const bt::SHA1Hash & dkey,const bt::SHA1Hash & ekey)
	{
		delete enc;
		enc = new RC4Encryptor(dkey,ekey);
	}
		
	void EncryptedPacketSocket::disableCrypt()
	{
		delete enc;
		enc = 0;
	}
	
	bool EncryptedPacketSocket::ok() const
	{
		return sock->ok();
	}

	QString EncryptedPacketSocket::getRemoteIPAddress() const
	{
		return sock->getPeerName().toString();
	}
	
	bt::Uint16 EncryptedPacketSocket::getRemotePort() const
	{
		return sock->getPeerName().port();
	}
	
	net::Address EncryptedPacketSocket::getRemoteAddress() const
	{
		return sock->getPeerName();
	}
	
	void EncryptedPacketSocket::setRC4Encryptor(RC4Encryptor* e)
	{
		delete enc;
		enc = e;
	}
	
	void EncryptedPacketSocket::reinsert(const Uint8* d,Uint32 size)
	{
//		Out() << "Reinsert : " << size << endl;
		Uint32 off = 0;
		if (reinserted_data)
		{
			off = reinserted_data_size;
			reinserted_data = (Uint8*)realloc(reinserted_data,reinserted_data_size + size);
			reinserted_data_size += size;
		}
		else
		{
			reinserted_data = new Uint8[size];
			reinserted_data_size = size;
		}
		memcpy(reinserted_data + off,d,size);
	}
	
	bool EncryptedPacketSocket::connecting() const
	{
		return sock->state() == net::SocketDevice::CONNECTING;
	}

	bool EncryptedPacketSocket::connectSuccesFull() const 
	{
		return sock->connectSuccesFull();
	}
	
	void EncryptedPacketSocket::setRemoteAddress(const net::Address & addr)
	{
		sock->setRemoteAddress(addr);
	}
	
	void EncryptedPacketSocket::preProcess(Packet::Ptr packet)
	{
		if (enc)
			enc->encryptReplace(packet->getData(), packet->getDataLength());
	}
	
	void EncryptedPacketSocket::postProcess(Uint8* data, Uint32 size)
	{
		if (enc)
			enc->decrypt(data, size);
	}
	
}
