/*
 * Copyright (C) 2010 Martin Willi
 * Copyright (C) 2010 revosec AG
 *
 * 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.  See <http://www.fsf.org/copyleft/gpl.txt>.
 *
 * 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.
 */

#include "tls_server.h"

#include <time.h>

#include <debug.h>
#include <credentials/certificates/x509.h>

typedef struct private_tls_server_t private_tls_server_t;


typedef enum {
	STATE_INIT,
	STATE_HELLO_RECEIVED,
	STATE_HELLO_SENT,
	STATE_CERT_SENT,
	STATE_CERTREQ_SENT,
	STATE_HELLO_DONE,
	STATE_CERT_RECEIVED,
	STATE_KEY_EXCHANGE_RECEIVED,
	STATE_CERT_VERIFY_RECEIVED,
	STATE_CIPHERSPEC_CHANGED_IN,
	STATE_FINISHED_RECEIVED,
	STATE_CIPHERSPEC_CHANGED_OUT,
	STATE_FINISHED_SENT,
} server_state_t;

/**
 * Private data of an tls_server_t object.
 */
struct private_tls_server_t {

	/**
	 * Public tls_server_t interface.
	 */
	tls_server_t public;

	/**
	 * TLS stack
	 */
	tls_t *tls;

	/**
	 * TLS crypto context
	 */
	tls_crypto_t *crypto;

	/**
	 * TLS alert handler
	 */
	tls_alert_t *alert;

	/**
	 * Server identity
	 */
	identification_t *server;

	/**
	 * Peer identity, NULL for no client authentication
	 */
	identification_t *peer;

	/**
	 * State we are in
	 */
	server_state_t state;

	/**
	 * Hello random data selected by client
	 */
	char client_random[32];

	/**
	 * Hello random data selected by server
	 */
	char server_random[32];

	/**
	 * Auth helper for peer authentication
	 */
	auth_cfg_t *peer_auth;

	/**
	 * Auth helper for server authentication
	 */
	auth_cfg_t *server_auth;

	/**
	 * Peer private key
	 */
	private_key_t *private;

	/**
	 * Selected TLS cipher suite
	 */
	tls_cipher_suite_t suite;
};

/**
 * Process client hello message
 */
static status_t process_client_hello(private_tls_server_t *this,
									 tls_reader_t *reader)
{
	u_int16_t version;
	chunk_t random, session, ciphers, compression, ext = chunk_empty;
	tls_cipher_suite_t *suites;
	int count, i;

	this->crypto->append_handshake(this->crypto,
								   TLS_CLIENT_HELLO, reader->peek(reader));

	if (!reader->read_uint16(reader, &version) ||
		!reader->read_data(reader, sizeof(this->client_random), &random) ||
		!reader->read_data8(reader, &session) ||
		!reader->read_data16(reader, &ciphers) ||
		!reader->read_data8(reader, &compression) ||
		(reader->remaining(reader) && !reader->read_data16(reader, &ext)))
	{
		DBG1(DBG_TLS, "received invalid ClientHello");
		this->alert->add(this->alert, TLS_FATAL, TLS_DECODE_ERROR);
		return NEED_MORE;
	}

	memcpy(this->client_random, random.ptr, sizeof(this->client_random));

	if (!this->tls->set_version(this->tls, version))
	{
		DBG1(DBG_TLS, "negotiated version %N not supported",
			 tls_version_names, version);
		this->alert->add(this->alert, TLS_FATAL, TLS_PROTOCOL_VERSION);
		return NEED_MORE;
	}
	count = ciphers.len / sizeof(u_int16_t);
	suites = alloca(count * sizeof(tls_cipher_suite_t));
	DBG2(DBG_TLS, "received %d TLS cipher suites:", count);
	for (i = 0; i < count; i++)
	{
		suites[i] = untoh16(&ciphers.ptr[i * sizeof(u_int16_t)]);
		DBG2(DBG_TLS, "  %N", tls_cipher_suite_names, suites[i]);
	}
	this->suite = this->crypto->select_cipher_suite(this->crypto, suites, count);
	if (!this->suite)
	{
		DBG1(DBG_TLS, "received cipher suites inacceptable");
		this->alert->add(this->alert, TLS_FATAL, TLS_HANDSHAKE_FAILURE);
		return NEED_MORE;
	}
	DBG1(DBG_TLS, "negotiated TLS version %N with suite %N",
		 tls_version_names, version, tls_cipher_suite_names, this->suite);
	this->state = STATE_HELLO_RECEIVED;
	return NEED_MORE;
}

/**
 * Process certificate
 */
