/* PipeWire */
/* SPDX-FileCopyrightText: Copyright © 2013 Tom Cocagne */
/* SPDX-License-Identifier: MIT */

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

#include <openssl/bn.h>
#include <openssl/evp.h>

#include <pipewire/log.h>

#include "srp.h"

#ifndef SHA512_DIGEST_LENGTH
#define SHA512_DIGEST_LENGTH 64
#endif

#ifndef SHA_DIGEST_LENGTH
#define SHA_DIGEST_LENGTH 20
#endif

#define SRP_SALT_BYTES 16

typedef struct {
	int N_len;
	BIGNUM *N;
	BIGNUM *g;
} NGConstant;

struct NGHex {
	int N_len;
	const char *n_hex;
	const char *g_hex;
};

// These constants here were pulled from Appendix A of RFC 5054
static struct NGHex Ng_constants[] = {
	{ /* 2048 */
		256,
		"AC6BDB41324A9A9BF166DE5E1389582FAF72B6651987EE07FC3192943DB56050A37329CBB4"
		"A099ED8193E0757767A13DD52312AB4B03310DCD7F48A9DA04FD50E8083969EDB767B0CF60"
		"95179A163AB3661A05FBD5FAAAE82918A9962F0B93B855F97993EC975EEAA80D740ADBF4FF"
		"747359D041D5C33EA71D281E446B14773BCA97B43A23FB801676BD207A436C6481F1D2B907"
		"8717461A5B9D32E688F87748544523B524B0D57D5EA77A2775D2ECFA032CFBDBF52FB37861"
		"60279004E57AE6AF874E7303CE53299CCC041C7BC308D82A5698F3A8D0C38271AE35F8E9DB"
		"FBB694B5C803D89F7AE435DE236D525F54759B65E372FCD68EF20FA7111F9E4AFF73",
		"2"
	},
	{ /* 3072 */
		384,
		"FFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B"
		"139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485"
		"B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1F"
		"E649286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23"
		"DCA3AD961C62F356208552BB9ED529077096966D670C354E4ABC9804F1746C08CA18217C32"
		"905E462E36CE3BE39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF69558"
		"17183995497CEA956AE515D2261898FA051015728E5A8AAAC42DAD33170D04507A33A85521"
		"ABDF1CBA64ECFB850458DBEF0A8AEA71575D060C7DB3970F85A6E1E4C7ABF5AE8CDB0933D7"
		"1E8C94E04A25619DCEE3D2261AD2EE6BF12FFA06D98A0864D87602733EC86A64521F2B1817"
		"7B200CBBE117577A615D6C770988C0BAD946E208E24FA074E5AB3143DB5BFCE0FD108E4B82"
		"D120A93AD2CAFFFFFFFFFFFFFFFF",
		"5"
	},
};

static NGConstant *new_ng(SRP_NGType ng_type, const char *n_hex, const char *g_hex)
{
	NGConstant *ng = calloc(1, sizeof(NGConstant));

	n_hex = Ng_constants[ng_type].n_hex;
	g_hex = Ng_constants[ng_type].g_hex;

	BN_hex2bn(&ng->N, n_hex);
	BN_hex2bn(&ng->g, g_hex);

	ng->N_len = BN_num_bytes(ng->N);

	assert(ng->N_len == Ng_constants[ng_type].N_len);

	return ng;
}

int get_N_len(SRP_NGType ng_type)
{
	return Ng_constants[ng_type].N_len;
}

struct SRPUser {
	SRP_HashAlgorithm  hash_alg;
	NGConstant        *ng;

	BIGNUM *a;
	BIGNUM *A;
	BIGNUM *S;

	const unsigned char *bytes_A;
	int                  len_A;
	int                  authenticated;

	char          *username;
	unsigned char *password;
	int            password_len;

	unsigned char M[SHA512_DIGEST_LENGTH];
	unsigned char H_AMK[SHA512_DIGEST_LENGTH];
	unsigned char session_key[SHA512_DIGEST_LENGTH];
	int           session_key_len;
};

