#include "lib6.h"
#include "lib5.h"
#include "lib4.h"
#include "lib3.h"
#include "lib2.h"
#include "lib.h"
#include "util/set.h"
#include <openssl/sha.h>


int rsa_sign_bignum(BIGNUM *message, BIGNUM *signed_message, struct rsa_key_bignum *private)
{
	rsa_encrypt_bignum(message, signed_message, private);
}

int rsa_verify_bignum(BIGNUM *signed_message, BIGNUM *org_message, struct rsa_key_bignum *public)
{
	BIGNUM *res = BN_new();
	int ret = -1;

	rsa_decrypt_bignum(signed_message, res, public);
	ret = BN_cmp(res, org_message);
	printf("\nverfied mess ret: %i, message:\n", ret);
	BN_print(out, res);
	printf("\n");
	BN_free(res);

	return ret == 0;
}
/**
  * construct a VALID pkcs_padding
  **/
void pkcs1_5_padding(char *message, char *result, unsigned int target_length_byte)
{
	SHA_CTX sha1;
	char sha1_hash[20];
	int i;

	memset(result, 0xff, target_length_byte);
	result[0] = 0x00;
	result[1] = 0x02;
	result[target_length_byte-21] = 0x00;	

	// TODO ASN.1 things

	SHA1_Init(&sha1);
	SHA1_Update(&sha1, message, strlen(message));
	SHA1_Final(sha1_hash, &sha1);

	for(i = 20;i>0;i--)
		result[target_length_byte-i] = sha1_hash[20-i];
}

int pkcs1_5_padding_verify(char *to_verify, int len, char *message)
{
	char result[1024/8];
	int i;

	// construct the padding how the expect it and than compare
	pkcs1_5_padding(message, result, 1024/8);
	// printf both paddings
	char buf[(1024/8)*2];
	hex_binary_to_string(result, buf, 1024/8);
	printf("expected padding:\n%s\n", buf);
	hex_binary_to_string(to_verify, buf, len);
	printf("got:\n%s\n", buf);

	return memcmp(to_verify, result, 128) == 0;
}

int shitty_pkcs1_5_padding_verify(char *to_verify, int len, char *message)
{
	int i = 2;
	SHA1Context sha1;
	char sha1_hash[20];

	if (len < 2 && to_verify[0] != 0x00 && to_verify[1] != 0x01)
		return 0;

	// search for the next 0x00 no matter what's in between
	while(to_verify[i] != 0x00)
		i++;

	i++;
	// TODO check asn.1 things
	// verfiy the hash
	SHA1Reset(&sha1);
	SHA1Input(&sha1, message, strlen(message));
	SHA1Result(&sha1);
	memcpy(sha1_hash, &(sha1.Message_Digest), 20);
	
	int j;
	for(j=0;j<20;j++, i++) {
		if (to_verify[i] != sha1_hash[j])
			return 0;
	}

	char buf[(1024/8)*2];
	hex_binary_to_string(to_verify, buf, len);
	printf("got:\n%s\n", buf);
	return 1;
}

void init_dsa_pub_param(struct dsa_public_params *p)
{
	char *p_str = "800000000000000089e1855218a0e7dac38136ffafa72eda7"
	"859f2171e25e65eac698c1702578b07dc2a1076da241c76c6"
	"2d374d8389ea5aeffd3226a0530cc565f3bf6b50929139ebe"
	"ac04f48c3c84afb796d61e5a4f9a8fda812ab59494232c7d2"
	"b4deb50aa18ee9e132bfa85ac4374d7f9091abc3d015efc87"
	"1a584471bb1";

	char *q_str = "f4f47f05794b256174bba6e9b396a7707e563c5b";

	char *g_str = "5958c9d3898b224b12672c0b98e06c60df923cb8bc999d119"
	"458fef538b8fa4046c8db53039db620c094c9fa077ef389b5"
	"322a559946a71903f990f1f7e0e025e2d7f7cf494aff1a047"
	"0f5b64c36b625a097f1651fe775323556fe00b3608c887892"
	"878480e99041be601a62166ca6894bdd41a7054ec89f756ba"
	"9fc95302291";

	p->p = BN_new();
	p->q = BN_new();
	p->g = BN_new();

	BN_hex2bn(&p->p, p_str);
	BN_hex2bn(&p->q, q_str);
	BN_hex2bn(&p->g, g_str);
	p->bits = BN_num_bytes(p->q)*16;
}

