/*
 * Copyright (c) 2013 INSIDE Secure Corporation
 * Copyright (c) PeerSec Networks, 2002-2011
 * All Rights Reserved
 *
 * The latest version of this code is available at http://www.matrixssl.org
 *
 * This software is open source; 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 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
 * http://www.gnu.org/copyleft/gpl.html
 */
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <stdarg.h>
#include <fcntl.h>
#include <stdio.h>
#include <time.h>
#include <poll.h>
#include <sys/socket.h>

#include "matrixssl/matrixsslApi.h"

//#warning "DO NOT USE THESE DEFAULT KEYS IN PRODUCTION ENVIRONMENTS."

/*
 * If supporting client authentication, pick ONE identity to auto select a
 * certificate and private key that support desired algorithms.
 */
#define ID_RSA /* RSA Certificate and Key */

#define USE_HEADER_KEYS

/* If the algorithm type is supported, load a CA for it */
#ifdef USE_HEADER_KEYS
/* CAs */
# include "sampleCerts/RSA/ALL_RSA_CAS.h"
/* Identity Certs and Keys for use with Client Authentication */
# ifdef ID_RSA
#  define EXAMPLE_RSA_KEYS
#  include "sampleCerts/RSA/2048_RSA.h"
#  include "sampleCerts/RSA/2048_RSA_KEY.h"
# endif
#endif

static ssize_t safe_write(int fd, const void *buf, size_t count)
{
	ssize_t n;

	do {
		n = write(fd, buf, count);
	} while (n < 0 && errno == EINTR);

	return n;
}

static ssize_t full_write(int fd, const void *buf, size_t len)
{
	ssize_t cc;
	ssize_t total;

	total = 0;

	while (len) {
		cc = safe_write(fd, buf, len);

		if (cc < 0) {
			if (total) {
				/* we already wrote some! */
				/* user can do another write to know the error code */
				return total;
			}
			return cc;  /* write() returns -1 on failure. */
		}

		total += cc;
		buf = ((const char *)buf) + cc;
		len -= cc;
	}

	return total;
}

static void say(const char *s, ...)
{
	char buf[256];
	va_list p;
	int sz;

	va_start(p, s);
	sz = vsnprintf(buf, sizeof(buf), s, p);
	full_write(STDERR_FILENO, buf, sz >= 0 && sz < sizeof(buf) ? sz : strlen(buf));
	va_end(p);
}

static void die(const char *s, ...)
{
	char buf[256];
	va_list p;
	int sz;

	va_start(p, s);
	sz = vsnprintf(buf, sizeof(buf), s, p);
	full_write(STDERR_FILENO, buf, sz >= 0 && sz < sizeof(buf) ? sz : strlen(buf));
	exit(1);
	va_end(p);
}

#if 0
# define dbg(...) say(__VA_ARGS__)
#else
# define dbg(...) ((void)0)
#endif

static struct pollfd pfd[2] = {
	{ -1, POLLIN|POLLERR|POLLHUP, 0 },
	{ -1, POLLIN|POLLERR|POLLHUP, 0 },
};
#define STDIN           pfd[0]
#define NETWORK         pfd[1]
#define STDIN_READY()   (pfd[0].revents & (POLLIN|POLLERR|POLLHUP))
#define NETWORK_READY() (pfd[1].revents & (POLLIN|POLLERR|POLLHUP))

static int wait_for_input(void)
{
	if (STDIN.fd == NETWORK.fd) /* means both are -1 */
		exit(0);
	dbg("polling\n");
	STDIN.revents = NETWORK.revents = 0;
	return poll(pfd, 2, -1);
}

static int32 certCb(ssl_t *ssl, psX509Cert_t *cert, int32 alert)
{
	/* Example to allow anonymous connections based on a define */
	if (alert > 0) {
		return SSL_ALLOW_ANON_CONNECTION; // = 254
	}
#if 0
	/* Validate the 'not before' and 'not after' dates, etc */
	return PS_FAILURE; /* if we don't like this cert */
#endif
	return PS_SUCCESS;
}