const EVP_MD *hash_func(SRP_HashAlgorithm alg)
{
	switch (alg) {
	case SRP_SHA1:
		return EVP_sha1();
	case SRP_SHA512:
		return EVP_sha512();
	default:
		return NULL;
	};
}

static int hash_length(SRP_HashAlgorithm alg)
{
	switch (alg) {
	case SRP_SHA1:
		return SHA_DIGEST_LENGTH;
	case SRP_SHA512:
		return SHA512_DIGEST_LENGTH;
	default:
		return -1;
	};
}

static BIGNUM *H_nn_pad(SRP_HashAlgorithm alg, const BIGNUM *n1, const BIGNUM *n2, int padded_len)
{
	unsigned char  buff[SHA512_DIGEST_LENGTH];
	int            len_n1    = BN_num_bytes(n1);
	int            len_n2    = BN_num_bytes(n2);
	int            nbytes    = 2 * padded_len;
	int            offset_n1 = padded_len - len_n1;
	int            offset_n2 = nbytes - len_n2;
	unsigned char *bin       = (unsigned char *) malloc(nbytes);
	EVP_MD_CTX    *ctx       = EVP_MD_CTX_new();

	pw_log_debug("srp: H_nn_pad");

	if (!bin)
		return 0;

	assert(len_n1 <= padded_len);
	assert(len_n2 <= padded_len);

	BN_bn2bin(n1, bin + offset_n1);
	BN_bn2bin(n2, bin + offset_n2);
	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);
	EVP_DigestUpdate(ctx, bin, nbytes);
	EVP_DigestFinal_ex(ctx, buff, NULL);
	EVP_MD_CTX_free(ctx);
	free(bin);
  
	return BN_bin2bn(buff, hash_length(alg), NULL);
}

static BIGNUM *calculate_x(SRP_HashAlgorithm alg, const uint8_t *salt, int salt_len,
	const char *username, const unsigned char *password, int password_len)
{
	unsigned char x[SHA512_DIGEST_LENGTH];
	EVP_MD_CTX *ctx;

	assert(salt_len == SRP_SALT_BYTES);

	pw_log_debug("srp: calculate_x");

	ctx = EVP_MD_CTX_new();
	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);
	EVP_DigestUpdate(ctx, username, strlen(username));
	EVP_DigestUpdate(ctx, (const uint8_t *) ":", 1);
	EVP_DigestUpdate(ctx, password, password_len);
	EVP_DigestFinal_ex(ctx, x, NULL);
	EVP_MD_CTX_free(ctx);

	ctx = EVP_MD_CTX_new(); 
	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);
	EVP_DigestUpdate(ctx, (const uint8_t*) salt, salt_len);
	EVP_DigestUpdate(ctx, x, hash_length(alg));
	EVP_DigestFinal_ex(ctx, x, NULL);
	EVP_MD_CTX_free(ctx);
	
	return BN_bin2bn(x, hash_length(alg), NULL);
}

static void hash_num(SRP_HashAlgorithm alg, const BIGNUM *n, unsigned char *dest)
{
	int            nbytes = BN_num_bytes(n);
	unsigned char *bin    = (unsigned char *) malloc(nbytes);
	EVP_MD_CTX    *ctx    = EVP_MD_CTX_new();

	if (!bin)
		return;

	BN_bn2bin(n, bin);
	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);
	EVP_DigestUpdate(ctx, bin, nbytes);
	EVP_DigestFinal_ex(ctx, dest, NULL);
	EVP_MD_CTX_free(ctx);
	free(bin);
}