void dsa_compute_per_user_keys(struct dsa_public_params *pub_param, struct
		dsa_per_user_param *priv_param)
{
	BN_pseudo_rand(priv_param->private, pub_param->bits, -1, -1);	
	BN_mod(priv_param->private, priv_param->private, pub_param->q, ctx);
	BN_mod_exp(priv_param->public, pub_param->g, priv_param->private,
			pub_param->p, ctx);
}

void dsa_sign(char *mess, struct dsa_public_params *pub_param, struct dsa_per_user_param
		*priv_param, BIGNUM *k)
{
	// random per message value
	BIGNUM *k_1 = BN_new();
	BIGNUM *hash_bn = BN_new();
	BIGNUM *tmp = BN_new();
	SHA_CTX sha1;
	char sha1_hash[20];

	SHA1_Init(&sha1);
	SHA1_Update(&sha1, mess, strlen(mess));
	SHA1_Final(sha1_hash, &sha1);
	BN_bin2bn(sha1_hash, 20, hash_bn);
	BN_zero(priv_param->r);
	BN_zero(priv_param->s);

	//while(BN_is_zero(priv_param->r) || BN_is_zero(priv_param->s)) {
		if(BN_is_zero(k)) {
			BN_pseudo_rand(k, pub_param->bits, -1, -1);
			BN_mod(k, k, pub_param->q, ctx);
		}
		BN_mod_exp(priv_param->r, pub_param->g, k, pub_param->p, ctx);
		BN_mod(priv_param->r, priv_param->r, pub_param->q, ctx);

		BN_mod_mul(tmp, priv_param->private, priv_param->r, pub_param->q, ctx);
		BN_mod_add(tmp, hash_bn, tmp, pub_param->q, ctx);
		BN_mod_inverse(k_1, k, pub_param->q, ctx);
		BN_mod_mul(priv_param->s, k_1, tmp, pub_param->q, ctx);
	//}
}


int dsa_verify(char *mess, struct dsa_public_params *pub, struct dsa_per_user_param
		*priv)
{
	SHA_CTX sha1;
	char sha1_hash[20];
	BIGNUM *w = BN_new();
	BIGNUM *u1 = BN_new();
	BIGNUM *u2 = BN_new();
	BIGNUM *tmp1 = BN_new();
	BIGNUM *tmp2 = BN_new();
	BIGNUM *v = BN_new();
	BIGNUM *hash_bn = BN_new();

	BN_mod_inverse(w, priv->s, pub->q, ctx);

	SHA1_Init(&sha1);
	SHA1_Update(&sha1, mess, strlen(mess));
	SHA1_Final(sha1_hash, &sha1);
	BN_bin2bn(sha1_hash, 20, hash_bn);
	
	BN_mod_mul(u1, w, hash_bn, pub->q, ctx);
	BN_mod_mul(u2, priv->r, w, pub->q, ctx);

	BN_mod_exp(tmp1, pub->g, u1, pub->p, ctx);
	BN_mod_exp(tmp2, priv->public, u2, pub->p, ctx);

	BN_mod_mul(v, tmp1, tmp2, pub->p, ctx);
	BN_mod(v, v, pub->q, ctx);

	return BN_cmp(v, priv->r);
}

void dsa_recover_x_from_known_k(struct dsa_public_params *pub, BIGNUM *k,
		struct dsa_per_user_param *priv, BIGNUM *mess_hash)
{
	BIGNUM *r_1 = BN_new();
	BN_mod_inverse(r_1, priv->r, pub->q, ctx);
	BN_mod_mul(priv->private, k, priv->s, pub->q, ctx);
	BN_mod_sub(priv->private, priv->private, mess_hash, pub->q, ctx);
	BN_mod_mul(priv->private, priv->private, r_1, pub->q, ctx);
}