static void close_conn_and_exit(ssl_t *ssl, int fd)
{
	unsigned char *buf;
	int len;

	fcntl(fd, F_SETFL, fcntl(fd, F_GETFL) | O_NONBLOCK);
	/* Quick attempt to send a closure alert, don't worry about failure */
	if (matrixSslEncodeClosureAlert(ssl) >= 0) {
		len = matrixSslGetOutdata(ssl, &buf);
		if (len > 0) {
			len = safe_write(fd, buf, len);
			//if (len > 0) {
			//	matrixSslSentData(ssl, len);
			//}
		}
	}
	//matrixSslDeleteSession(ssl);
	shutdown(fd, SHUT_WR);
	exit(0);
}

static int encode_data(ssl_t *ssl, const void *data, int len)
{
	unsigned char *buf;
	int available;

	available = matrixSslGetWritebuf(ssl, &buf, len);
	if (available < 0)
		die("matrixSslGetWritebuf\n");
	if (len > available)
		die("len > available\n");
	memcpy(buf, data, len);
	if (matrixSslEncodeWritebuf(ssl, len) < 0)
		die("matrixSslEncodeWritebuf\n");
	return len;
}

static void flush_to_net(ssl_t *ssl, int fd)
{
	int rc;
	int len;
	unsigned char *buf;

	while ((len = matrixSslGetOutdata(ssl, &buf)) > 0) {
		dbg("writing net %d bytes\n", len);
		if (full_write(fd, buf, len) != len)
			die("write to network\n");
		rc = matrixSslSentData(ssl, len);
		if (rc < 0)
			die("matrixSslSentData\n");
	}
}

static void do_io_until_eof_and_exit(int fd, sslKeys_t *keys)
{
	int rc;
	int len;
	uint32_t len32u;
	sslSessionId_t *sid;
	ssl_t *ssl;
	unsigned char *buf;

	NETWORK.fd = fd;
	/* Note! STDIN.fd is disabled (-1) until SSL handshake is over:
	 * we do not attempt to feed any user data to MatrixSSL
	 * before it is ready.
	 */

	matrixSslNewSessionId(&sid);
	rc = matrixSslNewClientSession(&ssl, keys, sid, 0, certCb, NULL, NULL, 0);
dbg("matrixSslNewClientSession:rc=%d\n", rc);
	if (rc != MATRIXSSL_REQUEST_SEND)
		die("matrixSslNewClientSession\n");

	len = 0; /* only to suppress compiler warning */
 again:
	switch (rc) {
	case MATRIXSSL_REQUEST_SEND:
		dbg("MATRIXSSL_REQUEST_SEND\n");
		flush_to_net(ssl, fd);
		goto poll_input;

	case 0:
		dbg("rc==0\n");
		flush_to_net(ssl, fd);
		goto poll_input;

	case MATRIXSSL_REQUEST_CLOSE:
		/* what does this mean if we are here? */
		dbg("MATRIXSSL_REQUEST_CLOSE\n");
		close_conn_and_exit(ssl, fd);

	case MATRIXSSL_HANDSHAKE_COMPLETE:
		dbg("MATRIXSSL_HANDSHAKE_COMPLETE\n");
		/* Init complete, can start reading local user's data: */
		STDIN.fd = STDIN_FILENO;
 poll_input:
		wait_for_input();
		if (STDIN_READY()) {
			char ibuf[4 * 1024];
			dbg("reading stdin\n");
			len = read(STDIN_FILENO, ibuf, sizeof(ibuf));
			if (len < 0)
				die("read error on stdin\n");
			if (len == 0)
				STDIN.fd = -1;
			else {
				len = encode_data(ssl, ibuf, len);
				if (len) {
					rc = MATRIXSSL_REQUEST_SEND;
dbg("rc=%d\n", rc);
					goto again;
				}
			}
		}
 read_network:
		if (NETWORK_READY()) {
			dbg("%s%s%s\n",
				(pfd[1].revents & POLLIN)  ? "POLLIN"  : "",
				(pfd[1].revents & POLLERR) ? "|POLLERR" : "",
				(pfd[1].revents & POLLHUP) ? "|POLLHUP" : ""
			);
			len = matrixSslGetReadbuf(ssl, &buf);
			if (len <= 0)
				die("matrixSslGetReadbuf\n");
			dbg("reading net up to %d\n", len);
			len = read(fd, buf, len);
			dbg("reading net:%d\n", len);
			if (len < 0)
				die("read error on network\n");
			if (len == 0) /*eof*/
				NETWORK.fd = -1;
			len32u = len;
			rc = matrixSslReceivedData(ssl, len, &buf, &len32u);
dbg("matrixSslReceivedData:rc=%d\n", rc);
			len = len32u;
			if (rc < 0)
				die("matrixSslReceivedData\n");
		}
		goto again;

	case MATRIXSSL_APP_DATA:
		dbg("MATRIXSSL_APP_DATA: writing stdout\n");
		do {
			if (full_write(STDOUT_FILENO, buf, len) != len)
				die("write to stdout\n");
			len32u = len;
			rc = matrixSslProcessedData(ssl, &buf, &len32u);
//this was seen returning rc=0:
dbg("matrixSslProcessedData:rc=%d\n", rc);
			len = len32u;
		} while (rc == MATRIXSSL_APP_DATA);
		if (pfd[1].fd == -1) {
			/* Already saw EOF on network, and we processed
			 * and wrote out all ssl data. Signal it:
			 */
			close(STDOUT_FILENO);
		}
		goto again;

	case MATRIXSSL_REQUEST_RECV:
		dbg("MATRIXSSL_REQUEST_RECV\n");
		wait_for_input();
		goto read_network;

	case MATRIXSSL_RECEIVED_ALERT:
		dbg("MATRIXSSL_RECEIVED_ALERT\n");
		/* The first byte of the buffer is the level */
		/* The second byte is the description */
		if (buf[0] == SSL_ALERT_LEVEL_FATAL)
			die("Fatal alert\n");
		/* Closure alert is normal (and best) way to close */
		if (buf[1] == SSL_ALERT_CLOSE_NOTIFY)
			close_conn_and_exit(ssl, fd);
		die("Warning alert\n");
		len32u = len;
		rc = matrixSslProcessedData(ssl, &buf, &len32u);
dbg("matrixSslProcessedData:rc=%d\n", rc);
		len = len32u;
		goto again;

	default:
		/* If rc < 0 it is an error */
		die("bad rc:%d\n", rc);
	}
}

