#include "tls.h"

#include "mem.h"

#include <string.h>

#include <openssl/bio.h>
#include <openssl/err.h>
#include <openssl/pem.h>
#include <openssl/ssl.h>
#include <openssl/x509v3.h>

typedef enum _direction_t
{
	k_direction_undetermined,
	k_direction_accept,
	k_direction_connect,
} direction_t;

typedef struct _tf_tls_context_t
{
	SSL_CTX* context;
} tf_tls_context_t;

typedef struct _tf_tls_session_t
{
	tf_tls_context_t* context;
	BIO* bio_in;
	BIO* bio_out;
	SSL* ssl;
	const char* hostname;
	direction_t direction;
} tf_tls_session_t;

tf_tls_context_t* tf_tls_context_create()
{
	tf_tls_context_t* context = tf_malloc(sizeof(tf_tls_context_t));
	memset(context, 0, sizeof(*context));
	OPENSSL_init_ssl(0, NULL);
	context->context = SSL_CTX_new(SSLv23_method());
	SSL_CTX_set_default_verify_paths(context->context);
	return context;
}

bool tf_tls_context_set_certificate(tf_tls_context_t* context, const char* certificate)
{
	int result = 0;
	BIO* bio = BIO_new(BIO_s_mem());
	BIO_puts(bio, certificate);
	X509* x509 = PEM_read_bio_X509(bio, 0, 0, 0);
	result = SSL_CTX_use_certificate(context->context, x509);
	X509_free(x509);
	while (true)
	{
		x509 = PEM_read_bio_X509(bio, 0, 0, 0);
		if (x509)
		{
			SSL_CTX_add_extra_chain_cert(context->context, x509);
			/* Docs say don't x509_free: https://www.openssl.org/docs/man3.2/man3/SSL_CTX_add_extra_chain_cert.html. */
		}
		else
		{
			break;
		}
	}
	BIO_free(bio);
	return result == 1;
}

bool tf_tls_context_set_private_key(tf_tls_context_t* context, const char* private_key)
{
	int result = 0;
	BIO* bio = BIO_new(BIO_s_mem());
	BIO_puts(bio, private_key);
	EVP_PKEY* key = PEM_read_bio_PrivateKey(bio, 0, 0, 0);
	result = SSL_CTX_use_PrivateKey(context->context, key);
	EVP_PKEY_free(key);
	BIO_free(bio);
	return result == 1;
}

bool tf_tls_context_add_trusted_certificate(tf_tls_context_t* context, const char* certificate)
{
	bool result = false;
	BIO* bio = BIO_new_mem_buf(certificate, -1);
	X509* x509 = PEM_read_bio_X509(bio, 0, 0, 0);
	BIO_free(bio);

	if (x509)
	{
		X509_STORE* store = SSL_CTX_get_cert_store(context->context);
		if (store && X509_STORE_add_cert(store, x509) == 1)
		{
			result = true;
		}
		X509_free(x509);
	}

	return result;
}

tf_tls_session_t* tf_tls_context_create_session(tf_tls_context_t* context)
{
	tf_tls_session_t* session = tf_malloc(sizeof(tf_tls_session_t));
	memset(session, 0, sizeof(*session));
	session->context = context;
	session->bio_in = BIO_new(BIO_s_mem());
	session->bio_out = BIO_new(BIO_s_mem());
	return session;
}

void tf_tls_context_destroy(tf_tls_context_t* context)
{
	SSL_CTX_free(context->context);
	OPENSSL_cleanup();
	tf_free(context);
}

void tf_tls_session_destroy(tf_tls_session_t* session)
{
	if (session->ssl)
	{
		SSL_free(session->ssl);
	}
	if (session->hostname)
	{
		tf_free((void*)session->hostname);
	}
	tf_free(session);
}