void dsa_recover_k_from_repeated_nonce(BIGNUM *mess1_hash, BIGNUM *mess2_hash,
		BIGNUM *s1, BIGNUM *s2, struct dsa_public_params *pub,
		struct dsa_per_user_param *priv, BIGNUM *k)
{

	BIGNUM *diff1 = BN_new();

	BN_mod_sub(diff1, mess1_hash, mess2_hash, pub->q, ctx);
	BN_mod_sub(k, s1, s2, pub->q, ctx);
	BN_mod_inverse(k, k, pub->q, ctx);
	BN_mod_mul(k, k, diff1, pub->q, ctx);
	printf("recoverd k is: \n");
	BN_print(out, k);
}

void dsa_generate_magic_signature(struct dsa_public_params *pub, struct dsa_per_user_param
		*priv, BIGNUM *mess_hash)
{
	BIGNUM *tmp = BN_new();

	BN_mod_exp(tmp, priv->public, mess_hash, pub->p, ctx);
	BN_mod(priv->r, tmp, pub->q, ctx);

	BN_mod_inverse(tmp, mess_hash, pub->q, ctx);
	BN_mod_mul(priv->s, priv->r, tmp, pub->q, ctx);
	
}

int rsa_parity_orcale(BIGNUM *message, struct rsa_key_bignum *private)
{
	BIGNUM *decrypted = BN_new();
	// decrypt and check last bit
	rsa_decrypt_bignum(message, decrypted, private);
	return BN_is_odd(decrypted);
}

int rsa_bleichenbacher_orcale(BIGNUM *ciphertext, struct rsa_key_bignum *priv)
{
	int ret = 0;
	BIGNUM *plaintext = BN_new();
	rsa_decrypt_bignum(ciphertext, plaintext, priv);
	char *p = malloc(BN_num_bytes(plaintext));
	BN_bn2bin(plaintext, p);

	if(p[0] == 0x02)
		ret = 1;

	free(p);
	BN_free(plaintext);
	return ret;
}

int bleichenbacher_prepare(BIGNUM *c_0, struct rsa_key_bignum *public,
		struct bb_attack *b)
{
	b->_2a = BN_new();
	b->one = BN_new();
	b->two = BN_new();
	b->three = BN_new();
	b->B = BN_new();
	b->_2B = BN_new();
	b->_3B = BN_new();
	b->c_0 = BN_new();
	BIGNUM *tmp2 = BN_new();

	BN_set_word(b->one, 1);
	BN_set_word(b->two, 2);
	BN_set_word(b->three, 3);
	BN_copy(b->c_0, c_0);

	unsigned long long k = (BN_num_bytes(public->modulo) - 2) * 8;
	BN_set_word(tmp2, k);
	BN_mod_exp(b->B, b->two, tmp2, public->modulo, ctx);
	BN_mod_mul(b->_3B, b->B, b->three, public->modulo, ctx);
	BN_mod_mul(b->_2B, b->B, b->two, public->modulo, ctx);
	BN_div(b->_2a, NULL, public->modulo, b->_3B, ctx);
	printf("\n2a\n");
	BN_print(out, b->_2a);
	printf("\n");
}

int bleichenbacher_step_2a(BIGNUM *s_i, struct rsa_key_bignum *public,
		struct rsa_key_bignum *private, struct bb_attack *b)
{
	printf("bleichenbacher attack: step 2a\n");
	BIGNUM *new_c = BN_new();

	// start from 2a = c_0/3B
	BN_copy(s_i, b->_2a);
	printf("start with s_i:\n");
	BN_print(out, s_i);

	while(1) {
		// c_0 * s^e = m_0 * s
		BN_mod_exp(new_c, s_i, public->exponent, public->modulo, ctx);
		BN_mod_mul(new_c, new_c, b->c_0, public->modulo, ctx);
		// asking orcale if this ciphertext has valid pkcs
		if(rsa_bleichenbacher_orcale(new_c, private)) {
			break;
		}
		// try next number
		BN_mod_add(s_i, s_i, b->one, public->modulo, ctx);
	}
	printf("\ns_i from step 2a:\n");
	BN_print(out, s_i);
}
/**
  * when methods get called, s_i contains s_i-1
  **/
