/***************************************************************************
 *   Copyright (C) 2012 by Timothy Pearson                                 *
 *   kb9vqf@pearsoncomputing.net                                           *
 *                                                                         *
 *   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.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include <stdlib.h>
#include <unistd.h>

#include <tqapplication.h>
#include <tqbuffer.h>

#include <sasl.h>
#include <saslplug.h>
#include <saslutil.h>

#include "tdekrbclientsocket.h"

#define NET_SEC_BUF_SIZE (2048)

/* exception handling */
struct exit_exception {
	int c;
	exit_exception(int c):c(c) { }
};

class SASLDataPrivate
{
	public:
		sasl_callback_t m_callbacks[N_CALLBACKS];
		sasl_conn_t *m_krbConnection;
};

static const char * safe_sasl_errdetail(sasl_conn_t *conn) {
	const char * str = sasl_errdetail(conn);
	if (str) {
		return str;
	}
	else {
		return "unknown error";
	}
}

static int logSASLMessages(void *context __attribute__((unused)), int priority, const char *message) {
	const char *label;

	if (!message) {
		return SASL_BADPARAM;
	}

	switch (priority) {
		case SASL_LOG_ERR:
			label = "Error";
			break;
		case SASL_LOG_NOTE:
			label = "Info";
			break;
		default:
			label = "Other";
			break;
	}

	printf("[SASL %s] %s\n\r", label, message);

	return SASL_OK;
}

TDEKerberosClientSocket::TDEKerberosClientSocket(TQObject *parent, const char *name) : TQSocket(parent, name), m_kerberosRequested(false), m_criticalSection(0), m_bufferLength(0), m_negotiatedMaxBufferSize(NET_SEC_BUF_SIZE) {
	saslData = new SASLDataPrivate;
	saslData->m_krbConnection = NULL;
	m_buffer = new TQBuffer();
	m_buffer->open(IO_ReadWrite);
}

TDEKerberosClientSocket::~TDEKerberosClientSocket() {
	m_buffer->close();
	delete m_buffer;
	delete saslData;
}

bool TDEKerberosClientSocket::open(int mode) {
	bool ret = TQSocket::open(mode);
	if (m_kerberosRequested) {
		initializeKerberosInterface();
	}
	return ret;
}

void TDEKerberosClientSocket::close() {
	TQSocket::close();
	if (m_criticalSection > 0) {
		throw exit_exception(-1);
	}
}

void TDEKerberosClientSocket::flush() {
	if (m_kerberosRequested) {
		// RAJA FIXME
	}
	else {
		TQSocket::flush();
	}
}

TQIODevice::Offset TDEKerberosClientSocket::size() const {
	TQIODevice::Offset ret;

	if (m_kerberosRequested) {
		// RAJA FIXME
	}
	else {
		ret = TQSocket::size();
	}

	return ret;
}

TQIODevice::Offset TDEKerberosClientSocket::at() const {
	return TQSocket::at();
}

bool TDEKerberosClientSocket::at(TQIODevice::Offset off) {
	bool ret;

	if (m_kerberosRequested) {
		// RAJA FIXME
	}
	else {
		ret = TQSocket::at(off);
	}

	return ret;
}

bool TDEKerberosClientSocket::atEnd() const {
	bool ret;

	if (m_kerberosRequested) {
		// RAJA FIXME
	}
	else {
		ret = TQSocket::atEnd();
	}

	return ret;
}

int TDEKerberosClientSocket::setUsingKerberos(bool krbactive) {
	int ret = 0;

	if (m_serviceName == "") {
		printf("[ERROR] No service name set!\n\r"); fflush(stdout);
		return -1;
	}

	if (krbactive) {
		m_kerberosRequested = true;
		if ((!saslData->m_krbConnection) && (state() == TQSocket::Connected)) {
			ret = initializeKerberosInterface();
		}
	}
	else {
		m_kerberosRequested = false;
		if (saslData->m_krbConnection) {
			freeKerberosConnection();
		}
	}

	return ret;
}

void TDEKerberosClientSocket::setServiceName(TQString name) {
	m_serviceName = name;
}

void TDEKerberosClientSocket::setServerFQDN(TQString name) {
	m_serverFQDN = name;
}

TQ_LONG TDEKerberosClientSocket::readBlock(char *data, TQ_ULONG maxlen) {
	TQ_LONG ret;

	if (m_kerberosRequested) {
		ret = receiveEncryptedData(data, maxlen);
	}
	else {
		ret = TQSocket::readBlock(data, maxlen);
	}

	return ret;
}

TQ_LONG TDEKerberosClientSocket::writeBlock(const char *data, TQ_ULONG len) {
	TQ_LONG ret;

	if (m_kerberosRequested) {
		ret = transmitEncryptedData(socket(), data, len);
	}
	else {
		ret = TQSocket::writeBlock(data, len);
	}

	return ret;
}

TQ_LONG TDEKerberosClientSocket::readLine(char *data, TQ_ULONG maxlen) {
	TQ_LONG ret;

	if (m_kerberosRequested) {
		ret = receiveEncryptedData(data, maxlen);
	}
	else {
		ret = TQSocket::readLine(data, maxlen);
	}

	return ret;
}

