/* HexChat
*
* 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.
*
* 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.
*
* 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., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
*/

/*
 * For Diffie-Hellman key-exchange a 1080bit germain prime is used, the
 * generator g=2 renders a field Fp from 1 to p-1. Therefore breaking it
 * means to solve a discrete logarithm problem with no less than 1080bit.
 *
 * Base64 format is used to send the public keys over IRC.
 *
 * The calculated secret key is hashed with SHA-256, the result is converted
 * to base64 for final use with blowfish.
 */

#include "config.h"
#include "dh1080.h"

#include <openssl/bn.h>
#include <openssl/sha.h>
#include <openssl/opensslv.h>

#if OPENSSL_VERSION_NUMBER >= 0x30000000L
#include <openssl/evp.h>
#include <openssl/param_build.h>
#include <openssl/core_names.h>
#else
#include <openssl/dh.h>
#endif

#include <string.h>
#include <glib.h>

#define DH1080_PRIME_BITS 1080
#define DH1080_PRIME_BYTES 135
#define SHA256_DIGEST_LENGTH 32
#define B64ABC "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
#define MEMZERO(x) memset(x, 0x00, sizeof(x))

/* All clients must use the same prime number to be able to keyx */
static const guchar prime1080[DH1080_PRIME_BYTES] =
{
	0xFB, 0xE1, 0x02, 0x2E, 0x23, 0xD2, 0x13, 0xE8, 0xAC, 0xFA, 0x9A, 0xE8, 0xB9, 0xDF, 0xAD, 0xA3, 0xEA,
	0x6B, 0x7A, 0xC7, 0xA7, 0xB7, 0xE9, 0x5A, 0xB5, 0xEB, 0x2D, 0xF8, 0x58, 0x92, 0x1F, 0xEA, 0xDE, 0x95,
	0xE6, 0xAC, 0x7B, 0xE7, 0xDE, 0x6A, 0xDB, 0xAB, 0x8A, 0x78, 0x3E, 0x7A, 0xF7, 0xA7, 0xFA, 0x6A, 0x2B,
	0x7B, 0xEB, 0x1E, 0x72, 0xEA, 0xE2, 0xB7, 0x2F, 0x9F, 0xA2, 0xBF, 0xB2, 0xA2, 0xEF, 0xBE, 0xFA, 0xC8,
	0x68, 0xBA, 0xDB, 0x3E, 0x82, 0x8F, 0xA8, 0xBA, 0xDF, 0xAD, 0xA3, 0xE4, 0xCC, 0x1B, 0xE7, 0xE8, 0xAF,
	0xE8, 0x5E, 0x96, 0x98, 0xA7, 0x83, 0xEB, 0x68, 0xFA, 0x07, 0xA7, 0x7A, 0xB6, 0xAD, 0x7B, 0xEB, 0x61,
	0x8A, 0xCF, 0x9C, 0xA2, 0x89, 0x7E, 0xB2, 0x8A, 0x61, 0x89, 0xEF, 0xA0, 0x7A, 0xB9, 0x9A, 0x8A, 0x7F,
	0xA9, 0xAE, 0x29, 0x9E, 0xFA, 0x7B, 0xA6, 0x6D, 0xEA, 0xFE, 0xFB, 0xEF, 0xBF, 0x0B, 0x7D, 0x8B
};

#if OPENSSL_VERSION_NUMBER >= 0x30000000L

/* OpenSSL 3.0+ implementation using EVP API */

static EVP_PKEY *g_dh_params = NULL;

int
dh1080_init (void)
{
	EVP_PKEY_CTX *pctx = NULL;
	OSSL_PARAM_BLD *param_bld = NULL;
	OSSL_PARAM *params = NULL;
	BIGNUM *p = NULL, *g = NULL;
	int ret = 0;

	g_return_val_if_fail (g_dh_params == NULL, 0);

	p = BN_bin2bn (prime1080, DH1080_PRIME_BYTES, NULL);
	g = BN_new ();

	if (p == NULL || g == NULL)
		goto cleanup;

	BN_set_word (g, 2);

	param_bld = OSSL_PARAM_BLD_new ();
	if (param_bld == NULL)
		goto cleanup;

	if (!OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_FFC_P, p) ||
	    !OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_FFC_G, g))
		goto cleanup;

	params = OSSL_PARAM_BLD_to_param (param_bld);
	if (params == NULL)
		goto cleanup;

	pctx = EVP_PKEY_CTX_new_from_name (NULL, "DH", NULL);
	if (pctx == NULL)
		goto cleanup;

	if (EVP_PKEY_fromdata_init (pctx) <= 0)
		goto cleanup;

	if (EVP_PKEY_fromdata (pctx, &g_dh_params, EVP_PKEY_KEY_PARAMETERS, params) <= 0)
		goto cleanup;

	ret = 1;