static status_t process_certificate(private_tls_server_t *this,
									tls_reader_t *reader)
{
	certificate_t *cert;
	tls_reader_t *certs;
	chunk_t data;
	bool first = TRUE;

	this->crypto->append_handshake(this->crypto,
								   TLS_CERTIFICATE, reader->peek(reader));

	if (!reader->read_data24(reader, &data))
	{
		DBG1(DBG_TLS, "certificate message header invalid");
		this->alert->add(this->alert, TLS_FATAL, TLS_DECODE_ERROR);
		return NEED_MORE;
	}
	certs = tls_reader_create(data);
	while (certs->remaining(certs))
	{
		if (!certs->read_data24(certs, &data))
		{
			DBG1(DBG_TLS, "certificate message invalid");
			this->alert->add(this->alert, TLS_FATAL, TLS_DECODE_ERROR);
			certs->destroy(certs);
			return NEED_MORE;
		}
		cert = lib->creds->create(lib->creds, CRED_CERTIFICATE, CERT_X509,
								  BUILD_BLOB_ASN1_DER, data, BUILD_END);
		if (cert)
		{
			if (first)
			{
				this->peer_auth->add(this->peer_auth,
									 AUTH_HELPER_SUBJECT_CERT, cert);
				DBG1(DBG_TLS, "received TLS peer certificate '%Y'",
					 cert->get_subject(cert));
				first = FALSE;
			}
			else
			{
				DBG1(DBG_TLS, "received TLS intermediate certificate '%Y'",
					 cert->get_subject(cert));
				this->peer_auth->add(this->peer_auth, AUTH_HELPER_IM_CERT, cert);
			}
		}
		else
		{
			DBG1(DBG_TLS, "parsing TLS certificate failed, skipped");
			this->alert->add(this->alert, TLS_WARNING, TLS_BAD_CERTIFICATE);
		}
	}
	certs->destroy(certs);
	this->state = STATE_CERT_RECEIVED;
	return NEED_MORE;
}

/**
 * Process Client Key Exchange
 */
static status_t process_key_exchange(private_tls_server_t *this,
									 tls_reader_t *reader)
{
	chunk_t encrypted, premaster;

	this->crypto->append_handshake(this->crypto,
								   TLS_CLIENT_KEY_EXCHANGE, reader->peek(reader));

	if (!reader->read_data16(reader, &encrypted))
	{
		DBG1(DBG_TLS, "received invalid Client Key Exchange");
		this->alert->add(this->alert, TLS_FATAL, TLS_DECODE_ERROR);
		return NEED_MORE;
	}

	if (!this->private ||
		!this->private->decrypt(this->private, ENCRYPT_RSA_PKCS1,
								encrypted, &premaster))
	{
		DBG1(DBG_TLS, "decrypting Client Key Exchange data failed");
		this->alert->add(this->alert, TLS_FATAL, TLS_DECRYPT_ERROR);
		return NEED_MORE;
	}
	this->crypto->derive_secrets(this->crypto, premaster,
								 chunk_from_thing(this->client_random),
								 chunk_from_thing(this->server_random));
	chunk_clear(&premaster);

	this->state = STATE_KEY_EXCHANGE_RECEIVED;
	return NEED_MORE;
}

/**
 * Process Certificate verify
 */
static status_t process_cert_verify(private_tls_server_t *this,
									tls_reader_t *reader)
{
	bool verified = FALSE;
	enumerator_t *enumerator;
	public_key_t *public;
	auth_cfg_t *auth;
	tls_reader_t *sig;

	enumerator = lib->credmgr->create_public_enumerator(lib->credmgr,
										KEY_ANY, this->peer, this->peer_auth);
	while (enumerator->enumerate(enumerator, &public, &auth))
	{
		sig = tls_reader_create(reader->peek(reader));
		verified = this->crypto->verify_handshake(this->crypto, public, sig);
		sig->destroy(sig);
		if (verified)
		{
			break;
		}
		DBG1(DBG_TLS, "signature verification failed, trying another key");
	}
	enumerator->destroy(enumerator);

	if (!verified)
	{
		DBG1(DBG_TLS, "no trusted certificate found for '%Y' to verify TLS peer",
			 this->peer);
		this->alert->add(this->alert, TLS_FATAL, TLS_CERTIFICATE_UNKNOWN);
		return NEED_MORE;
	}

	this->crypto->append_handshake(this->crypto,
								   TLS_CERTIFICATE_VERIFY, reader->peek(reader));
	this->state = STATE_CERT_VERIFY_RECEIVED;
	return NEED_MORE;
}

/**
 * Process finished message
 */
