#define _POSIX_C_SOURCE 200809L

#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>

#include "tlse.h"
#include "chacha.h"
#include "buffer.h"

static int tls_cipher_is_fs(struct TLSContext *context, unsigned short cipher) {
	if (!context) {
		return 0;
	}

	if (context->tlsver == TLS_VERSION13) {
		switch (cipher) {
			case TLS_AES_128_GCM_SHA256:
			case TLS_AES_256_GCM_SHA384:
			case TLS_CHACHA20_POLY1305_SHA256:
				return 1;
		}
		return 0;
	}

	switch (cipher) {
		case TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:
		case TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:
		case TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256:
			if (context && context->certificates
					&& context->certificates_count
					&& context->ec_private_key) {
				return 1;
			}
			return 0;
		case TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256:
		case TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA384:
		case TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256:
		case TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384:
			if (context->version == TLS_V12
					|| context->version == DTLS_V12) {
				if (context && context->certificates
						&& context->certificates_count
						&& context->ec_private_key) {
					return 1;
				}
			}
		return 0;
	case TLS_DHE_RSA_WITH_AES_128_CBC_SHA:
	case TLS_DHE_RSA_WITH_AES_256_CBC_SHA:
	case TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:
	case TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:
		return 1;
	case TLS_DHE_RSA_WITH_AES_128_CBC_SHA256:
	case TLS_DHE_RSA_WITH_AES_256_CBC_SHA256:
	case TLS_DHE_RSA_WITH_AES_128_GCM_SHA256:
	case TLS_DHE_RSA_WITH_AES_256_GCM_SHA384:
	case TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:
	case TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256:
	case TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384:
	case TLS_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256:
	case TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256:
		if (context->tlsver == TLS_VERSION12) {
			return 1;
		}
		break;
	}
	return 0;
}

static uint16_t get16(const unsigned char *buf) {
	uint16_t res;

	res = ((*buf) << 8) + (*(buf+1));
	return res;
}

int tls_choose_cipher(struct TLSContext *context, const unsigned char *buf,
		      int buf_len, int *scsv_set) {
	int i;
	if (scsv_set) {
		*scsv_set = 0;
	}
	if (!context) {
		return 0;
	}
	int selected_cipher = TLS_NO_COMMON_CIPHER;

	if (selected_cipher == TLS_NO_COMMON_CIPHER) {
		for (i = 0; i < buf_len; i += 2) {
			uint16_t cipher = get16(&buf[i]);
			if (tls_cipher_is_fs(context, cipher)) {
				selected_cipher = cipher;
				break;
			}
		}
	}

	for (i = 0; i < buf_len; i += 2) {
		uint16_t cipher = get16(&buf[i]);
		if (cipher == TLS_FALLBACK_SCSV) {
			if (scsv_set) {
				*scsv_set = 1;
			}
			if (selected_cipher != TLS_NO_COMMON_CIPHER) {
				break;
			}
		}
	}
	return selected_cipher;
}