cleanup:
	OSSL_PARAM_free (params);
	OSSL_PARAM_BLD_free (param_bld);
	EVP_PKEY_CTX_free (pctx);
	BN_free (p);
	BN_free (g);

	if (!ret && g_dh_params) {
		EVP_PKEY_free (g_dh_params);
		g_dh_params = NULL;
	}

	return ret;
}

void
dh1080_deinit (void)
{
	if (g_dh_params) {
		EVP_PKEY_free (g_dh_params);
		g_dh_params = NULL;
	}
}

static int
DH_verifyPubKey (BIGNUM *pk)
{
	/* In OpenSSL 3.0+, verification happens during key derivation */
	/* Basic sanity check: pub key should be > 1 and < p-1 */
	BIGNUM *p = NULL;
	BIGNUM *p_minus_1 = NULL;
	BIGNUM *one = NULL;
	int ret = 0;

	if (!EVP_PKEY_get_bn_param (g_dh_params, OSSL_PKEY_PARAM_FFC_P, &p))
		return 0;

	p_minus_1 = BN_new ();
	one = BN_new ();
	if (!p_minus_1 || !one)
		goto cleanup;

	BN_one (one);
	BN_sub (p_minus_1, p, one);

	/* pub_key must be > 1 and < p-1 */
	if (BN_cmp (pk, one) > 0 && BN_cmp (pk, p_minus_1) < 0)
		ret = 1;

cleanup:
	BN_free (p);
	BN_free (p_minus_1);
	BN_free (one);
	return ret;
}

#else

/* OpenSSL < 3.0 implementation using deprecated DH API */

static DH *g_dh;

int
dh1080_init (void)
{
	g_return_val_if_fail (g_dh == NULL, 0);

	if ((g_dh = DH_new()))
	{
		int codes;
		BIGNUM *p, *g;

		p = BN_bin2bn (prime1080, DH1080_PRIME_BYTES, NULL);
		g = BN_new ();

		if (p == NULL || g == NULL)
			return 1;

		BN_set_word (g, 2);

#ifndef HAVE_DH_SET0_PQG
		g_dh->p = p;
		g_dh->g = g;
#else
		if (!DH_set0_pqg (g_dh, p, NULL, g))
			return 1;
#endif

		if (DH_check (g_dh, &codes))
			return codes == 0;
	}

	return 0;
}

void
dh1080_deinit (void)
{
	g_clear_pointer (&g_dh, DH_free);
}

static inline int
DH_verifyPubKey (BIGNUM *pk)
{
	int codes;
	return DH_check_pub_key (g_dh, pk, &codes) && codes == 0;
}

#endif /* OPENSSL_VERSION_NUMBER >= 0x30000000L */

static guchar *
dh1080_decode_b64 (const char *data, gsize *out_len)
{
	GString *str = g_string_new (data);
	guchar *ret;

	if (str->len % 4 == 1 && str->str[str->len - 1] == 'A')
		g_string_truncate (str, str->len - 1);

	while (str->len % 4 != 0)
		g_string_append_c (str, '=');

	ret = g_base64_decode_inplace (str->str, out_len);
	g_string_free (str, FALSE);
  	return ret;
}

static char *
dh1080_encode_b64 (const guchar *data, gsize data_len)
{
	char *ret = g_base64_encode (data, data_len);
	char *p;

	if (!(p = strchr (ret, '=')))
	{
		char *new_ret = g_new(char, strlen(ret) + 2);
		strcpy (new_ret, ret);
		strcat (new_ret, "A");
		g_free (ret);
		ret = new_ret;
	}
	else
	{
		*p = '\0';
	}

  	return ret;
}

#if OPENSSL_VERSION_NUMBER >= 0x30000000L

/* OpenSSL 3.0+ implementation using EVP API */