static void calculate_M(SRP_HashAlgorithm alg, NGConstant *ng, unsigned char *dest, const char *I,
	const uint8_t *salt, int salt_len, const uint8_t *A, int A_len, const uint8_t *B, int B_len,
	const unsigned char *K)
{
	unsigned char H_N[SHA512_DIGEST_LENGTH];
	unsigned char H_g[SHA512_DIGEST_LENGTH];
	unsigned char H_I[SHA512_DIGEST_LENGTH];
	unsigned char H_xor[SHA512_DIGEST_LENGTH];
	EVP_MD_CTX   *ctx;
	int           i = 0;
	int           hash_len = hash_length(alg);

	pw_log_debug("srp: calculate_M");

	hash_num(alg, ng->N, H_N);
	hash_num(alg, ng->g, H_g);

	ctx = EVP_MD_CTX_new();
	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);
	EVP_DigestUpdate(ctx, I, strlen(I));
	EVP_DigestFinal_ex(ctx, H_I, NULL);
	EVP_MD_CTX_free(ctx);

	for (i=0; i < hash_len; i++)
		H_xor[i] = H_N[i] ^ H_g[i];

	ctx = EVP_MD_CTX_new();
	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);
	
	EVP_DigestUpdate(ctx, H_xor, hash_len);
	EVP_DigestUpdate(ctx, H_I, hash_len);
	EVP_DigestUpdate(ctx, salt, salt_len);
	EVP_DigestUpdate(ctx, A, A_len);
	EVP_DigestUpdate(ctx, B, B_len);
	EVP_DigestUpdate(ctx, K, hash_len);

	EVP_DigestFinal_ex(ctx, dest, NULL);
	EVP_MD_CTX_free(ctx);
}

static void calculate_H_AMK(SRP_HashAlgorithm alg, unsigned char *dest, const uint8_t *A, int A_len,
	const unsigned char *M, const unsigned char *K)
{
	EVP_MD_CTX *ctx = EVP_MD_CTX_new();

	EVP_DigestInit_ex(ctx, hash_func(alg), NULL);

	EVP_DigestUpdate(ctx, A, A_len);
	EVP_DigestUpdate(ctx, M, hash_length(alg));
	EVP_DigestUpdate(ctx, K, hash_length(alg));

	EVP_DigestFinal_ex(ctx, dest, NULL);
	EVP_MD_CTX_free(ctx);
}

struct SRPUser *srp_user_new(SRP_HashAlgorithm alg, SRP_NGType ng_type, const char *username,
	const unsigned char *bytes_password, int len_password, const char *n_hex, const char *g_hex)
{
	struct SRPUser *usr  = (struct SRPUser *) malloc(sizeof(struct SRPUser));
	int             ulen = strlen(username) + 1;

	if (!usr)
		goto err_exit;

	usr->hash_alg = alg;
	usr->ng  = new_ng(ng_type, n_hex, g_hex);

	usr->a = BN_new();
	usr->A = BN_new();
	usr->S = BN_new();

	if (!usr->ng || !usr->a || !usr->A || !usr->S)
		goto err_exit;

	usr->username     = (char *) malloc(ulen);
	usr->password     = (unsigned char *) malloc(len_password);
	usr->password_len = len_password;

	if (!usr->username || !usr->password)
		goto err_exit;

	memcpy(usr->username, username,       ulen);
	memcpy(usr->password, bytes_password, len_password);

	usr->authenticated = 0;
	usr->bytes_A = 0;
	usr->len_A = 0;

	return usr;

err_exit:
	if (!usr)
		return NULL;

	BN_free(usr->a);
	BN_free(usr->A);
	BN_free(usr->S);
	if (usr->username)
		free(usr->username);
	if (usr->password) {
		memset(usr->password, 0, usr->password_len);
		free(usr->password);
	}
	free(usr);

	return 0;
}

void srp_user_start_authentication(struct SRPUser *usr, const char **username, const unsigned char **bytes_A, int *len_A)
{
	BN_CTX *bn_ctx = BN_CTX_new();

	pw_log_debug("srp: srp_user_start_authentication");

	BN_rand(usr->a, 256, -1, 0);
	BN_mod_exp(usr->A, usr->ng->g, usr->a, usr->ng->N, bn_ctx);

	BN_CTX_free(bn_ctx);

	*len_A   = BN_num_bytes(usr->A);
	*bytes_A = (const unsigned char *) malloc(*len_A);

	if (!*bytes_A) {
		pw_log_debug("srp: srp_user_start_authentication something went wrong");
		*len_A = 0;
		*bytes_A = 0;
		*username = 0;
		return;
	}

	BN_bn2bin(usr->A, (unsigned char *) *bytes_A);

	usr->bytes_A = *bytes_A;
	usr->len_A = *len_A;
	*username = usr->username;
}