static status_t process_finished(private_tls_server_t *this,
								 tls_reader_t *reader)
{
	chunk_t received;
	char buf[12];

	if (!reader->read_data(reader, sizeof(buf), &received))
	{
		DBG1(DBG_TLS, "received client finished too short");
		this->alert->add(this->alert, TLS_FATAL, TLS_DECODE_ERROR);
		return NEED_MORE;
	}
	if (!this->crypto->calculate_finished(this->crypto, "client finished", buf))
	{
		DBG1(DBG_TLS, "calculating client finished failed");
		this->alert->add(this->alert, TLS_FATAL, TLS_INTERNAL_ERROR);
		return NEED_MORE;
	}
	if (!chunk_equals(received, chunk_from_thing(buf)))
	{
		DBG1(DBG_TLS, "received client finished invalid");
		this->alert->add(this->alert, TLS_FATAL, TLS_DECRYPT_ERROR);
		return NEED_MORE;
	}

	this->crypto->append_handshake(this->crypto, TLS_FINISHED, received);
	this->state = STATE_FINISHED_RECEIVED;
	return NEED_MORE;
}

METHOD(tls_handshake_t, process, status_t,
	private_tls_server_t *this, tls_handshake_type_t type, tls_reader_t *reader)
{
	tls_handshake_type_t expected;

	switch (this->state)
	{
		case STATE_INIT:
			if (type == TLS_CLIENT_HELLO)
			{
				return process_client_hello(this, reader);
			}
			expected = TLS_CLIENT_HELLO;
			break;
		case STATE_HELLO_DONE:
			if (type == TLS_CERTIFICATE)
			{
				return process_certificate(this, reader);
			}
			if (this->peer)
			{
				expected = TLS_CERTIFICATE;
				break;
			}
			/* otherwise fall through to next state */
		case STATE_CERT_RECEIVED:
			if (type == TLS_CLIENT_KEY_EXCHANGE)
			{
				return process_key_exchange(this, reader);
			}
			expected = TLS_CLIENT_KEY_EXCHANGE;
			break;
		case STATE_KEY_EXCHANGE_RECEIVED:
			if (type == TLS_CERTIFICATE_VERIFY)
			{
				return process_cert_verify(this, reader);
			}
			if (this->peer)
			{
				expected = TLS_CERTIFICATE_VERIFY;
				break;
			}
			else
			{
				return INVALID_STATE;
			}
		case STATE_CIPHERSPEC_CHANGED_IN:
			if (type == TLS_FINISHED)
			{
				return process_finished(this, reader);
			}
			expected = TLS_FINISHED;
			break;
		default:
			DBG1(DBG_TLS, "TLS %N not expected in current state",
				 tls_handshake_type_names, type);
			this->alert->add(this->alert, TLS_FATAL, TLS_UNEXPECTED_MESSAGE);
			return NEED_MORE;
	}
	DBG1(DBG_TLS, "TLS %N expected, but received %N",
		 tls_handshake_type_names, expected, tls_handshake_type_names, type);
	this->alert->add(this->alert, TLS_FATAL, TLS_UNEXPECTED_MESSAGE);
	return NEED_MORE;
}

/**
 * Send ServerHello message
 */
static status_t send_server_hello(private_tls_server_t *this,
							tls_handshake_type_t *type, tls_writer_t *writer)
{
	tls_version_t version;
	rng_t *rng;

	htoun32(&this->server_random, time(NULL));
	rng = lib->crypto->create_rng(lib->crypto, RNG_WEAK);
	if (!rng)
	{
		DBG1(DBG_TLS, "no suitable RNG found to generate server random");
		this->alert->add(this->alert, TLS_FATAL, TLS_INTERNAL_ERROR);
		return FAILED;
	}
	rng->get_bytes(rng, sizeof(this->server_random) - 4, this->server_random + 4);
	rng->destroy(rng);

	/* TLS version */
	version = this->tls->get_version(this->tls);
	writer->write_uint16(writer, version);
	writer->write_data(writer, chunk_from_thing(this->server_random));

	/* session identifier => none, we don't support session resumption */
	writer->write_data8(writer, chunk_empty);

	/* add selected TLS cipher suite */
	writer->write_uint16(writer, this->suite);

	/* NULL compression only */
	writer->write_uint8(writer, 0);

	*type = TLS_SERVER_HELLO;
	this->state = STATE_HELLO_SENT;
	this->crypto->append_handshake(this->crypto, *type, writer->get_buf(writer));
	return NEED_MORE;
}

/**
 * Send Certificate
 */