int
dh1080_generate_key (char **priv_key, char **pub_key)
{
	guchar buf[DH1080_PRIME_BYTES];
	int len;
	EVP_PKEY_CTX *ctx = NULL;
	EVP_PKEY *keypair = NULL;
	BIGNUM *priv_bn = NULL, *pub_bn = NULL;
	int ret = 0;

	g_assert (priv_key != NULL);
	g_assert (pub_key != NULL);

	ctx = EVP_PKEY_CTX_new_from_pkey (NULL, g_dh_params, NULL);
	if (!ctx)
		goto cleanup;

	if (EVP_PKEY_keygen_init (ctx) <= 0)
		goto cleanup;

	if (EVP_PKEY_keygen (ctx, &keypair) <= 0)
		goto cleanup;

	if (!EVP_PKEY_get_bn_param (keypair, OSSL_PKEY_PARAM_PRIV_KEY, &priv_bn))
		goto cleanup;

	if (!EVP_PKEY_get_bn_param (keypair, OSSL_PKEY_PARAM_PUB_KEY, &pub_bn))
		goto cleanup;

	MEMZERO (buf);
	len = BN_bn2bin (priv_bn, buf);
	*priv_key = dh1080_encode_b64 (buf, len);

	MEMZERO (buf);
	len = BN_bn2bin (pub_bn, buf);
	*pub_key = dh1080_encode_b64 (buf, len);

	OPENSSL_cleanse (buf, sizeof (buf));
	ret = 1;

cleanup:
	BN_free (priv_bn);
	BN_free (pub_bn);
	EVP_PKEY_free (keypair);
	EVP_PKEY_CTX_free (ctx);
	return ret;
}

int
dh1080_compute_key (const char *priv_key, const char *pub_key, char **secret_key)
{
	guchar *pub_key_data = NULL;
	gsize pub_key_len;
	BIGNUM *pk = NULL;
	EVP_PKEY *local_key = NULL;
	EVP_PKEY *peer_key = NULL;
	EVP_PKEY_CTX *derive_ctx = NULL;
	OSSL_PARAM_BLD *param_bld = NULL;
	OSSL_PARAM *params = NULL;
	EVP_PKEY_CTX *pctx = NULL;
	guchar *priv_key_data = NULL;
	gsize priv_key_len;
	BIGNUM *priv_bn = NULL, *p_bn = NULL, *g_bn = NULL;
	int ret = 0;

	g_assert (secret_key != NULL);

	/* Verify base64 strings */
	if (strspn (priv_key, B64ABC) != strlen (priv_key)
	    || strspn (pub_key, B64ABC) != strlen (pub_key))
		return 0;

	pub_key_data = dh1080_decode_b64 (pub_key, &pub_key_len);
	pk = BN_bin2bn (pub_key_data, pub_key_len, NULL);

	if (!DH_verifyPubKey (pk))
		goto cleanup;

	/* Get p and g from params */
	if (!EVP_PKEY_get_bn_param (g_dh_params, OSSL_PKEY_PARAM_FFC_P, &p_bn))
		goto cleanup;
	if (!EVP_PKEY_get_bn_param (g_dh_params, OSSL_PKEY_PARAM_FFC_G, &g_bn))
		goto cleanup;

	/* Build local key with our private key */
	priv_key_data = dh1080_decode_b64 (priv_key, &priv_key_len);
	priv_bn = BN_bin2bn (priv_key_data, priv_key_len, NULL);

	param_bld = OSSL_PARAM_BLD_new ();
	if (!param_bld)
		goto cleanup;

	if (!OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_FFC_P, p_bn) ||
	    !OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_FFC_G, g_bn) ||
	    !OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_PRIV_KEY, priv_bn) ||
	    !OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_PUB_KEY, pk))
		goto cleanup;

	params = OSSL_PARAM_BLD_to_param (param_bld);
	if (!params)
		goto cleanup;

	pctx = EVP_PKEY_CTX_new_from_name (NULL, "DH", NULL);
	if (!pctx)
		goto cleanup;

	if (EVP_PKEY_fromdata_init (pctx) <= 0)
		goto cleanup;

	if (EVP_PKEY_fromdata (pctx, &local_key, EVP_PKEY_KEYPAIR, params) <= 0)
		goto cleanup;

	OSSL_PARAM_free (params);
	params = NULL;
	OSSL_PARAM_BLD_free (param_bld);
	param_bld = NULL;
	EVP_PKEY_CTX_free (pctx);
	pctx = NULL;

	/* Build peer key with their public key */
	param_bld = OSSL_PARAM_BLD_new ();
	if (!param_bld)
		goto cleanup;

	if (!OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_FFC_P, p_bn) ||
	    !OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_FFC_G, g_bn) ||
	    !OSSL_PARAM_BLD_push_BN (param_bld, OSSL_PKEY_PARAM_PUB_KEY, pk))
		goto cleanup;

	params = OSSL_PARAM_BLD_to_param (param_bld);
	if (!params)
		goto cleanup;

	pctx = EVP_PKEY_CTX_new_from_name (NULL, "DH", NULL);
	if (!pctx)
		goto cleanup;

	if (EVP_PKEY_fromdata_init (pctx) <= 0)
		goto cleanup;

	if (EVP_PKEY_fromdata (pctx, &peer_key, EVP_PKEY_PUBLIC_KEY, params) <= 0)
		goto cleanup;

	/* Derive shared secret */
	derive_ctx = EVP_PKEY_CTX_new_from_pkey (NULL, local_key, NULL);
	if (!derive_ctx)
		goto cleanup;

	if (EVP_PKEY_derive_init (derive_ctx) <= 0)
		goto cleanup;

	if (EVP_PKEY_derive_set_peer (derive_ctx, peer_key) <= 0)
		goto cleanup;

	{
		guchar shared_key[DH1080_PRIME_BYTES] = { 0 };
		guchar sha256[SHA256_DIGEST_LENGTH] = { 0 };
		size_t shared_len = sizeof (shared_key);

		if (EVP_PKEY_derive (derive_ctx, shared_key, &shared_len) <= 0)
			goto cleanup;

		SHA256 (shared_key, shared_len, sha256);
		*secret_key = dh1080_encode_b64 (sha256, sizeof (sha256));
		OPENSSL_cleanse (shared_key, sizeof (shared_key));
	}

	ret = 1;

