#include "lib5.h"
#include "lib4.h"
#include "lib3.h"
#include "lib2.h"
#include "lib.h"

void mod_bignums(unsigned char *number, unsigned char *mod, unsigned int base, unsigned char **erg)
{
	mpz_t number_mp, mod_mp, erg_mp;

	mpz_init_set_str(number_mp, number, base);
	mpz_init_set_str(mod_mp, mod, base);
	mpz_init(erg_mp);
	mpz_mod(erg_mp, number_mp, mod_mp);

	(*erg) = malloc(mpz_sizeinbase(erg_mp,16)+2);
	mpz_get_str(*erg, base, erg_mp);
}

void modexp_mpz(mpz_t *base_mp, unsigned char *exp, unsigned char *mod, int string_base,
		mpz_t *erg_mp)
{
	mpz_t exp_mp, mod_mp;
 
	mpz_init_set_str(exp_mp, exp, string_base);
	mpz_init_set_str(mod_mp, mod, string_base);
	mpz_init(*erg_mp);
 
	mpz_powm(*erg_mp, *base_mp, exp_mp, mod_mp);
 
	mpz_clear(exp_mp);
	mpz_clear(mod_mp);
}


void modexp_bignums(unsigned char *base, unsigned char *exp, unsigned char *mod, int string_base,
		mpz_t *erg_mp)
{
	mpz_t base_mp, exp_mp, mod_mp;
 
	mpz_init_set_str(base_mp, base, string_base);
	mpz_init_set_str(exp_mp, exp, string_base);
	mpz_init_set_str(mod_mp, mod, string_base);
	mpz_init(*erg_mp);
 
	mpz_powm(*erg_mp, base_mp, exp_mp, mod_mp);
	mpz_clear(base_mp);
	mpz_clear(exp_mp);
	mpz_clear(mod_mp);

	//gmp_printf("%Zd\n", *erg_mp); 
}

void dh_init(struct dh_param *dh)
{
	dh->p = "ffffffffffffffffc90fdaa22168c234c4c6628b80dc1cd129024"
		"e088a67cc74020bbea63b139b22514a08798e3404ddef9519b3cd"
		"3a431b302b0a6df25f14374fe1356d6d51c245e485b576625e7ec"
		"6f44c42e9a637ed6b0bff5cb6f406b7edee386bfb5a899fa5ae9f"
		"24117c4b1fe649286651ece45b3dc2007cb8a163bf0598da48361"
		"c55d39a69163fa8fd24cf5f83655d23dca3ad961c62f356208552"
		"bb9ed529077096966d670c354e4abc9804f1746c08ca237327fff"
		"fffffffffffff";
	dh->g = "2";
}
void dh_generate_secret_keys(struct dh_param *dh)
{
	unsigned char b_tmp[1000];
	unsigned char a_tmp[1000];
	generate_random_hex(a_tmp,1000);
	a_tmp[999] = '\0';
	mod_bignums(a_tmp,dh->p,16,&(dh->a));
	generate_random_hex(b_tmp,1000);
	b_tmp[999] = '\0';
	mod_bignums(b_tmp,dh->p,16,&(dh->b));
}

void dh_generate_public_keys(struct dh_param *dh)
{
	modexp_bignums(dh->g, dh->a, dh->p, 16, &(dh->A));
	modexp_bignums(dh->g, dh->b, dh->p, 16, &(dh->B));
}

void dh_get_session_key(struct dh_param *dh)
{
	modexp_mpz(&(dh->B), dh->a, dh->p, 16, &(dh->s1));
	modexp_mpz(&(dh->A), dh->b, dh->p, 16, &(dh->s2));
	//printf("sessino keys are:\n");
	gmp_printf("%Zd\n", dh->s1);
	//gmp_printf("%Zd\n", dh->s1);
}

void do_dh_key_exchange(struct dh_param *dh)
{
	dh_init(dh);
	dh_generate_secret_keys(dh);
	dh_generate_public_keys(dh);
	dh_get_session_key(dh);
}

void sha1_key_from_dh(struct dh_param *dh, unsigned char *key)
{
	char *s1_char;
	SHA1Context sha1;

	s1_char = malloc(mpz_sizeinbase(dh->s1,16)+2);
	memset(s1_char, 0, 16);
	mpz_get_str(s1_char, 16, dh->s1);
	SHA1Reset(&sha1);
	// only use the first 16 bytes accoring to the challenge
	SHA1Input(&sha1, s1_char, 16);
	SHA1Result(&sha1);

	memcpy(key, &(sha1.Message_Digest), 20);
}