int bleichenbacher_step_2c(struct interval *m, BIGNUM *s_i, struct rsa_key_bignum *public,
	struct rsa_key_bignum *private, struct bb_attack *b)
{
	BIGNUM *r_i = BN_new();
	BIGNUM *r_i_start = BN_new();
	BIGNUM *s_i_1 = BN_new();
	BIGNUM *tmp= BN_new();
	BIGNUM *si_upper = BN_new();
	BIGNUM *si_lower = BN_new();
	BIGNUM *new_c = BN_new();

	BN_copy(s_i_1, s_i);
	// TODO set it here to look if it works for the first step
	BN_copy(m->lower, b->_2B);
	BN_copy(m->upper, b->_3B);

	BN_mul(tmp, s_i_1, m->upper, ctx);
	BN_sub(tmp, tmp, b->_2B);
	BN_mul(tmp, tmp, b->two, ctx);
	BN_div(r_i_start, NULL, tmp, public->modulo, ctx);
	printf("tmp:\n");
	BN_print(out, tmp);

	printf("\nr_i start is:\n");
	BN_print(out, r_i_start);

	printf("\nout out out\n");
	// compute s_i intervall, increment r_i in this loop
	while(1) {
		BN_mul(tmp, r_i, public->modulo, ctx);
		BN_add(si_lower, b->_2B, tmp);
		BN_div(si_lower, NULL, si_lower, m->upper, ctx);
		BN_add(si_upper, b->_3B, tmp);
		BN_div(si_upper, NULL, si_upper, m->lower, ctx);
		printf("\nsi_lower\n");
		BN_print(out, si_lower);
		printf("\nsi_upper\n");
		BN_print(out, si_upper);
		// try all s_i in [si_lower, si_upper]
		while(BN_cmp(s_i, si_upper) != 1) {
			// for all s_i between [si_lower, si_upper]
			BN_mod_exp(tmp, s_i, public->exponent, public->modulo, ctx);
			BN_mod_mul(new_c, tmp, b->c_0, public->modulo, ctx);
			// try if we get valid pkcs padding
			if(rsa_bleichenbacher_orcale(new_c, private))
				goto found_si;

			BN_add(s_i, s_i, b->one);
		}
	}
		

found_si:
	printf("found s_i in step 2c:\n");
	BN_print(out, s_i);
}

/**
  * Narrowing down the solution
 **/

/**
int bleichenbacher_step_3(BIGNUM *s_i, set_t *m_i1)
{
	struct list *intervall, *head;
	// compute m_i from m_i-1
	LIST_TO_END(intervall, head) {
		
	}

}
**/

BIGNUM *BN_max(BIGNUM *e1, BIGNUM *e2)
{
	if(BN_cmp(e1, e2) == 1)
		return e1;
	else
		return e2;
}

/**
  * returns 1 if the intervall have been merged (one intervall),
  * 2 if intevalls not overlapping (two intervalls)
  **/
int iv_merge2(struct interval *left, struct interval *right)
{
	if(!(BN_cmp(left->upper, right->lower) == -1)) {
		BN_copy(left->upper, BN_max(left->upper, right->upper));
		return 1;
	}

	return 2;
}

/**
  * returns 1 if all three intervalls have merged to one, res in left
  * returns 2 if the left + middle has been merged and the right not, res in left
  * returns 3 if none have been merged
  * retuern 4 if the middle and the right has been merged and left not, res in middle
  *
  **/

int iv_merge3(struct interval *left, struct interval *middle, struct interval *rigth)
{
	if(iv_merge2(left, middle) == 1) {
		if (iv_merge2(left, rigth) == 1)
			return 1;
		else
			return 2;
	} else {
		if(iv_merge2(middle, rigth) == 1)
			return 4;
		else
			return 3;
	}
		
}

void iv_printf(struct interval *e1)
{
	printf("lower:");
	BN_print(out, e1->lower);
	printf("\nupper:");
	BN_print(out, e1->upper);
	printf("\n");
}