#include <sys/types.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <limits.h>
#include <assert.h>
#include <ssl_utils.h>

#include <stdio.h>
#include <stdlib.h>

#define ALLOC(type, number) ((type *)calloc(sizeof(type), number))
#define	FREE(p)	{free(*p);*p=NULL;}
#define TFE_SYMBOL_MAX	64

void ssl_chello_free(struct ssl_chello* chello)
{
	if(chello == NULL)
	{
		return;
	}
	FREE(&(chello->sni));
	chello->sni = NULL;
	FREE(&(chello->alpn));
	chello->alpn = NULL;
	FREE(&(chello->sign_algos));
	chello->sign_algos = NULL;
	FREE(&(chello->supported_groups));
	chello->supported_groups = NULL;
	FREE(&(chello->cipher_suites));
	chello->cipher_suites = NULL;
	FREE(&chello);
}
static int cipher_is_grease(uint16_t cipher)
{
	uint16_t a=cipher>>8;
	uint16_t b=cipher&0x00ff;
	//https://tools.ietf.org/html/draft-davidben-tls-grease-01#section-5
	if(a==b && (a&0x0f)==0x0a)
	{
		return 1;
	}
	else
	{
		return 0;
	}
}

static int  parse_server_name_extension(const unsigned char *buff, uint16_t buff_len, struct ssl_chello* chello)
{
    size_t pos = 2; /* skip server name list length */
    uint16_t len;
	char *sni = NULL;
    while (pos + 3 < buff_len)
	{
        len = (buff[pos + 1] << 8) | buff[pos + 2];
        if (pos + 3 + len > buff_len)
		{
			chello->sni = sni;
			return CHELLO_PARSE_INVALID_FORMAT;
		}
        switch (buff[pos])
		{
            case 0x00: /* host_name */
				sni = ALLOC(char, len + 1);
                strncpy(sni, (const char*)buff + pos + 3, len);
                sni[len] = '\0';
				chello->sni = sni;
				return CHELLO_PARSE_SUCCESS;
        }
        pos += 3 + len;
    }
	chello->sni = sni;
    if (pos != buff_len)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	return CHELLO_PARSE_SUCCESS;
}

