#include <string>

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <string.h>
#include <unistd.h>

#include <gnutls/gnutls.h>

#include "users.h"
#include "channels.h"
#include "modules.h"
#include "helperfuncs.h"
#include "socket.h"
#include "hashcomp.h"

/* $ModDesc: Provides SSL support for clients */
/* $CompileFlags: `libgnutls-config --cflags` */
/* $LinkerFlags: -L/usr/local/lib -Wl,--rpath -Wl,/usr/local/lib -L/usr/lib -Wl,--rpath -Wl,/usr/lib -lgnutls */

#define KEYFILE "key.pem"
#define CERTFILE "cert.pem"
#define CAFILE "ca.pem"
#define CRLFILE "crl.pem"

#define DH_BITS 1024

enum issl_status { ISSL_NONE, ISSL_HANDSHAKING_READ, ISSL_HANDSHAKING_WRITE, ISSL_HANDSHAKEN, ISSL_CLOSING, ISSL_CLOSED };

class issl_session
{
public:
	gnutls_session_t sess;
	issl_status status;
	std::string outbuf;
	int fd;
};

class ModuleSSL : public Module
{
	Server* Srv;
	ServerConfig* SrvConf;
	
	issl_session sessions[MAX_DESCRIPTORS];
	
	gnutls_certificate_credentials x509_cred;
	gnutls_dh_params dh_params;
	
 public:
	
	ModuleSSL(Server* Me)
		: Module::Module(Me)
	{
		Srv = Me;
		SrvConf = Srv->GetConfig();
		SrvConf->AddIOHook(6661, this);
		
		gnutls_global_init(); // This must be called once in the program

		gnutls_certificate_allocate_credentials(&x509_cred);
		gnutls_certificate_set_x509_trust_file(x509_cred, CAFILE, GNUTLS_X509_FMT_PEM);
		gnutls_certificate_set_x509_crl_file (x509_cred, CRLFILE, GNUTLS_X509_FMT_PEM);
		gnutls_certificate_set_x509_key_file (x509_cred, CERTFILE, KEYFILE, GNUTLS_X509_FMT_PEM);

		generate_dh_params();
		
		gnutls_certificate_set_dh_params(x509_cred, dh_params);
	}
	
	virtual ~ModuleSSL()
	{
		gnutls_certificate_free_credentials(x509_cred);
		gnutls_global_deinit();
	}
	
	virtual Version GetVersion()
	{
		return Version(1, 0, 0, 0, 0);
	}

	void Implements(char* List)
	{
		List[I_OnRawSocketAccept] = List[I_OnRawSocketClose] = List[I_OnRawSocketRead] = List[I_OnRawSocketWrite] = 1;
		List[I_OnSyncUserMetaData] = List[I_OnDecodeMetaData] = List[I_OnUserQuit] = 1;
	}

	virtual void OnRawSocketAccept(int fd, std::string ip, int localport)
	{
		issl_session* session = &sessions[fd];
	
		session->fd = fd;	
	
		gnutls_init(&session->sess, GNUTLS_SERVER);

		gnutls_set_default_priority(session->sess); // Avoid calling all the priority functions, defaults are adequate.
		gnutls_credentials_set(session->sess, GNUTLS_CRD_CERTIFICATE, x509_cred);
		gnutls_certificate_server_set_request(session->sess, GNUTLS_CERT_REQUEST); // Request client certificate if any.
		gnutls_dh_set_prime_bits(session->sess, DH_BITS);
		gnutls_transport_set_ptr(session->sess, (gnutls_transport_ptr_t) fd); // Give gnutls the fd for the socket.
		
		Handshake(session);
	}

	virtual void OnRawSocketClose(int fd)
	{
		CloseSession(&sessions[fd]);
	}
	
	virtual int OnRawSocketRead(int fd, char* buffer, unsigned int count, int &readresult)
	{
		issl_session* session = &sessions[fd];
		
		if(session->status == ISSL_HANDSHAKING_WRITE || session->status == ISSL_HANDSHAKING_READ)
		{
			// The handshake isn't finished, try to finish it.
			
			if((session->status == ISSL_HANDSHAKING_READ) && Handshake(session))
			{
				// Handshake successfully resumed.
				log(DEBUG, "m_ssl.so: OnRawSocketRead(): successfully resumed handshake");
			}
			else
			{
				// If the handshake still isn't finished, don't try and write anything.
				
				if(session->status == ISSL_HANDSHAKING_WRITE)
					log(DEBUG, "m_ssl.so: OnRawSocketRead(): handshake wants to write data but we are currently reading");
				else
					log(DEBUG, "m_ssl.so: OnRawSocketRead(): failed to resume handshake"); 
				
				// If the handshake isn't complete I don't think we can read with gnutls_record_recv(),
				// so I don't think we can do anything.
				
				return 1;
			}
		}
		
		// Is this right? Not sure if the unencrypted data is garaunteed to be the same length.
		int ret = gnutls_record_recv(session->sess, buffer, count);

		if(ret == 0)
		{
			// Client closed connection.
			log(DEBUG, "m_ssl: Client closed the connection");
		}
		else if(ret < 0)
		{
			if(ret != GNUTLS_E_AGAIN && ret != GNUTLS_E_INTERRUPTED)
			{
				log(DEBUG, "m_ssl.so: Error reading SSL data: %s", gnutls_strerror(ret));
			}
			else
			{
				log(DEBUG, "m_ssl.so: Not all SSL data read: %s", gnutls_strerror(ret));
			}
	    }
		else if (ret > 0)
		{
			readresult = ret;
		}
		
		return 1;
	}
	