static status_t send_certificate(private_tls_server_t *this,
							tls_handshake_type_t *type, tls_writer_t *writer)
{
	enumerator_t *enumerator;
	certificate_t *cert;
	auth_rule_t rule;
	tls_writer_t *certs;
	chunk_t data;

	this->private = lib->credmgr->get_private(lib->credmgr,
									KEY_ANY, this->server, this->server_auth);
	if (!this->private)
	{
		DBG1(DBG_TLS, "no TLS server certificate found for '%Y'", this->server);
		this->alert->add(this->alert, TLS_FATAL, TLS_INTERNAL_ERROR);
		return FAILED;
	}

	/* generate certificate payload */
	certs = tls_writer_create(256);
	cert = this->server_auth->get(this->server_auth, AUTH_RULE_SUBJECT_CERT);
	if (cert)
	{
		if (cert->get_encoding(cert, CERT_ASN1_DER, &data))
		{
			DBG1(DBG_TLS, "sending TLS server certificate '%Y'",
				 cert->get_subject(cert));
			certs->write_data24(certs, data);
			free(data.ptr);
		}
	}
	enumerator = this->server_auth->create_enumerator(this->server_auth);
	while (enumerator->enumerate(enumerator, &rule, &cert))
	{
		if (rule == AUTH_RULE_IM_CERT)
		{
			if (cert->get_encoding(cert, CERT_ASN1_DER, &data))
			{
				DBG1(DBG_TLS, "sending TLS intermediate certificate '%Y'",
					 cert->get_subject(cert));
				certs->write_data24(certs, data);
				free(data.ptr);
			}
		}
	}
	enumerator->destroy(enumerator);

	writer->write_data24(writer, certs->get_buf(certs));
	certs->destroy(certs);

	*type = TLS_CERTIFICATE;
	this->state = STATE_CERT_SENT;
	this->crypto->append_handshake(this->crypto, *type, writer->get_buf(writer));
	return NEED_MORE;
}

/**
 * Create a list of supported certificate types and hash/sig algorithms
 */
static void get_supported_algorithms(private_tls_server_t *this,
									 tls_writer_t *writer)
{
	tls_writer_t *supported;

	supported = tls_writer_create(4);
	/* we propose both RSA and ECDSA */
	supported->write_uint8(supported, TLS_RSA_SIGN);
	supported->write_uint8(supported, TLS_ECDSA_SIGN);
	writer->write_data8(writer, supported->get_buf(supported));
	supported->destroy(supported);

	if (this->tls->get_version(this->tls) >= TLS_1_2)
	{
		enumerator_t *enumerator;
		hash_algorithm_t alg;
		tls_hash_algorithm_t hash;

		supported = tls_writer_create(32);
		enumerator = lib->crypto->create_hasher_enumerator(lib->crypto);
		while (enumerator->enumerate(enumerator, &alg))
		{
			switch (alg)
			{
				case HASH_MD5:
					hash = TLS_HASH_MD5;
					break;
				case HASH_SHA1:
					hash = TLS_HASH_SHA1;
					break;
				case HASH_SHA224:
					hash = TLS_HASH_SHA224;
					break;
				case HASH_SHA256:
					hash = TLS_HASH_SHA256;
					break;
				case HASH_SHA384:
					hash = TLS_HASH_SHA384;
					break;
				case HASH_SHA512:
					hash = TLS_HASH_SHA512;
					break;
				default:
					continue;
			}
			supported->write_uint8(supported, hash);
			supported->write_uint8(supported, TLS_SIG_RSA);
			if (alg != HASH_MD5 && alg != HASH_SHA224)
			{
				supported->write_uint8(supported, hash);
				supported->write_uint8(supported, TLS_SIG_ECDSA);
			}
		}
		enumerator->destroy(enumerator);

		writer->write_data16(writer, supported->get_buf(supported));
		supported->destroy(supported);
	}
}

/**
 * Send Certificate Request
 */
static status_t send_certificate_request(private_tls_server_t *this,
							tls_handshake_type_t *type, tls_writer_t *writer)
{
	tls_writer_t *authorities;
	enumerator_t *enumerator;
	certificate_t *cert;
	x509_t *x509;
	identification_t *id;

	get_supported_algorithms(this, writer);

	authorities = tls_writer_create(64);
	enumerator = lib->credmgr->create_cert_enumerator(lib->credmgr,
												CERT_X509, KEY_RSA, NULL, TRUE);
	while (enumerator->enumerate(enumerator, &cert))
	{
		x509 = (x509_t*)cert;
		if (x509->get_flags(x509) & X509_CA)
		{
			id = cert->get_subject(cert);
			DBG1(DBG_TLS, "sending TLS cert request for '%Y'", id);
			authorities->write_data16(authorities, id->get_encoding(id));
		}
	}
	enumerator->destroy(enumerator);
	writer->write_data16(writer, authorities->get_buf(authorities));
	authorities->destroy(authorities);

	*type = TLS_CERTIFICATE_REQUEST;
	this->state = STATE_CERTREQ_SENT;
	this->crypto->append_handshake(this->crypto, *type, writer->get_buf(writer));
	return NEED_MORE;
}