TQString TDEKerberosClientSocket::readLine() {
	TQString ret;
	char *buf;

	if (m_kerberosRequested) {
		buf = (char*)malloc(m_negotiatedMaxBufferSize);
		receiveEncryptedData(buf, m_negotiatedMaxBufferSize);
		ret = TQString(buf);
		free(buf);
	}
	else {
		ret = TQSocket::readLine();
	}

	return ret;
}

void TDEKerberosClientSocket::writeLine(TQString str) {
	if (m_kerberosRequested) {
		transmitEncryptedData(socket(), str.ascii(), str.length());
	}
	else {
		TQSocket::writeBlock(str.ascii(), str.length());
	}
}

void TDEKerberosClientSocket::freeKerberosConnection(void) {
	if (saslData->m_krbConnection) {
		sasl_dispose(&saslData->m_krbConnection);
	}
	saslData->m_krbConnection = 0;
}

void TDEKerberosClientSocket::sendSASLDataToNetwork(const char *buffer, unsigned length, int netfd) {
	char *buf;
	unsigned len, alloclen;
	int result;

	alloclen = ((length / 3) + 1) * 4 + 1;
	buf = (char*)malloc(alloclen);
	if (!buf) {
		printf("[ERROR] Unable to malloc()!\n\r");
		return;
	}

	result = sasl_encode64(buffer, length, buf, alloclen, &len);
	if (result != SASL_OK) {
		printf("[ERROR] Encoding data in base64 returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
		return;
	}

	len = strlen(buf);
	buf[len] = '\n';
	buf[len+1] = 0;
	if (write(netfd, buf, len+1) < 0) {
		// ERROR
	}

	free(buf);
}

int TDEKerberosClientSocket::getSASLDataFromNetwork(char *buf, int trunclen) {
	m_criticalSection++;
	try {
		unsigned int len;
		int result;
	
		TQByteArray ba(2048);
	
		len = 0;
		while (1) {
			tqApp->processEvents();
			if (state() != TQSocket::Connected) {
				m_criticalSection--;
				return -1;
			}
			if (TQSocket::readBlock(ba.data()+len, 1) > 0) {
				if (ba.data()[len] == '\n') {
					ba.data()[len] = 0;
					break;
				}
				if (ba.data()[len] != '\r') {
					len++;
				}
			}
			else {
				usleep(1000);
			}
			if (len >= (ba.size()-1)) {
				ba.resize(ba.size()+2048);
			}
		}
	
		len = strlen(ba.data());
		result = sasl_decode64(ba.data(), strlen(ba.data()), buf, trunclen, &len);
		if (result != SASL_OK) {
			printf("[ERROR] Decoding data from base64 returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
			m_criticalSection--;
			return -1;
		}
		buf[len] = '\0';

		m_criticalSection--;
		return len;
	}
	catch(exit_exception& e) {
		m_criticalSection--;
		return -1;
	}
}

int TDEKerberosClientSocket::transmitEncryptedData(int fd, const char* readbuf, int cc) {
	int result = 0;
	unsigned int len;
	const char *data;

	result=sasl_encode(saslData->m_krbConnection, readbuf, cc, &data, &len);
	if (result != SASL_OK) {
		printf("[ERROR] Encrypting data returned %s (%d)\n\r", safe_sasl_errdetail(saslData->m_krbConnection), result);
		return -1;
	}
	sendSASLDataToNetwork(data, len, fd);

	return 0;
}

int TDEKerberosClientSocket::receiveEncryptedData(char *buf, unsigned int trunclen) {
	unsigned int recv_len;
	const char *recv_data;
	int result;
	int len;

	char *encbuf = (char*)malloc(m_negotiatedMaxBufferSize);
	len = getSASLDataFromNetwork(encbuf, m_negotiatedMaxBufferSize);
	if (len < 0) {
		return -1;
	}
	if (len >= 0) {
		result=sasl_decode(saslData->m_krbConnection, encbuf, len, &recv_data, &recv_len);
		if (result != SASL_OK) {
			free(encbuf);
			printf("[ERROR] Decrypting data returned %s (%d)\n\r", safe_sasl_errdetail(saslData->m_krbConnection), result);
			return -1;
		}
		if (recv_len > trunclen) {
			recv_len = trunclen;
		}
		memcpy(buf, recv_data, recv_len);
	}

	free(encbuf);
	return recv_len;
}

int TDEKerberosClientSocket::initializeKerberosInterface() {
	if (state() != TQSocket::Connected) {
		saslData->m_krbConnection = false;
		return -1;
	}

	sasl_callback_t *callback;
	char buf[NET_SEC_BUF_SIZE];
	int result = 0;
	int serverlast = 0;
	sasl_security_properties_t secprops;
	const char *chosenmech;
	unsigned int len;
	int slen;
	const char *data;
	sasl_ssf_t *ssf;
	char *iplocal = NULL;
	char *ipremote = NULL;
	const char *service = m_serviceName.ascii();
	const char *fqdn = m_serverFQDN.ascii();

	callback = saslData->m_callbacks;

	// log
	callback->id = SASL_CB_LOG;
	callback->proc = (sasl_callback_ft)&logSASLMessages;
	callback->context = NULL;
	++callback;

	// end of callback list
	callback->id = SASL_CB_LIST_END;
	callback->proc = NULL;
	callback->context = NULL;
	++callback;

	// Initialize default data structures
	memset(&secprops, 0L, sizeof(secprops));
	secprops.maxbufsize = NET_SEC_BUF_SIZE;
	secprops.max_ssf = UINT_MAX;

	result = sasl_client_init(saslData->m_callbacks);
	if (result != SASL_OK) {
		printf("[ERROR] Initializing libsasl returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
		return -1;
	}

	result = sasl_client_new(service, fqdn, iplocal, ipremote, NULL, serverlast, &saslData->m_krbConnection);
	if (result != SASL_OK) {
		printf("[ERROR] Allocating sasl connection state returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
		return -1;
	}

	result = sasl_setprop(saslData->m_krbConnection, SASL_SEC_PROPS, &secprops);
	if (result != SASL_OK) {
		printf("[ERROR] Setting security properties returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
		freeKerberosConnection();
		return -1;
	}

	printf("[DEBUG] Waiting for mechanism list from server...\n\r");
	slen = getSASLDataFromNetwork(buf, NET_SEC_BUF_SIZE);
	if (slen < 0) {
		return -2;
	}
	len = slen;

	printf("Choosing best mechanism from: %s\n", buf);

	result = sasl_client_start(saslData->m_krbConnection, buf, NULL, &data, &len, &chosenmech);
	if (result != SASL_OK && result != SASL_CONTINUE) {
		printf("[ERROR] Starting SASL negotiation returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
		freeKerberosConnection();
		return -1;
	}

	printf("[DEBUG] Using mechanism %s\n\r", chosenmech);
	strcpy(buf, chosenmech);
	if (data) {
		if (NET_SEC_BUF_SIZE - strlen(buf) - 1 < len) {
			printf("[ERROR] Insufficient buffer space to construct initial response!\n\r");
			freeKerberosConnection();
			return -1;
		}
		printf("[DEBUG] Preparing initial response...\n\r");
		memcpy(buf + strlen(buf) + 1, data, len);
		len += (unsigned) strlen(buf) + 1;
		data = NULL;
	}
	else {
		len = (unsigned) strlen(buf);
	}

	printf("[DEBUG] Sending initial response...\n\r");
	sendSASLDataToNetwork(buf, len, socket());

	while (result == SASL_CONTINUE) {
		printf("[DEBUG] Waiting for server reply...\n\r");
		slen = getSASLDataFromNetwork(buf, NET_SEC_BUF_SIZE);
		if (slen < 0) {
			return -2;
		}
		len = slen;
		result = sasl_client_step(saslData->m_krbConnection, buf, len, NULL, &data, &len);
		if (result != SASL_OK && result != SASL_CONTINUE) {
			printf("[ERROR] Performing SASL negotiation returned %s (%d)\n\r", sasl_errstring(result, NULL, NULL), result);
			freeKerberosConnection();
			return -1;
		}
		if (data && len) {
			printf("[DEBUG] Sending response...\n\r");
			sendSASLDataToNetwork(data, len, socket());
		}
		else if (result != SASL_OK || !serverlast) {
			sendSASLDataToNetwork("", 0, socket());
		}
	}
	printf("[DEBUG] Negotiation complete!\n\r");

	result = sasl_getprop(saslData->m_krbConnection, SASL_USERNAME, (const void **)&data);
	if (result != SASL_OK) {
		printf("[WARNING] Unable to determine authenticated username!\n\r");
	}
	else {
		printf("[DEBUG] Authenticated username: %s\n\r", data ? data : "(NULL)");
	}

	result = sasl_getprop(saslData->m_krbConnection, SASL_DEFUSERREALM, (const void **)&data);
	if (result != SASL_OK) {
		printf("[WARNING] Unable to determine authenticated realm!\n\r");
	}
	else {
		printf("[DEBUG] Authenticated realm: %s\n\r", data ? data : "(NULL)");
	}

	result = sasl_getprop(saslData->m_krbConnection, SASL_SSF, (const void **)&ssf);
	if (result != SASL_OK) {
		printf("[WARNING] Unable to determine SSF!\n\r");
	}
	else {
		printf("[DEBUG] Authenticated SSF: %d\n", *ssf);
	}

	result = sasl_getprop(saslData->m_krbConnection, SASL_MAXOUTBUF, (const void **)&m_negotiatedMaxBufferSize);
	if (result != SASL_OK) {
		printf("[WARNING] Unable to determine maximum buffer size!\n\r");
		m_negotiatedMaxBufferSize = NET_SEC_BUF_SIZE;
	}
	else {
		printf("[DEBUG] Maximum buffer size: %d\n", m_negotiatedMaxBufferSize);
	}

	return 0;
}