void tf_tls_session_set_hostname(tf_tls_session_t* session, const char* hostname)
{
	if (session->hostname)
	{
		tf_free((void*)session->hostname);
		session->hostname = NULL;
	}
	if (hostname)
	{
		session->hostname = tf_strdup(hostname);
	}
}

void tf_tls_session_start_accept(tf_tls_session_t* session)
{
	session->direction = k_direction_accept;
	session->ssl = SSL_new(session->context->context);
	SSL_set_bio(session->ssl, session->bio_in, session->bio_out);
	SSL_accept(session->ssl);
	tf_tls_session_handshake(session);
}

void tf_tls_session_start_connect(tf_tls_session_t* session)
{
	session->direction = k_direction_connect;
	session->ssl = SSL_new(session->context->context);
	X509_VERIFY_PARAM* param = SSL_get0_param(session->ssl);
	X509_VERIFY_PARAM_set_hostflags(param, X509_CHECK_FLAG_NO_PARTIAL_WILDCARDS);
	X509_VERIFY_PARAM_set1_host(param, session->hostname, 0);
	SSL_set_tlsext_host_name(session->ssl, session->hostname);
	SSL_set_bio(session->ssl, session->bio_in, session->bio_out);
	SSL_connect(session->ssl);
	tf_tls_session_handshake(session);
}

void tf_tls_session_shutdown(tf_tls_session_t* session)
{
	SSL_shutdown(session->ssl);
}

int tf_tls_session_get_peer_certificate(tf_tls_session_t* session, char* buffer, size_t bytes)
{
	int result = -1;
#if OPENSSL_VERSION_NUMBER < 0x30000000L
	X509* certificate = SSL_get_peer_certificate(session->ssl);
#else
	X509* certificate = SSL_get1_peer_certificate(session->ssl);
#endif
	BIO* bio = BIO_new(BIO_s_mem());
	PEM_write_bio_X509(bio, certificate);
	X509_free(certificate);
	BUF_MEM* mem;
	BIO_get_mem_ptr(bio, &mem);
	if (mem->length <= bytes)
	{
		memcpy(buffer, mem->data, mem->length);
		result = mem->length;
	}
	BIO_free(bio);
	return result;
}

#if OPENSSL_VERSION_NUMBER < 0x10100000L
static bool _tls_session_wildcard_match(const char* pattern, size_t pattern_length, const char* name)
{
	const char* it = pattern;
	while (it - pattern < pattern_length && *name)
	{
		if (*it == '*')
		{
			for (const char* p = name; *p; ++p)
			{
				if (_tls_session_wildcard_match(it + 1, pattern_length - 1, p))
				{
					return true;
				}
			}
			return false;
		}
		else if (tolower(*it) == tolower(*name))
		{
			++it;
			++name;
		}
		else
		{
			break;
		}
	}
	return it - pattern <= pattern_length && *name == 0;
}

static bool _tls_session_verify_hostname(X509* certificate, const char* hostname)
{
	bool verified = false;
	void* names = X509_get_ext_d2i(certificate, NID_subject_alt_name, 0, 0);
	if (names)
	{
		int count = sk_GENERAL_NAME_num(names);
		for (int i = 0; i < count; ++i)
		{
			const GENERAL_NAME* check = sk_GENERAL_NAME_value(names, i);
			if (!verified)
			{
#if OPENSSL_VERSION_NUMBER <= 0x1000211fL
				const unsigned char* name = ASN1_STRING_data(check->d.ia5);
#else
				const char* name = ASN1_STRING_get0_data(check->d.ia5);
#endif
				size_t length = ASN1_STRING_length(check->d.ia5);
				if (_tls_session_wildcard_match((const char*)name, length, hostname))
				{
					verified = true;
				}
			}
		}
		sk_GENERAL_NAMES_free(names);
	}

	if (!verified)
	{
		int index = X509_NAME_get_index_by_NID(X509_get_subject_name(certificate), NID_commonName, -1);
		if (index >= 0)
		{
			X509_NAME_ENTRY* entry = X509_NAME_get_entry(X509_get_subject_name(certificate), index);
			if (entry)
			{
				ASN1_STRING* asn1 = X509_NAME_ENTRY_get_data(entry);
				if (asn1)
				{
#if OPENSSL_VERSION_NUMBER <= 0x1000211fL
					const unsigned char* commonName = ASN1_STRING_data(asn1);
#else
					const char* commonName = ASN1_STRING_get0_data(asn1);
#endif
					if ((size_t)(ASN1_STRING_length(asn1)) == strlen((const char*)commonName))
					{
						verified = _tls_session_wildcard_match((const char*)commonName, ASN1_STRING_length(asn1), hostname);
					}
				}
			}
		}
	}

	return verified;
}
#endif