void dh_mitm(struct dh_param *dh)
{
	dh_init(dh);
	dh_generate_secret_keys(dh);
	dh_generate_public_keys(dh);
	// swap the public keys with p
	// p mod p will always be 0; s = 0
	mpz_init_set_str(dh->A, dh->p, 16);
	mpz_init_set_str(dh->B, dh->p, 16);
	dh_get_session_key(dh);
}

void srp_compute_x(int salt, unsigned char *password, char *sha1_hash)
{
	unsigned char *to_hash;
	SHA1Context sha1;

	to_hash = malloc(strlen(password) + sizeof(int));

	memcpy(to_hash, &salt, sizeof(int));
	memcpy(&to_hash[sizeof(int)], password, strlen(password));

	SHA1Reset(&sha1);
	SHA1Input(&sha1, to_hash, strlen(to_hash));
	SHA1Result(&sha1);
	
	memcpy(sha1_hash, &(sha1.Message_Digest), 20);
}

void srp_server_init(char *email, char *password, char *g, char *N)
{
	int salt;
	char sha1_hash[20];
	mpz_t sha1_as_number;
	mpz_t v;

	generate_random_bytes((char *)&salt, sizeof(int));

	srp_compute_x(salt, password, sha1_hash);
	modexp_bignums(g, sha1_hash, N, 16, &v);
}

void srp_client_send1(char *g)
{
	// send email

	// compute public key A
	//char *a
}

void srp_server_send1()
{
	// send salt

	// compute public key B
}

void srp_compute_uH(unsigned char *A, unsigned char *B)
{
	SHA1Context sha1;
	unsigned char uH[20];
	mpz_t u;
	unsigned char *res = malloc(strlen(A) + strlen(B));

	memcpy(res, A, strlen(A));
	memcpy(&res[strlen(A)], B, strlen(B));

	SHA1Reset(&sha1);
	SHA1Input(&sha1, res, (strlen(A) + strlen(B)));
	SHA1Result(&sha1);

	memcpy(uH, &(sha1.Message_Digest), 20);

	mpz_init_set_str(u, uH, 16);
}
/* 
void srp_client(unsigned char *salt, unsigned char *password, unsigned char *g,
		unsigned char *N, unsigned char *B, unsigned char *k)
{
	char sha1_hash[20];
	mpz_t g_mp, N_mp, B_mp, k_mp, tmp_mp;

	srp_compute_x(salt, password, sha1_hash);

	mpz_init_set_str(g_mp, g, 16);
	mpz_init_set_str(N_mp, N, 16);
	mpz_init_set_str(B_mp, B, 16);
	mpz_init_set_str(k_mp, k, 16);

	mpz_pow_
}
*/

/**
  * in C the % operator is more the remainder than the modulo
  * so implement modulo which also works fine with negative numbers
  */
int modulo(int a, int b)
{
	int mod = a % b;

	if (mod*b < 0)
		return mod + b;
	else
		return mod;
}

void extended_euclid_algo(int a, int b, struct extended_euclid *e)
{
	struct extended_euclid *tmp = malloc(sizeof(struct extended_euclid));

	if (b == 0) {
		e->d=a;
		e->s=1;
		e->t=0;
		return;
	}
	extended_euclid_algo(b, a % b, tmp);
	e->d = tmp->d;
	e->s = tmp->t;
	e->t = tmp->s - (a / b) * tmp->t;
	free(tmp);
	return;
}

int modular_multiplicative_inverse(int number, int _modulo)
{
	struct extended_euclid tmp;
	extended_euclid_algo(number, _modulo, &tmp);
	// only has a inverse iff gcd = 1
	if ( tmp.d != 1)
		return INT_MIN;

	// mod works not fine for negytive numbers in c
	return modulo(tmp.s, _modulo);
}

int rsa_encrypt(int message, struct rsa_key *public)
{
	return modulo((message^public->exponent), public->modulo);
}

int rsa_decrpyt(int message, struct rsa_key *private)
{
	return modulo((message^private->exponent), private->modulo);
}

void die(char *message)
{
	printf("%s\n", message);
	exit(1);
}

int rsa_encrypt_bignum(BIGNUM *message, BIGNUM *res, struct rsa_key_bignum *public)
{
	return BN_mod_exp(res, message, public->exponent, public->modulo, ctx);
}

int rsa_decrypt_bignum(BIGNUM *message, BIGNUM *res, struct rsa_key_bignum *private)
{
	return BN_mod_exp(res, message, private->exponent, private->modulo, ctx);
}