void srp_user_process_challenge(struct SRPUser *usr, const unsigned char *bytes_s, int len_s,
	const unsigned char *bytes_B, int len_B, const unsigned char **bytes_M, int *len_M)
{
	BIGNUM *B      = BN_bin2bn(bytes_B, len_B, NULL);
	BIGNUM *u      = 0;
	BIGNUM *x      = 0;
	BIGNUM *k      = 0;
	BIGNUM *v      = BN_new();
	BIGNUM *tmp1   = BN_new();
	BIGNUM *tmp2   = BN_new();
	BIGNUM *tmp3   = BN_new();
	BN_CTX *bn_ctx = BN_CTX_new();

	pw_log_debug("srp: srp_user_process_challenge");

	*len_M = 0;
	*bytes_M = 0;

	if (!bytes_s || !B || !v || !tmp1 || !tmp2 || !tmp3 || !bn_ctx) {
		pw_log_debug("srp: BIGNUM init failed");
		goto cleanup_and_exit;
	}
	k = H_nn_pad(usr->hash_alg, usr->ng->N, usr->ng->g, usr->ng->N_len);
	if (!k) {
		pw_log_debug("srp: couldn't calc k");
		goto cleanup_and_exit;
	}
	u = H_nn_pad(usr->hash_alg, usr->A, B, usr->ng->N_len);
	if (!u) {
	   pw_log_debug("srp: couldn't calc u");
	   goto cleanup_and_exit;
	}
	x = calculate_x(usr->hash_alg, bytes_s, len_s, usr->username, usr->password, usr->password_len);
	if (!x) {
		pw_log_debug("srp: couldn't calc x");
		goto cleanup_and_exit;
	}
	/* SRP-6a safety check */
	if (!BN_is_zero(B) && !BN_is_zero(u)) {
  		BN_mod_exp(v, usr->ng->g, x, usr->ng->N, bn_ctx);

		/* S = (B - k*(g^x)) ^ (a + ux) */
  		BN_mul(tmp1, u, x, bn_ctx);
		BN_add(tmp2, usr->a, tmp1);     /* tmp2 = (a + ux)      */
		BN_mod_exp(tmp1, usr->ng->g, x, usr->ng->N, bn_ctx);
		BN_mul(tmp3, k, tmp1, bn_ctx);     /* tmp3 = k*(g^x)       */
		BN_sub(tmp1, B, tmp3);          /* tmp1 = (B - K*(g^x)) */
		BN_mod_exp(usr->S, tmp1, tmp2, usr->ng->N, bn_ctx);

		hash_num(usr->hash_alg, usr->S, usr->session_key);

		calculate_M(usr->hash_alg, usr->ng, usr->M, usr->username, bytes_s, len_s, usr->bytes_A, usr->len_A, bytes_B, len_B, usr->session_key);
		calculate_H_AMK(usr->hash_alg, usr->H_AMK, usr->bytes_A, usr->len_A, usr->M, usr->session_key);

		*bytes_M = usr->M;
		if (len_M)
			*len_M = hash_length(usr->hash_alg);
	}

cleanup_and_exit:
	BN_free(B);
	BN_free(u);
	BN_free(x);
	BN_free(k);
	BN_free(v);
	BN_free(tmp1);
	BN_free(tmp2);
	BN_free(tmp3);
	BN_CTX_free(bn_ctx);
}

int srp_user_is_authenticated(struct SRPUser *usr)
{
	return usr->authenticated;
}

const unsigned char *srp_user_get_session_key(struct SRPUser *usr, int *key_length)
{
	if (key_length)
		*key_length = hash_length(usr->hash_alg);

	return usr->session_key;
}

void srp_user_verify_session(struct SRPUser *usr, const unsigned char *bytes_HAMK)
{
	if (memcmp(usr->H_AMK, bytes_HAMK, hash_length(usr->hash_alg)) == 0 )
		usr->authenticated = 1;
}