int tls_parse_client_hello(struct TLSContext *context,
	       	const unsigned char *buf, int buf_len,
	       	unsigned int *write_packets) {
	*write_packets = 0;
	if (context->connection_status != 0 && context->connection_status != 4) {
		DEBUG_PRINT("UNEXPECTED HELLO MESSAGE\n");
		return TLS_UNEXPECTED_MESSAGE;
	}

	int res = 0;
	int downgraded = 0;

	int hello_min_size = TLS_CLIENT_HELLO_MINSIZE;

	if (buf_len < hello_min_size) {
		return TLS_NEED_MORE_DATA;
	}

	/* big endian */
	int bytes_to_follow = buf[0] * 0x10000 + buf[1] * 0x100 + buf[2];
	res += 3;
	if (buf_len - res < bytes_to_follow) {
		return TLS_NEED_MORE_DATA;
	}

	if (buf_len - res < 2) {
		return TLS_NEED_MORE_DATA;
	}

	unsigned short version = get16(&buf[res]);

	res += 2;
	if (!tls_supported_version(version)) {
		return TLS_NOT_SAFE;
	}
	DEBUG_PRINT("VERSION REQUIRED BY REMOTE %x, VERSION NOW %x\n",
			(int) version, (int) context->version);
	memcpy(context->remote_random, &buf[res], TLS_CLIENT_RANDOM_SIZE);
	res += TLS_CLIENT_RANDOM_SIZE;

	unsigned char session_len = buf[res++];
	if (buf_len - res < session_len) return TLS_NEED_MORE_DATA;

	if (session_len && session_len <= TLS_MAX_SESSION_ID) {
		memcpy(context->session, &buf[res], session_len);
		context->session_size = session_len;
		DEBUG_DUMP_HEX_LABEL("REMOTE SESSION ID: ",
				context->session,
				context->session_size);
	} else {
		context->session_size = 0;
	}
	res += session_len;

	const unsigned char *cipher_buffer = NULL;
	unsigned short cipher_len = 0;
	int scsv_set = 0;
	if (context->is_server) {
		if (buf_len - res < 2) {
			return TLS_NEED_MORE_DATA;
		}

		cipher_len = get16(&buf[res]);
		res += 2;
		if (buf_len - res < cipher_len) return TLS_NEED_MORE_DATA;
		/* faster than cipher_len % 2 */
		/* TODO unlikely, let the compiler worry about it */
		if (cipher_len & 1) {
			return TLS_BROKEN_PACKET;
		}

		cipher_buffer = &buf[res];
		res += cipher_len;

		if (buf_len - res < 1) {
			return TLS_NEED_MORE_DATA;
		}

		unsigned char compression_list_size = buf[res++];
		if (buf_len - res < compression_list_size) {
			return TLS_NEED_MORE_DATA;
		}
		/* no compression support */
		res += compression_list_size;
	} else {
		/* client */
		if (buf_len - res < 2) {
			return TLS_NEED_MORE_DATA;
		}

		unsigned short cipher = get16(&buf[res]);
		res += 2;
		context->cipher = cipher;
		if (!tls_cipher_supported(context, cipher)) {
			context->cipher = 0;
			DEBUG_PRINT("NO CIPHER SUPPORTED\n");
			return TLS_NO_COMMON_CIPHER;
		}
		DEBUG_PRINT("CIPHER: %s\n", tls_cipher_name(context));
		if (buf_len - res < 1) return TLS_NEED_MORE_DATA;
		unsigned char compression = buf[res++];
		if (compression != 0) {
			DEBUG_PRINT("COMPRESSION NOT SUPPORTED\n");
			return TLS_COMPRESSION_NOT_SUPPORTED;
		}
	}

	if (res > 0) {
		if (context->is_server) {
			*write_packets = 2;
		}
		if (context->connection_status != 4) {
			context->connection_status = 1;
		}
	}


	if (res > 2) {
		res += 2;
	}
	const unsigned char *key_share = NULL;
	unsigned short key_size = 0;
	while (buf_len - res >= 4) {
		/* have extensions */
		unsigned short extension_type = get16(&buf[res]);
		res += 2;
		unsigned short extension_len = get16(&buf[res]);
		res += 2;
		DEBUG_PRINT("Extension: 0x0%x (%i), len: %i\n",
			    (int) extension_type, (int) extension_type,
			    (int) extension_len);
		if (extension_len) {
			/* SNI extension */
			if (buf_len - res < extension_len) {
				return TLS_NEED_MORE_DATA;
			}
			if (extension_type == 0x00) {
				/* unsigned char sni_type = buf[res + 2]; */
				uint16_t sni_host_len = get16(&buf[res+3]);
				if (buf_len - res - 5 < sni_host_len) {
					return TLS_NEED_MORE_DATA;
				}

				if (sni_host_len) {
					free(context->sni);
					context->sni = malloc(sni_host_len + 1);
					if (context->sni) {
						memcpy(context->sni,
								&buf[res + 5],
								sni_host_len);
						context->sni[sni_host_len] = 0;
						DEBUG_PRINT("SNI HOST INDICATOR: [%s]\n", context->sni);
					}
				}
			} else if (extension_type == 0x0A) {
				/* supported groups */
				if (buf_len - res > 2) {
					uint16_t group_len = get16(&buf[res]);
					if (buf_len - res >= group_len + 2) {
						DEBUG_DUMP_HEX_LABEL
						    ("SUPPORTED GROUPS",
						     &buf[res + 2],
						     group_len);
						int i;
						int selected = 0;
						for (i = 0; i < group_len;
						     i += 2) {
							uint16_t iana_n = get16(&buf[res + 2 + i]);
							switch (iana_n) {
							case 23:
								context->
								    curve =
								    &secp256r1;
								selected =
								    1;
								break;
							case 24:
								context->
								    curve =
								    &secp384r1;
								selected =
								    1;
								break;
#if 0
								/* needs different implementation */
								case 29:
								context->curve = &curve25519;
								selected = 1;
								break;
#endif
#if 0
								/* do not use it anymore */
								case 25:
								context->curve = &secp521r1;
								selected = 1;
								break;
#endif
							}
							if (selected) {
								DEBUG_PRINT
								    ("SELECTED CURVE %s\n",
								     context->
								     curve->
								     name);
								break;
							}
						}
					}
				}
			} else
				if (extension_type == 0x10 && context->alpn &&
						context->alpn_count) {
				if (buf_len - res > 2) {
					uint16_t alpn_len = get16(&buf[res]);
					if (alpn_len && alpn_len <=
							extension_len - 2) {
						unsigned char *alpn =
						    (unsigned char *)
						    &buf[res + 2];
						int alpn_pos = 0;
						while (alpn_pos < alpn_len) {
							unsigned char
							 alpn_size =
							    alpn
							    [alpn_pos++];
							if (alpn_size +
							    alpn_pos >=
							    extension_len)
								break;
							if ((alpn_size)
							    &&
							    (tls_alpn_contains
							     (context,
							      (char *)
							      &alpn
							      [alpn_pos],
							      alpn_size)))
							{
								free
								    (context->
								     negotiated_alpn);
								context->
								    negotiated_alpn
								    =
								    malloc
								    (alpn_size
								     + 1);
								if (context->negotiated_alpn) {
									memcpy
									    (context->
									     negotiated_alpn,
									     &alpn
									     [alpn_pos],
									     alpn_size);
									context->
									    negotiated_alpn
									    [alpn_size]
									    =
									    0;
									DEBUG_PRINT
									    ("NEGOTIATED ALPN: %s\n",
									     context->
									     negotiated_alpn);
								}
								break;
							}
							alpn_pos +=
							    alpn_size;
							/* ServerHello contains just one alpn */
							if (!context->
							    is_server)
								break;
						}
					}
				}
			} else if (extension_type == 0x0D) {
				/* supported signatures */
				DEBUG_DUMP_HEX_LABEL
				    ("SUPPORTED SIGNATURES", &buf[res],
				     extension_len);
			} else if (extension_type == 0x0B) {
				/* supported point formats */
				DEBUG_DUMP_HEX_LABEL
				    ("SUPPORTED POINT FORMATS", &buf[res],
				     extension_len);
			}
			else if (extension_type == 0x2B) {
				/* supported versions */
				if ((buf[res] == extension_len - 1)
				    && (extension_len > 4)) {
					DEBUG_DUMP_HEX_LABEL
					    ("SUPPORTED VERSIONS",
					     &buf[res], extension_len);
					/* tls 1.3 draft version 28 */
					int i;
					int limit = (int) buf[res];
					if (limit == extension_len - 1) {
						limit--;
						for (i = 1; i < limit;
						     i += 2) {
							if ((get16(&buf[res + i]) == 0x7F1C)
							    ||
							    (get16(&buf[res + i]) == TLS_V13)) {
								context->
								    version
								    =
								    TLS_V13;
								context->
								    tls13_version
								    =
								    get16(&buf[res + i]);
								DEBUG_PRINT
								    ("TLS 1.3 SUPPORTED\n");
								break;
							}
						}
					}
				}
			} else if (extension_type == 0x2A) {
				/* early data */
				DEBUG_DUMP_HEX_LABEL
				    ("EXTENSION, EARLY DATA", &buf[res],
				     extension_len);
			} else if (extension_type == 0x29) {
				/* pre shared key */
				DEBUG_DUMP_HEX_LABEL
				    ("EXTENSION, PRE SHARED KEY",
				     &buf[res], extension_len);
			} else if (extension_type == 0x33) {
				/* key share */
				key_size = get16(&buf[res]);
				if (key_size > extension_len - 2) {
					DEBUG_PRINT("BROKEN KEY SHARE\n");
					return TLS_BROKEN_PACKET;
				}
				DEBUG_DUMP_HEX_LABEL
				    ("EXTENSION, KEY SHARE", &buf[res],
				     extension_len);
				key_share = &buf[res + 2];
			} else if (extension_type == 0x0D) {
				/* signature algorithms */
				DEBUG_DUMP_HEX_LABEL
				    ("EXTENSION, SIGNATURE ALGORITHMS",
				     &buf[res], extension_len);
			} else if (extension_type == 0x2D) {
				/* psk key exchange modes */
				DEBUG_DUMP_HEX_LABEL
				    ("EXTENSION, PSK KEY EXCHANGE MODES",
				     &buf[res], extension_len);
			}
			res += extension_len;
		}
	}

	if (buf_len != res) {
		return TLS_NEED_MORE_DATA;
	}

	if (context->is_server && cipher_buffer && cipher_len) {
		int cipher =
		    tls_choose_cipher(context, cipher_buffer, cipher_len,
				      &scsv_set);
		if (cipher < 0) {
			DEBUG_PRINT("NO COMMON CIPHERS\n");
			return cipher;
		}
		if (downgraded && scsv_set) {
			DEBUG_PRINT("NO DOWNGRADE (SCSV SET)\n");
			tls_alert(context, 1, inappropriate_fallback);
			context->critical_error = 1;
			return TLS_NOT_SAFE;
		}
		context->cipher = cipher;
	}
	if (key_share && key_size && context->tlsver == TLS_VERSION13) {
		int key_share_err =
		    tls_parse_key_share(context, key_share, key_size);
		if (key_share_err) {
			/* request hello retry */
			if (context->connection_status != 4) {
				*write_packets = 5;
				context->hs_messages[1] = 0;
				context->connection_status = 4;
				return res;
			} else
				return key_share_err;
		}
		/* we have key share */
		context->connection_status = 3;
	}
	return res;
}