cleanup:
	if (priv_key_data) {
		OPENSSL_cleanse (priv_key_data, priv_key_len);
		g_free (priv_key_data);
	}
	g_free (pub_key_data);
	BN_free (pk);
	BN_free (priv_bn);
	BN_free (p_bn);
	BN_free (g_bn);
	OSSL_PARAM_free (params);
	OSSL_PARAM_BLD_free (param_bld);
	EVP_PKEY_CTX_free (pctx);
	EVP_PKEY_CTX_free (derive_ctx);
	EVP_PKEY_free (local_key);
	EVP_PKEY_free (peer_key);
	return ret;
}

#else

/* OpenSSL < 3.0 implementation using deprecated DH API */

int
dh1080_generate_key (char **priv_key, char **pub_key)
{
	guchar buf[DH1080_PRIME_BYTES];
	int len;
	DH *dh;
	const BIGNUM *dh_priv_key, *dh_pub_key;

  	g_assert (priv_key != NULL);
	g_assert (pub_key != NULL);

  	dh = DHparams_dup (g_dh);
	if (!dh)
		return 0;

	if (!DH_generate_key (dh))
	{
		DH_free (dh);
		return 0;
	}

#ifndef HAVE_DH_GET0_KEY
	dh_pub_key = dh->pub_key;
	dh_priv_key = dh->priv_key;
#else
	DH_get0_key (dh, &dh_pub_key, &dh_priv_key);
#endif

	MEMZERO (buf);
	len = BN_bn2bin (dh_priv_key, buf);
	*priv_key = dh1080_encode_b64 (buf, len);

	MEMZERO (buf);
	len = BN_bn2bin (dh_pub_key, buf);
	*pub_key = dh1080_encode_b64 (buf, len);

	OPENSSL_cleanse (buf, sizeof (buf));
	DH_free (dh);
	return 1;
}

int
dh1080_compute_key (const char *priv_key, const char *pub_key, char **secret_key)
{
	char *pub_key_data;
	gsize pub_key_len;
	BIGNUM *pk;
	DH *dh;
#ifdef HAVE_DH_SET0_KEY
	BIGNUM *temp_pub_key = BN_new();
#endif

	g_assert (secret_key != NULL);

	/* Verify base64 strings */
	if (strspn (priv_key, B64ABC) != strlen (priv_key)
	    || strspn (pub_key, B64ABC) != strlen (pub_key))
		return 0;

	dh = DHparams_dup (g_dh);

	pub_key_data = dh1080_decode_b64 (pub_key, &pub_key_len);
	pk = BN_bin2bn (pub_key_data, pub_key_len, NULL);

	if (DH_verifyPubKey (pk))
	{
		guchar shared_key[DH1080_PRIME_BYTES] = { 0 };
		guchar sha256[SHA256_DIGEST_LENGTH] = { 0 };
		char *priv_key_data;
		gsize priv_key_len;
		int shared_len;
		BIGNUM *priv_key_num;

		priv_key_data = dh1080_decode_b64 (priv_key, &priv_key_len);
		priv_key_num = BN_bin2bn(priv_key_data, priv_key_len, NULL);
#ifndef HAVE_DH_SET0_KEY
		dh->priv_key = priv_key_num;
#else
		DH_set0_key (dh, temp_pub_key, priv_key_num);
#endif

		shared_len = DH_compute_key (shared_key, pk, dh);
		SHA256(shared_key, shared_len, sha256);
		*secret_key = dh1080_encode_b64 (sha256, sizeof(sha256));

		OPENSSL_cleanse (priv_key_data, priv_key_len);
		g_free (priv_key_data);
	}

	BN_free (pk);
	DH_free (dh);
	g_free (pub_key_data);
	return 1;
}

#endif /* OPENSSL_VERSION_NUMBER >= 0x30000000L */