static bool _tls_session_verify_peer_certificate(tf_tls_session_t* session)
{
	bool verified = false;
#if OPENSSL_VERSION_NUMBER < 0x30000000L
	X509* certificate = SSL_get_peer_certificate(session->ssl);
#else
	X509* certificate = SSL_get1_peer_certificate(session->ssl);
#endif
	if (certificate)
	{
		if (SSL_get_verify_result(session->ssl) == X509_V_OK)
		{
#if OPENSSL_VERSION_NUMBER < 0x10100000L
			if (_tls_session_verify_hostname(certificate, session->hostname))
			{
				verified = true;
			}
#else
			verified = true;
#endif
		}
		X509_free(certificate);
	}
	return verified;
}

tf_tls_handshake_t tf_tls_session_handshake(tf_tls_session_t* session)
{
	tf_tls_handshake_t result = k_tls_handshake_done;
	if (!SSL_is_init_finished(session->ssl))
	{
		int value = SSL_do_handshake(session->ssl);
		if (value <= 0)
		{
			int error = SSL_get_error(session->ssl, value);
			if (error != SSL_ERROR_WANT_READ && error != SSL_ERROR_WANT_WRITE)
			{
				result = k_tls_handshake_failed;
			}
			else
			{
				result = k_tls_handshake_more;
			}
		}
	}
	if (result == k_tls_handshake_done && session->direction == k_direction_connect && !_tls_session_verify_peer_certificate(session))
	{
		result = k_tls_handshake_failed;
	}
	return result;
}

int tf_tls_session_read_plain(tf_tls_session_t* session, char* buffer, size_t bytes)
{
	int result = SSL_read(session->ssl, buffer, bytes);
	if (result <= 0)
	{
		int error = SSL_get_error(session->ssl, result);
		if (error == SSL_ERROR_WANT_READ || error == SSL_ERROR_WANT_WRITE)
		{
			result = 0;
		}
		else if (error == SSL_ERROR_ZERO_RETURN)
		{
			if ((SSL_get_shutdown(session->ssl) & SSL_RECEIVED_SHUTDOWN) != 0)
			{
				result = k_tls_read_zero;
			}
			else
			{
				result = 0;
			}
		}
		else
		{
			result = k_tls_read_failed;
		}
	}
	return result;
}

int tf_tls_session_write_plain(tf_tls_session_t* session, const char* buffer, size_t bytes)
{
	return SSL_write(session->ssl, buffer, bytes);
}

int tf_tls_session_read_encrypted(tf_tls_session_t* session, char* buffer, size_t bytes)
{
	return BIO_read(session->bio_out, buffer, bytes);
}

int tf_tls_session_write_encrypted(tf_tls_session_t* session, const char* buffer, size_t bytes)
{
	return BIO_write(session->bio_in, buffer, bytes);
}

bool tf_tls_session_get_error(tf_tls_session_t* session, char* buffer, size_t bytes)
{
	unsigned long error = ERR_get_error();
	if (error != 0)
	{
		ERR_error_string_n(error, buffer, bytes);
	}
	return error != 0;
}