void extended_euclid_algo_bignum(BIGNUM *a, BIGNUM *b, struct extended_euclid_bignum *e)
{
	struct extended_euclid_bignum tmp;

	if (BN_is_zero(b)) {
		BN_copy(e->d, a);
		BN_one(e->s);
		BN_zero(e->t);
		return;
	}
	tmp.d = BN_new();
	tmp.s = BN_new();
	tmp.t = BN_new();

	BIGNUM *mod = BN_new();
	BN_mod(mod, a, b, ctx);

	extended_euclid_algo_bignum(b, mod, &tmp);
	BN_copy(e->d, tmp.d);
	BN_copy(e->s, tmp.t);
	BN_div(mod, NULL, a, b, ctx);
	BN_mul(mod, mod, tmp.t, ctx);
	BN_sub(e->t, tmp.s, mod);
	
	
	BN_free(mod);
	BN_free(tmp.d);
	BN_free(tmp.s);
	BN_free(tmp.t);
	return;
}

int modular_multiplicative_inverse_bignum_my(BIGNUM *res, BIGNUM *number, BIGNUM *modulo)
{
	struct extended_euclid_bignum tmp;
	tmp.d = BN_new();
	tmp.s = BN_new();
	tmp.t = BN_new();
	extended_euclid_algo_bignum(number, modulo, &tmp);
	// only has a invese iff gcd = 1
	if (!BN_is_one(tmp.d))
		return -1;

	return BN_nnmod(res, tmp.s, modulo, ctx);
}

int rsa_generate_key_bignum(struct rsa_key_bignum *public, struct rsa_key_bignum *private)
{
	// RSA with bignum
	// using openssl'S BN
	BIGNUM *p = BN_new();
	// well should check here for error but asusme infinte memory here
	BIGNUM *q = BN_new();

	BN_generate_prime_ex(p, 256, 1, NULL, NULL, NULL);
	do {
		BN_generate_prime_ex(q, 256, 1, NULL, NULL, NULL);
	} while (!BN_cmp(p, q));
	BIGNUM *n = BN_new();

	if(!BN_mul(n,p,q,ctx))
		die("error multipling p and q");

	BIGNUM *et = BN_new();
	BIGNUM *p_1 = BN_new();
	BIGNUM *q_1= BN_new();
	BIGNUM *one = BN_new();
	BN_one(one);

	if(!BN_sub(p_1, p, one))
		die("could not substract one from p");

	if(!BN_sub(q_1, q, one))
		die("could not substract one from q");

	if(!BN_mul(et, p_1, q_1, ctx))
		die("could not multiply p*q");


	BIGNUM *e = BN_new();
	BN_set_word(e, 3);
	//BIGNUM *d = BN_mod_inverse(NULL, e, et, ctx);
	BIGNUM *d = BN_new();
	modular_multiplicative_inverse_bignum_my(d, e, et);
	public->exponent = BN_new();
	public->modulo = BN_new();
	private->exponent = BN_new();
	private->modulo = BN_new();
	public->exponent = BN_dup(e);
	BN_copy(public->modulo, n);
	BN_copy(private->exponent, d);
	BN_copy(private->modulo, n);

}

int free_rsa_key_bignum(struct rsa_key_bignum *t)
{
	BN_free(t->exponent);
	BN_free(t->modulo);
}

/**
  * computes the nth root of number.
  * Note that the computed root is always an integer
  * does not work good for numbers which are not divisible by n :-(
  **/ 
int nth_root_bignum(BIGNUM *res, BIGNUM *number, BIGNUM *n)
{
	BIGNUM *n_1 = BN_new();
	BIGNUM *r = BN_new();
	BIGNUM *d = BN_new();
	BIGNUM *zero = BN_new();
	BN_zero(zero);
	BN_set_word(r, 1);
	BN_set_word(d, 1);
	BN_sub(n_1, n, d);

	do {
		BN_exp(res, r, n_1, ctx);
		BN_div(res, NULL, number, res, ctx);
		BN_sub(res, res, r);
		BN_div(d, NULL, res, n, ctx);
		BN_add(r, r, d);
	} while (BN_cmp(d, zero));

	BN_copy(res, r);
	BN_free(zero);
	BN_free(r);
	BN_free(d);
	BN_free(n_1);
}
int rsa_broadcast_cube(BIGNUM *res, BIGNUM **a, BIGNUM **n)
{
	BIGNUM *tmp = BN_new();
	BIGNUM *N= BN_new();
	BIGNUM *N_ni = BN_new();
	BIGNUM *sum = BN_new();
	BIGNUM *n_3 = BN_new();

	BN_set_word(n_3, 3);
	BN_one(N);
	BN_zero(sum);
	int i;

	for(i=0;i<3;i++)
		BN_mul(N, N, n[i], ctx);

	for(i=0;i<3;i++) {
		BN_div(N_ni, NULL, N, n[i], ctx);
		BN_mod_inverse(tmp, N_ni, n[i], ctx);
		modular_multiplicative_inverse_bignum_my(tmp, N_ni, n[i]);
		BN_mul(tmp, tmp, N_ni, ctx);
		BN_mul(tmp, tmp, a[i], ctx);
		BN_add(sum, sum, tmp);
	}

	BN_nnmod(sum, sum, N, ctx);
	nth_root_bignum(res, sum, n_3);
}