static int parse_alpn_extension(const unsigned char* buff, uint16_t buff_len, struct ssl_chello* chello)
{
	int pos = 0;
	uint16_t len = (buff[pos] << 8) | buff[pos + 1];
	char *alpn = NULL;
	if(2 + len != buff_len)
	{
		chello->alpn = alpn;
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	alpn = ALLOC(char, len + 1);
	strncpy((char*)alpn, (const char*)buff + 2, len);
	alpn[len] = '\0';
	chello->alpn = alpn;
	return CHELLO_PARSE_SUCCESS;
}

static int parse_supported_versions_extension(const unsigned char* buff, uint16_t buff_len, struct ssl_chello* chello)
{
	if(buff_len < 1)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	uint8_t len = buff[0];
	if(len + 1 > buff_len || len % 2 != 0)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	uint16_t max_version = 0x0000, min_version = 0xffff;
	for(int i = 1; i < len; i+=2)
	{
		uint16_t version = (buff[i] << 8) | buff[i + 1];
		//unknown version
		if(version < 0x0300 || version > 0x0304)
		{
			continue;
		}
		if(version > max_version)
		{
			max_version = version;
		}
		if(version < min_version)
		{
			min_version = version;
		}
	}
	if(max_version != 0x0000)
	{
		chello->max_version.ossl_format = max_version;
		chello->max_version.major = (max_version >> 8);
		chello->max_version.minor = max_version & 0x00ff;
		chello->min_version.ossl_format = min_version;
		chello->min_version.major = (min_version >> 8);
		chello->min_version.minor = min_version & 0x00ff;
	}
	return CHELLO_PARSE_SUCCESS;
}

static int parse_sign_algos_extension(const unsigned char* buff, uint16_t buff_len, struct ssl_chello* chello)
{
	if(buff_len < 2)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	uint16_t len = (buff[0] << 8) | buff[1];
	if(len + 2 > buff_len)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	char *sign_algos = ALLOC(char, len);
	memcpy(sign_algos, (void*)(buff + 2), len);
	chello->sign_algos = sign_algos;
	chello->sign_algos_len = len;
	return CHELLO_PARSE_SUCCESS;
}

static int parse_supported_groups_extension(const unsigned char* buff, uint16_t buff_len, struct ssl_chello* chello)
{
	if(buff_len < 2)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	uint16_t len = (buff[0] << 8) | buff[1];
	if(len + 2 > buff_len)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	chello->supported_groups = ALLOC(char, len);
	uint16_t* known_groups = (uint16_t*) chello->supported_groups;
	uint16_t* raw_groups= (uint16_t*) (buff + 2);
	size_t i=0, j=0;
	for(i=0; i<len/2; i++)
	{
		if(!cipher_is_grease(raw_groups[i]))
		{
			known_groups[j++]=raw_groups[i];
		}
	}
	chello->supported_groups_len = j*2;
	return CHELLO_PARSE_SUCCESS;
}

static int parse_extensions(const unsigned char *buff, uint16_t buff_len, struct ssl_chello *chello) {
    uint16_t pos = 0;
    /* Parse each 4 bytes for the extension header */
    while (pos + 4 <= buff_len)
	{
        uint16_t len = (buff[pos + 2] << 8) | buff[pos + 3];
		if (pos + 4 + len > buff_len)
		{
			return CHELLO_PARSE_INVALID_FORMAT;
		}
		uint16_t extension_type = (buff[pos] << 8) | buff[pos + 1];
		int result = CHELLO_PARSE_SUCCESS;
		switch(extension_type){
			//server name extension
			case 0x0000:
				result = parse_server_name_extension(buff + pos + 4, len, chello);
				break;
			//alpn extension
			case 0x0010:
				result = parse_alpn_extension(buff + pos + 4, len, chello);
				break;
			//supported versions extension
			case 0x002b:
				result = parse_supported_versions_extension(buff + pos + 4, len, chello);
				break;
			//signature algorithms extension
			case 0x000d:
				result = parse_sign_algos_extension(buff + pos + 4, len, chello);
				break;
			//supported groups extensions
			case 0x000a:
				result = parse_supported_groups_extension(buff + pos + 4, len, chello);
				break;
			default:
				break;
		}
		if(result != CHELLO_PARSE_SUCCESS)
		{
			return result;
		}
        pos += (4 + len);
    }
    /* Check we ended where we expected to */
    if (pos != buff_len)
	{
		return CHELLO_PARSE_INVALID_FORMAT;
	}
	return CHELLO_PARSE_SUCCESS;
}


struct cipher_suite
{
	uint16_t value;
	const char* name;
};

static struct cipher_suite cipher_suite_list[] =
{
	{0xC030, "ECDHE-RSA-AES256-GCM-SHA384"},
	{0xC02C, "ECDHE-ECDSA-AES256-GCM-SHA384"},
	{0xC028, "ECDHE-RSA-AES256-SHA384"},
	{0xC024, "ECDHE-ECDSA-AES256-SHA384"},
	{0xC014, "ECDHE-RSA-AES256-SHA"},
	{0xC00A, "ECDHE-ECDSA-AES256-SHA"},
	{0x00A5, "DH-DSS-AES256-GCM-SHA384"},
	{0x00A3, "DHE-DSS-AES256-GCM-SHA384"},
	{0x00A1, "DH-RSA-AES256-GCM-SHA384"},
	{0x009F, "DHE-RSA-AES256-GCM-SHA384"},
	{0x006B, "DHE-RSA-AES256-SHA256"},
	{0x006A, "DHE-DSS-AES256-SHA256"},
	{0x0069, "DH-RSA-AES256-SHA256"},
	{0x0068, "DH-DSS-AES256-SHA256"},
	{0x0039, "DHE-RSA-AES256-SHA"},
	{0x0038, "DHE-DSS-AES256-SHA"},
	{0x0037, "DH-RSA-AES256-SHA"},
	{0x0036, "DH-DSS-AES256-SHA"},
	{0x0088, "DHE-RSA-CAMELLIA256-SHA"},
	{0x0087, "DHE-DSS-CAMELLIA256-SHA"},
	{0x0086, "DH-RSA-CAMELLIA256-SHA"},
	{0x0085, "DH-DSS-CAMELLIA256-SHA"},
	{0xC019, "AECDH-AES256-SHA"},
	{0x00A7, "ADH-AES256-GCM-SHA384"},
	{0x006D, "ADH-AES256-SHA256"},
	{0x003A, "ADH-AES256-SHA"},
	{0x0089, "ADH-CAMELLIA256-SHA"},
	{0xC032, "ECDH-RSA-AES256-GCM-SHA384"},
	{0xC02E, "ECDH-ECDSA-AES256-GCM-SHA384"},
	{0xC02A, "ECDH-RSA-AES256-SHA384"},
	{0xC026, "ECDH-ECDSA-AES256-SHA384"},
	{0xC00F, "ECDH-RSA-AES256-SHA"},
	{0xC005, "ECDH-ECDSA-AES256-SHA"},
	{0x009D, "AES256-GCM-SHA384"},
	{0x003D, "AES256-SHA256"},
	{0x0035, "AES256-SHA"},
	{0x0084, "CAMELLIA256-SHA"},
	{0x008D, "PSK-AES256-CBC-SHA"},
	{0xC02F, "ECDHE-RSA-AES128-GCM-SHA256"},
	{0xC02B, "ECDHE-ECDSA-AES128-GCM-SHA256"},
	{0xC027, "ECDHE-RSA-AES128-SHA256"},
	{0xC023, "ECDHE-ECDSA-AES128-SHA256"},
	{0xC013, "ECDHE-RSA-AES128-SHA"},
	{0xC009, "ECDHE-ECDSA-AES128-SHA"},
	{0x00A4, "DH-DSS-AES128-GCM-SHA256"},
	{0x00A2, "DHE-DSS-AES128-GCM-SHA256"},
	{0x00A0, "DH-RSA-AES128-GCM-SHA256"},
	{0x009E, "DHE-RSA-AES128-GCM-SHA256"},
	{0x0067, "DHE-RSA-AES128-SHA256"},
	{0x0040, "DHE-DSS-AES128-SHA256"},
	{0x003F, "DH-RSA-AES128-SHA256"},
	{0x003E, "DH-DSS-AES128-SHA256"},
	{0x0033, "DHE-RSA-AES128-SHA"},
	{0x0032, "DHE-DSS-AES128-SHA"},
	{0x0031, "DH-RSA-AES128-SHA"},
	{0x0030, "DH-DSS-AES128-SHA"},
	{0x009A, "DHE-RSA-SEED-SHA"},
	{0x0099, "DHE-DSS-SEED-SHA"},
	{0x0098, "DH-RSA-SEED-SHA"},
	{0x0097, "DH-DSS-SEED-SHA"},
	{0x0045, "DHE-RSA-CAMELLIA128-SHA"},
	{0x0044, "DHE-DSS-CAMELLIA128-SHA"},
	{0x0043, "DH-RSA-CAMELLIA128-SHA"},
	{0x0042, "DH-DSS-CAMELLIA128-SHA"},
	{0xC018, "AECDH-AES128-SHA"},
	{0x00A6, "ADH-AES128-GCM-SHA256"},
	{0x006C, "ADH-AES128-SHA256"},
	{0x0034, "ADH-AES128-SHA"},
	{0x009B, "ADH-SEED-SHA"},
	{0x0046, "ADH-CAMELLIA128-SHA"},
	{0xC031, "ECDH-RSA-AES128-GCM-SHA256"},
	{0xC02D, "ECDH-ECDSA-AES128-GCM-SHA256"},
	{0xC029, "ECDH-RSA-AES128-SHA256"},
	{0xC025, "ECDH-ECDSA-AES128-SHA256"},
	{0xC00E, "ECDH-RSA-AES128-SHA"},
	{0xC004, "ECDH-ECDSA-AES128-SHA"},
	{0x009C, "AES128-GCM-SHA256"},
	{0x003C, "AES128-SHA256"},
	{0x002F, "AES128-SHA"},
	{0x0096, "SEED-SHA"},
	{0x0041, "CAMELLIA128-SHA"},
	{0x008C, "PSK-AES128-CBC-SHA"},
	{0xC012, "ECDHE-RSA-DES-CBC3-SHA"},
	{0xC008, "ECDHE-ECDSA-DES-CBC3-SHA"},
	{0x0016, "EDH-RSA-DES-CBC3-SHA"},
	{0x0013, "EDH-DSS-DES-CBC3-SHA"},
	{0x0010, "DH-RSA-DES-CBC3-SHA"},
	{0x000D, "DH-DSS-DES-CBC3-SHA"},
	{0xC017, "AECDH-DES-CBC3-SHA"},
	{0x001B, "ADH-DES-CBC3-SHA"},
	{0xC00D, "ECDH-RSA-DES-CBC3-SHA"},
	{0xC003, "ECDH-ECDSA-DES-CBC3-SHA"},
	{0x000A, "DES-CBC3-SHA"},
	{0x0007, "IDEA-CBC-SHA"},
	{0x008B, "PSK-3DES-EDE-CBC-SHA"},
	{0x0021, "KRB5-IDEA-CBC-SHA"},
	{0x001F, "KRB5-DES-CBC3-SHA"},
	{0x0025, "KRB5-IDEA-CBC-MD5"},
	{0x0023, "KRB5-DES-CBC3-MD5"},
	{0xC011, "ECDHE-RSA-RC4-SHA"},
	{0xC007, "ECDHE-ECDSA-RC4-SHA"},
	{0xC016, "AECDH-RC4-SHA"},
	{0x0018, "ADH-RC4-MD5"},
	{0xC00C, "ECDH-RSA-RC4-SHA"},
	{0xC002, "ECDH-ECDSA-RC4-SHA"},
	{0x0005, "RC4-SHA"},
	{0x0004, "RC4-MD5"},
	{0x008A, "PSK-RC4-SHA"},
	{0x0020, "KRB5-RC4-SHA"},
	{0x0024, "KRB5-RC4-MD5"},
	{0xC010, "ECDHE-RSA-NULL-SHA"},
	{0xC006, "ECDHE-ECDSA-NULL-SHA"},
	{0xC015, "AECDH-NULL-SHA"},
	{0xC00B, "ECDH-RSA-NULL-SHA"},
	{0xC001, "ECDH-ECDSA-NULL-SHA"},
	{0x003B, "NULL-SHA256"},
	{0x0002, "NULL-SHA"},
	{0x0001, "NULL-MD5"}
};

static struct cipher_suite cipher_suite_list_tls13[] =
{
	{0x1301, "TLS_AES_128_GCM_SHA256"},
	{0x1302, "TLS_AES_256_GCM_SHA384"},
	{0x1303, "TLS_CHACHA20_POLY1305_SHA256"},
	{0x1304, "TLS_AES_128_CCM_SHA256"},
	{0x1305, "TLS_AES_128_CCM_8_SHA256"}
};

static int cipher_suites_convert_helper(uint16_t value, char *name, size_t name_sz)
{
	int n1 = sizeof(cipher_suite_list) / sizeof(struct cipher_suite);
	int n2 = sizeof(cipher_suite_list_tls13) / sizeof(struct cipher_suite);
	for(int i = 0; i < n1; i++)
	{
		if(value == cipher_suite_list[i].value)
		{
			if(name) memcpy(name, cipher_suite_list[i].name, strnlen(cipher_suite_list[i].name, name_sz));
			return 1;
		}
	}
	for(int i = 0; i < n2; i++)
	{
		if(value == cipher_suite_list_tls13[i].value)
		{
			if(name) memcpy(name, cipher_suite_list_tls13[i].name, strnlen(cipher_suite_list_tls13[i].name, name_sz));
			return 2;
		}
	}
	return 0;
}

void ssl_cipher_suites_to_name(const char *source, int source_len, char *result_common, size_t sz_common, char *result_tls13, size_t sz_tls13)
{
	int target_common_reach_max = 0;
	int target_tls13_reach_max = 0;
	char name[TFE_SYMBOL_MAX] = "";
	for(int i = 0; i < source_len - 1;)
	{
		uint16_t val = (source[i] << 8) | source[i + 1];
		memset(name, 0, sizeof(name));
		int ret = cipher_suites_convert_helper(val, name, sizeof(name));
		//target common
		if(ret == 1 && target_common_reach_max == 0)
		{
			if(strnlen(name, sizeof(name)) + strnlen(result_common, sz_common) + 1 > sz_common)
			{
				target_common_reach_max = 1;
			}
			else
			{
				strncat(result_common, name, sz_common);
				strncat(result_common, ":", sz_common);
			}
		}
		//result_tls13
		if(ret == 2 && target_tls13_reach_max == 0)
		{
			if(strnlen(name, sizeof(name)) + strnlen(result_tls13, sz_tls13) + 1 > sz_tls13)
			{
				target_tls13_reach_max = 1;
			}
			else
			{
				strncat(result_tls13, name, sz_tls13);
				strncat(result_tls13, ":", sz_tls13);
			}
		}
		i += 2;
	}
	int len1 = strnlen(result_common, sz_common);
	if(len1 > 0)
	{
		result_common[len1 - 1] = '\0';
	}
	int len2 = strnlen(result_tls13, sz_tls13);
	if(len2 > 0)
	{
		result_tls13[len2 - 1] = '\0';
	}
	return;
}

struct ssl_chello* ssl_chello_parse(const unsigned char* buff, size_t buff_len, enum chello_parse_result* result)
{
	if(buff == NULL)
	{
		*result = CHELLO_PARSE_INVALID_FORMAT;
		return NULL;
	}
	if(buff_len < 1)
	{
		*result = CHELLO_PARSE_NOT_ENOUGH_BUFF;
		return NULL;
	}
	if(buff[0] != 0x80 && buff[0] != 0x16)
	{
		*result = CHELLO_PARSE_INVALID_FORMAT;
		return NULL;
	}
	 /* SSL 2.0 compatible Client Hello
     * High bit of first byte (length) and content type is Client Hello
     * See RFC5246 Appendix E.2
	 * if it is SSL 2.0, only parse version
     */
	if(buff[0] == 0x80)
	{
		struct ssl_chello* _chello = (struct ssl_chello*)ALLOC(struct ssl_chello, 1);
		_chello->min_version.major = 0x02;
		if(buff_len < 2)
		{
			*result = CHELLO_PARSE_NOT_ENOUGH_BUFF;
			return _chello;
		}
		size_t len = (size_t)buff[1];
		if (buff_len < len + 2)
		{
			*result = CHELLO_PARSE_NOT_ENOUGH_BUFF;
        	return _chello;
		}
		buff_len = len + 2;
		size_t pos = 2;
		/* Handshark Message Type: Client Hello */
		if (pos + 1 > buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
    	}
		if (buff[pos] != 0x01)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
    	}
		pos += 1;
		/*  Version */
		if(pos + 2 > buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
		_chello->max_version.major = buff[pos];
		_chello->max_version.minor = buff[pos + 1];
		_chello->max_version.ossl_format=(uint16_t)_chello->max_version.major<<8|_chello->max_version.minor;
		*result = CHELLO_PARSE_SUCCESS;
		return _chello;
	}
	else
	{
		if (buff_len < 5)
		{
        	*result = CHELLO_PARSE_NOT_ENOUGH_BUFF;
			return NULL;
		}
		if(buff[1] != 3 || buff[2] > 4 || buff[2] < 0)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return NULL;
		}

		struct ssl_chello* _chello = (struct ssl_chello*)ALLOC(struct ssl_chello, 1);
		_chello->min_version.major = buff[1];
		_chello->min_version.minor = buff[2];
		_chello->min_version.ossl_format=(uint16_t)_chello->min_version.major<<8|_chello->min_version.minor;
		_chello->max_version.major = -1;
		_chello->max_version.minor = -1;
		_chello->sni = NULL;
		_chello->alpn = NULL;
		_chello->cipher_suites = NULL;
		_chello->cipher_suites_len = 0;
		_chello->sign_algos = NULL;
		_chello->sign_algos_len = 0;
		_chello->supported_groups = NULL;
		_chello->supported_groups_len = 0;

		/* TLS record length */
    	size_t len = ((size_t)buff[3] << 8) + (size_t)buff[4] + 5;
		if (buff_len < len)
		{
			*result = CHELLO_PARSE_NOT_ENOUGH_BUFF;
        	return _chello;
		}
		buff_len = len;
		size_t pos = 5;
    	if (pos + 1 > buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
    	}
		if (buff[pos] != 0x01)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
    	}
		pos += 4;
		if(pos + 2 > buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
		_chello->max_version.major = buff[pos];
		_chello->max_version.minor = buff[pos+1];
		_chello->max_version.ossl_format=(uint16_t)_chello->max_version.major<<8|_chello->max_version.minor;

		pos += 34;
		/* Session ID */
    	if (pos + 1 > buff_len)
		{
        	*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
    	len = (size_t)buff[pos];
    	pos += 1 + len;
		/* Cipher Suites */
    	if (pos + 2 > buff_len)
		{
        	*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
    	len = ((size_t)buff[pos] << 8) + (size_t)buff[pos + 1];
		pos += 2;
		if(pos + len > buff_len || len % 2 != 0)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
		size_t i=0, j=0;
		_chello->cipher_suites = ALLOC(char, len);
		uint16_t* known_cipher=(uint16_t*)_chello->cipher_suites;
		uint16_t* raw_cipher=(uint16_t*)(buff + pos);
		for(i=0, j=0; i<len/2; i++)
		{
			//https://security.stackexchange.com/questions/176951/google-chrome-weird-random-cipher-suite
			if(!cipher_is_grease(raw_cipher[i]))
			{
				known_cipher[j++]=raw_cipher[i];
			}
		}
		_chello->cipher_suites_len = j*2;

		pos += len;

		/* Compression Methods */
    	if (pos >= buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
    	len = (size_t)buff[pos];
    	pos += 1 + len;

    	/* No extension */
        if(pos == buff_len)
        {
            *result = CHELLO_PARSE_SUCCESS;
            return _chello;
        }

		/* Extensions */
    	if (pos + 2 > buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
    	len = ((size_t)buff[pos] << 8) + (size_t)buff[pos + 1];
    	pos += 2;
		if (pos + len > buff_len)
		{
			*result = CHELLO_PARSE_INVALID_FORMAT;
        	return _chello;
		}
		int ret = parse_extensions(buff + pos, len, _chello);
		*result = (enum chello_parse_result)ret;
		return _chello;
	}
}