/**
 * Send Hello Done
 */
static status_t send_hello_done(private_tls_server_t *this,
							tls_handshake_type_t *type, tls_writer_t *writer)
{
	*type = TLS_SERVER_HELLO_DONE;
	this->state = STATE_HELLO_DONE;
	this->crypto->append_handshake(this->crypto, *type, writer->get_buf(writer));
	return NEED_MORE;
}

/**
 * Send Finished
 */
static status_t send_finished(private_tls_server_t *this,
							  tls_handshake_type_t *type, tls_writer_t *writer)
{
	char buf[12];

	if (!this->crypto->calculate_finished(this->crypto, "server finished", buf))
	{
		DBG1(DBG_TLS, "calculating server finished data failed");
		this->alert->add(this->alert, TLS_FATAL, TLS_INTERNAL_ERROR);
		return FAILED;
	}

	writer->write_data(writer, chunk_from_thing(buf));

	*type = TLS_FINISHED;
	this->state = STATE_FINISHED_SENT;
	this->crypto->derive_eap_msk(this->crypto,
								 chunk_from_thing(this->client_random),
								 chunk_from_thing(this->server_random));
	return NEED_MORE;
}

METHOD(tls_handshake_t, build, status_t,
	private_tls_server_t *this, tls_handshake_type_t *type, tls_writer_t *writer)
{
	switch (this->state)
	{
		case STATE_HELLO_RECEIVED:
			return send_server_hello(this, type, writer);
		case STATE_HELLO_SENT:
			return send_certificate(this, type, writer);
		case STATE_CERT_SENT:
			if (this->peer)
			{
				return send_certificate_request(this, type, writer);
			}
			/* otherwise fall through to next state */
		case STATE_CERTREQ_SENT:
			return send_hello_done(this, type, writer);
		case STATE_CIPHERSPEC_CHANGED_OUT:
			return send_finished(this, type, writer);
		case STATE_FINISHED_SENT:
			return INVALID_STATE;
		default:
			return INVALID_STATE;
	}
}

METHOD(tls_handshake_t, cipherspec_changed, bool,
	private_tls_server_t *this)
{
	if (this->state == STATE_FINISHED_RECEIVED)
	{
		this->crypto->change_cipher(this->crypto, FALSE);
		this->state = STATE_CIPHERSPEC_CHANGED_OUT;
		return TRUE;
	}
	return FALSE;
}

METHOD(tls_handshake_t, change_cipherspec, bool,
	private_tls_server_t *this)
{
	if ((this->peer && this->state == STATE_CERT_VERIFY_RECEIVED) ||
	   (!this->peer && this->state == STATE_KEY_EXCHANGE_RECEIVED))
	{
		this->crypto->change_cipher(this->crypto, TRUE);
		this->state = STATE_CIPHERSPEC_CHANGED_IN;
		return TRUE;
	}
	return FALSE;
}

METHOD(tls_handshake_t, finished, bool,
	private_tls_server_t *this)
{
	return this->state == STATE_FINISHED_SENT;
}

METHOD(tls_handshake_t, destroy, void,
	private_tls_server_t *this)
{
	DESTROY_IF(this->private);
	this->peer_auth->destroy(this->peer_auth);
	this->server_auth->destroy(this->server_auth);
	free(this);
}

/**
 * See header
 */
tls_server_t *tls_server_create(tls_t *tls,
						tls_crypto_t *crypto, tls_alert_t *alert,
						identification_t *server, identification_t *peer)
{
	private_tls_server_t *this;

	INIT(this,
		.public = {
			.handshake = {
				.process = _process,
				.build = _build,
				.cipherspec_changed = _cipherspec_changed,
				.change_cipherspec = _change_cipherspec,
				.finished = _finished,
				.destroy = _destroy,
			},
		},
		.tls = tls,
		.crypto = crypto,
		.alert = alert,
		.server = server,
		.peer = peer,
		.state = STATE_INIT,
		.peer_auth = auth_cfg_create(),
		.server_auth = auth_cfg_create(),
	);

	return &this->public;
}