int chinese_remainder_theorem_bignum(BIGNUM *solution, BIGNUM *sol_no_mod, BIGNUM **a, BIGNUM **n, int len)
{
	int i,j;

	for(i=0;i<len;i++) {
		for(j=i+1;j<len;j++) {
			if(!check_co_prime_bignum(n[i], n[j]))
				return -1;
		}
	}

	__chinese_remainder_theorem_bignum(solution, sol_no_mod, a, n, len);
}

int check_co_prime_bignum(BIGNUM *a, BIGNUM *b)
{
	struct extended_euclid_bignum e;
	e.d = BN_new();
	e.s = BN_new();
	e.t = BN_new();

	BIGNUM *one = BN_new();
	BN_one(one);
	extended_euclid_algo_bignum(a, b, &e);
	int ret = BN_cmp(e.d, one);
	BN_free(one);
	BN_free(e.d);
	BN_free(e.s);
	BN_free(e.t);
	return (ret == 0);
}

int chinese_remainder_theorem(int *a, int *n, int len)
{
	// check if n[i] are pair-wise co-prime since it is a pre-condition
	// for the chinese remainder theorem
	int i, j;

	for(i=0;i<len;i++) {
		for(j=i+1;j<len;j++) {
			if(!check_co_prime(n[i], n[j]))
				return -1;
		}
	}

	return __chinese_remainder_theorem(a, n, len);

}
int __chinese_remainder_theorem_bignum(BIGNUM *solution, BIGNUM *sol_no_mod, BIGNUM **a, BIGNUM **n, int len)
{
	BIGNUM *N = BN_new();
	BIGNUM *N_ni = BN_new();
	BIGNUM *tmp  = BN_new();
	BN_one(N);
	BN_zero(solution);
	BN_zero(sol_no_mod);
	int i;
	struct extended_euclid_bignum e;
	e.d = BN_new();
	e.s = BN_new();
	e.t = BN_new();

	for(i=0;i<len;i++)
		BN_mul(N, N, n[i], ctx);


	for(i=0;i<len;i++) {
		BN_div(N_ni, NULL, N, n[i], ctx);
		extended_euclid_algo_bignum(n[i], N_ni, &e);
		BN_mul(tmp, a[i], e.t, ctx);
		BN_mul(tmp, tmp, N_ni, ctx);
		BN_add(solution, solution, tmp);
	}
	BN_copy(sol_no_mod, solution);
	BN_nnmod(solution, solution, N, ctx);
	BN_free(N);
	BN_free(N_ni);
	BN_free(tmp);
	BN_free(e.d);
	BN_free(e.s);
	BN_free(e.t);
}

int check_co_prime(int a, int b)
{
	struct extended_euclid e;
	extended_euclid_algo(a, b, &e);
	return (e.d == 1);
}

/** assumes that the system is sovleable with crt, aka. that the n_i are co-prime
  *
  **/
int __chinese_remainder_theorem(int *a, int *n, int len)
{
	int N = 1;
	int i = 0;
	int solution = 0;
	struct extended_euclid e;

	for(i=0;i<len;i++)
		N *= n[i];

	for(i=0;i<len;i++) {
		int N_ni = N / n[i];
		extended_euclid_algo(n[i], N_ni, &e);
		solution += a[i] * e.t * N_ni;
	}

	return modulo(solution, N);
}


#define NTH_ROOT_PRECISION 0.00001
double nth_root_wr(double x, int n)
{
	double r = 1;
	double d = 1;

	do {
		d = (x / pow(r, n - 1) - r) / n;
		r += d;
	}
	while (d > NTH_ROOT_PRECISION || d < -NTH_ROOT_PRECISION);
	return r;
}