	virtual int OnRawSocketWrite(int fd, char* buffer, int count)
	{		
		issl_session* session = &sessions[fd];
		const char* sendbuffer = buffer;
		
		if(session->status == ISSL_HANDSHAKING_WRITE || session->status == ISSL_HANDSHAKING_READ)
		{
			// The handshake isn't finished, try to finish it.
			
			if((session->status == ISSL_HANDSHAKING_WRITE) && Handshake(session))
			{
				// Handshake successfully resumed.
				log(DEBUG, "m_ssl.so: OnRawSocketWrite(): successfully resumed handshake");
			}
			else
			{
				// If the handshake still isn't finished, don't try and write anything.
				
				if(session->status == ISSL_HANDSHAKING_READ)
					log(DEBUG, "m_ssl.so: OnRawSocketWrite(): handshake wants to read data but we are currently writing");
				else
					log(DEBUG, "m_ssl.so: OnRawSocketWrite(): failed to resume handshake"); 
				
				// If the handshake isn't complete I don't think we can write with gnutls_record_send(),
				// so I don't think there's any option other than buffering what insp wants us to send.
				
				log(DEBUG, "m_ssl.so: OnRawSocketWrite(): buffering data as it cannot be written (handshake pending) Data: %s", sendbuffer);
				session->outbuf.append(sendbuffer, count);
				
				return 1;
			}
		}
		
		
		if(session->sess)
		{
			if(session->outbuf.size())
			{
				// If there's stuff in the outgoing buffer..
				session->outbuf.append(sendbuffer, count);
				sendbuffer = session->outbuf.data();
				count = session->outbuf.size();
			}
			
			int ret = gnutls_record_send(session->sess, sendbuffer, count);
		
			if(ret < 0)
			{
				if(ret != GNUTLS_E_AGAIN && ret != GNUTLS_E_INTERRUPTED)
				{
					log(DEBUG, "m_ssl.so: Error writing SSL data: %s", gnutls_strerror(ret));
				}
				else
				{
					log(DEBUG, "m_ssl.so: Not all SSL data read: %s", gnutls_strerror(ret));
				}
			}
		}
		else
		{
			log(DEBUG, "m_ssl.so: No session to write to");
		}
		
		return 1;	      	
	}

	virtual void OnUserQuit(userrec* user, std::string reason)
	{
		// Remove SSL flag
		// Free SSL stuff
	}
	
	virtual void OnSyncUserMetaData(userrec* user, Module* proto, void* opaque, std::string extname)
	{
		// check if the linking module wants to know about OUR metadata
		if(extname == "ssl")
		{
			// check if this user has an swhois field to send
			if(user->GetExt(extname))
			{
				// call this function in the linking module, let it format the data how it
				// sees fit, and send it on its way. We dont need or want to know how.
				proto->ProtoSendMetaData(opaque, TYPE_USER, user, extname, "ON");
			}
		}
	}
	
	virtual void OnDecodeMetaData(int target_type, void* target, std::string extname, std::string extdata)
	{
		// check if its our metadata key, and its associated with a user
		if ((target_type == TYPE_USER) && (extname == "ssl"))
		{
			userrec* dest = (userrec*)target;
			// if they dont already have an swhois field, accept the remote server's
			if (!dest->GetExt(extname))
			{
				dest->Extend(extname, "ON");
			}
		}
	}
	
	bool Handshake(issl_session* session)
	{		
		int ret = gnutls_handshake(session->sess);
      
      if(ret < 0)
		{
			if(ret == GNUTLS_E_AGAIN || ret == GNUTLS_E_INTERRUPTED)
			{
				// Handshake needs resuming later, read() or write() would have blocked.
				
				if(gnutls_record_get_direction(session->sess) == 0)
				{
					// gnutls_handshake() wants to read() again.
					session->status = ISSL_HANDSHAKING_READ;
					log(DEBUG, "m_ssl.so: Handshake needs resuming (reading) later, error string: %s", gnutls_strerror(ret));
				}
				else
				{
					// gnutls_handshake() wants to write() again.
					session->status = ISSL_HANDSHAKING_WRITE;
					log(DEBUG, "m_ssl.so: Handshake needs resuming (writing) later, error string: %s", gnutls_strerror(ret));
				}
			}
			else
			{
				// Handshake failed.
				CloseSession(session);
				
			   log(DEBUG, "m_ssl.so: Handshake failed, error string: %s", gnutls_strerror(ret));
			   
			   session->status = ISSL_CLOSING;
			}
			
			return false;
		}
		else
		{
			// Handshake complete.
			log(DEBUG, "m_ssl.so: Handshake completed");
			
			session->status = ISSL_HANDSHAKEN;
			
			return true;
		}
	}
	
	void CloseSession(issl_session* session)
	{
		if(session->sess)
		{
			gnutls_bye(session->sess, GNUTLS_SHUT_WR);
			gnutls_deinit(session->sess);
		}
			
		session->sess = NULL;
		session->status = ISSL_NONE;
	}
		
	int generate_dh_params()
	{
  		// Generate Diffie Hellman parameters - for use with DHE
		// kx algorithms. These should be discarded and regenerated
		// once a day, once a week or once a month. Depending on the
		// security requirements.
		 
		gnutls_dh_params_init(&dh_params);
		gnutls_dh_params_generate2(dh_params, DH_BITS);

		return 0;
	}
};

class ModuleSSLFactory : public ModuleFactory
{
 public:
	ModuleSSLFactory()
	{
	}
	
	~ModuleSSLFactory()
	{
	}
	
	virtual Module * CreateModule(Server* Me)
	{
		return new ModuleSSL(Me);
	}
};


extern "C" void * init_module( void )
{
	return new ModuleSSLFactory;
}