static sslKeys_t* make_keys(void)
{
	int rc, CAstreamLen;
	char *CAstream;
	sslKeys_t *keys;

	if (matrixSslNewKeys(&keys) < 0)
		die("matrixSslNewKeys\n");

#ifdef USE_HEADER_KEYS
	/*
	 * In-memory based keys
	 * Build the CA list first for potential client auth usage
	 */
	CAstream = NULL;
	CAstreamLen = sizeof(RSACAS);
	if (CAstreamLen > 0) {
		CAstream = psMalloc(NULL, CAstreamLen);
		memcpy(CAstream, RSACAS, sizeof(RSACAS));
	}

 #ifdef ID_RSA
	rc = matrixSslLoadRsaKeysMem(keys, RSA2048, sizeof(RSA2048),
			RSA2048KEY, sizeof(RSA2048KEY), (unsigned char*)CAstream,
			CAstreamLen);
	if (rc < 0)
		die("matrixSslLoadRsaKeysMem\n");
 #endif

	if (CAstream)
		psFree(CAstream);
#endif /* USE_HEADER_KEYS */
	return keys;
}

int main(int argc, char **argv)
{
	int fd;
	char *fd_str;

	if (!argv[1])
		die("Syntax error\n");
	if (argv[1][0] != '-')
		die("Syntax error\n");
	if (argv[1][1] != 'd')
		die("Syntax error\n");
	fd_str = argv[1] + 2;
	if (!fd_str[0])
		fd_str = argv[2];
	if (!fd_str || fd_str[0] < '0' || fd_str[0] > '9')
		die("Syntax error\n");

	fd = atoi(fd_str);
	if (fd < 3)
		die("Syntax error\n");

	if (matrixSslOpen() < 0)
		die("matrixSslOpen\n");

	do_io_until_eof_and_exit(fd, make_keys());
	/* does not return */

	return 0;
}