#include "rsa.h"

void bigint_set_random_bytes(bigint_t n, size_t len) {
	int fd = open("/dev/urandom", O_RDONLY);
	read(fd, n.data, len * sizeof(uint32_t));
	close(fd);
}

void bigint_set_msb_and_lsb_to_one(bigint_t n, size_t len) {
	n.data[0] |= 1;
	n.data[len - 1] |= 1 << 31;
}

void bigint_bitwise_right_shift(bigint_t n) {
	size_t size = sizeof(uint32_t) * 8 - 1;
	for (size_t i = 0; i < n.len - 1; i++) {
		n.data[i] = n.data[i] >> 1 | (n.data[i + 1] & 1) << size;
	}
	n.data[n.len - 1] >>= 1;
}

void bigint_bitwise_left_shift(bigint_t n) {
	size_t size = sizeof(uint32_t) * 8 - 1;
	for (int i = n.len - 1; i > 0; i--) {
		n.data[i] = n.data[i] << 1 | ((n.data[i - 1] & (1 << size)) >> size);
	}
	n.data[0] <<= 1;
}

void move_bigint_bitwise_left_shift(bigint_t n, bigint_t result) {
	memcpy(result.data, n.data, n.len * sizeof(uint32_t));
	size_t size = sizeof(uint32_t) * 8 - 1;
	for (int i = result.len - 1; i > 0; i--) {
		result.data[i] = result.data[i] << 1 | ((result.data[i - 1] & (1 << size)) >> size);
	}
	result.data[0] <<= 1;
}

// Will underflow
void bigint_decrement(bigint_t n) {
	size_t cursor = 0;
	size_t size = sizeof(uint32_t) * 8;
	while (cursor < size * n.len) {
		n.data[cursor / size] = n.data[cursor / size] ^ (1 << (cursor % size));
		if (((n.data[cursor / size] >> (cursor % size)) & 1) == 0) {
			return;
		}
		cursor += 1;
	}
}

// TODO refactor/clean assume same length ?
int64_t bigint_cmp(bigint_t a, bigint_t b) {
	int cursor = a.len - 1;
	while (cursor >= 0) {
		if (a.data[cursor] > b.data[cursor]) {
			return 1;
		}
		if (b.data[cursor] > a.data[cursor]) {
			return -1;
		}
		cursor -= 1;
	}

	return 0;
}

// TODO refactor/clean assume same length ?
int bigint_dif(bigint_t a, bigint_t b) {
	int cursor = a.len;
	while (--cursor >= 0) {
		if (a.data[cursor] ^ b.data[cursor]) {
			return 1;
		}
	}

	return 0;
}

int is_not_zero(bigint_t n) {
	for (size_t i = 0; i < n.len; i++) {
		if (n.data[i]) {
			return 1;
		}
	}
	return 0;
}

// TODO check opti
void bigint_substraction(bigint_t a, bigint_t b, bigint_t borrow, bigint_t y) {
	memcpy(borrow.data, b.data, b.len * sizeof(uint32_t));
	memcpy(y.data, b.data, b.len * sizeof(uint32_t));
	while (is_not_zero(y)) {
		for (size_t i = 0; i < a.len; i++) {
			borrow.data[i] = ~a.data[i] & y.data[i];
			a.data[i] = a.data[i] ^ y.data[i];
		}
		move_bigint_bitwise_left_shift(borrow, y);
	}
}

// TODO check opti
void custom_bigint_modulo(bigint_t a, bigint_t b, bigint_t result, bigint_t mod, bigint_t borrow_sub, bigint_t y_sub) {
	bigint_set_zeros(result);
	memcpy(result.data, a.data, a.len * sizeof(uint32_t));
	if (bigint_cmp(result, b) < 0) {
		return ;
	}
	bigint_bitwise_left_shift(mod);
	while (bigint_cmp(b, mod) < 0) {
		while (bigint_cmp(result, mod) > 0) {
			bigint_bitwise_left_shift(mod);
		}
		bigint_bitwise_right_shift(mod);
		if (bigint_cmp(result, mod) > 0) {
			bigint_substraction(result, mod, borrow_sub, y_sub);
		}
	}
	while (bigint_cmp(result, b) > 0) {
		bigint_substraction(result, b, borrow_sub, y_sub);
	}
}

bigint_t bigint_new(size_t len) {
	bigint_t bigint;

	bigint.len = len;
	bigint.data = (uint32_t *)protected_malloc(len * sizeof(uint32_t));

	return bigint;
}

bigint_t bigint_zero(size_t len) {
	bigint_t bigint;

	bigint = bigint_new(len);
	for (size_t i = 0; i < len; i++) {
		bigint.data[i] = 0;
	}

	return bigint;
}

bigint_t bigint_clone(bigint_t src) {
	bigint_t dst;

	dst.len = src.len;
	dst.data = (uint32_t *)protected_malloc(src.len * sizeof(uint32_t));
	memcpy(dst.data, src.data, src.len * sizeof(uint32_t));

	return dst;
}

void bigint_destroy(bigint_t n) {
	free(n.data);
	n.data = NULL;
}

void custom_bigint_add(bigint_t a, bigint_t b, int index) {
	uint64_t carriage = 0;

	for (size_t cursor = 0; cursor < a.len; cursor++) {
		uint64_t tmp = (uint64_t)a.data[cursor] + carriage;
		if ((int)cursor - index >= 0) {
			tmp += (uint64_t)b.data[cursor - index];
		}
		a.data[cursor] = (uint32_t)tmp;
		carriage = tmp >> 32;
	}
}

void bigint_set_zeros(bigint_t n) {
	for (size_t i = 0; i < n.len; i++) {
		n.data[i] = 0;
	}
}

void custom_bigint_mul(bigint_t a, bigint_t b, bigint_t result) {
	int width = a.len * 32;
	bigint_set_zeros(result);
	bigint_t *b_tool = (bigint_t *)malloc(32 * sizeof(bigint_t));

	b_tool[0] = bigint_zero(a.len + b.len);
	bigint_set_zeros(b_tool[0]);
	memcpy(b_tool[0].data, b.data, b.len * sizeof(uint32_t));

	for (int i = 1; i < 32; i++) {
		b_tool[i] = bigint_zero(a.len + b.len);
		bigint_set_zeros(b_tool[i]);
		memcpy(b_tool[i].data, b_tool[i - 1].data, b.len * sizeof(uint32_t));
		bigint_bitwise_left_shift(b_tool[i]);
	}

	for (int cursor = 0; cursor < width; cursor++) {
		int offset = cursor % 32;
		int index = cursor >> 5;
		if (a.data[index] >> offset & 1) {
			custom_bigint_add(result, b_tool[offset], index);
		}
	}
	for (int i = 0; i < 32; i++) {
		bigint_destroy(b_tool[i]);
	}
}

// a^e mod n
// clean memory tricks !!!
void custom_bigint_pow_mod(bigint_t a, bigint_t e, bigint_t n, bigint_t result, bigint_t custom, bigint_t custom2, bigint_t mod, bigint_t borrow_sub, bigint_t y_sub) {
	bigint_set_zeros(result);
	bigint_set_zeros(custom);
	bigint_set_zeros(custom2);
	memcpy(result.data, a.data, a.len * sizeof(uint32_t));
	size_t size = sizeof(uint32_t) * 8;
	int cursor = e.len * size - 1;
	while (!(e.data[cursor / 32] & 1 << (cursor % 32))) {
		cursor--;
	}
	cursor--;
	while (cursor >= 0) {
		custom_bigint_mul(result, result, custom);
		custom_bigint_modulo(custom, n, custom2, mod, borrow_sub, y_sub);
		bigint_set_zeros(result);
		memcpy(result.data, custom2.data, custom2.len * sizeof(uint32_t));
		if (e.data[cursor / 32] & 1 << (cursor % 32)) {
			custom_bigint_mul(result, a, custom);
			custom_bigint_modulo(custom, n, custom2, mod, borrow_sub, y_sub);
			memcpy(result.data, custom2.data, custom2.len * sizeof(uint32_t));
		}
		cursor -= 1;
	}
}

void bigint_print(bigint_t n) {
	for (int i = n.len - 1; i >= 0; i--) {
		printf("bigint %ud\n", n.data[i]);
	}
}


void bulk_destroy(bigint_t x, bigint_t y, bigint_t n, bigint_t d, bigint_t two, bigint_t one, bigint_t n_minus_two, bigint_t n_minus_one) {
	bigint_destroy(x);
	bigint_destroy(y);
	bigint_destroy(n);
	bigint_destroy(d);
	bigint_destroy(two);
	bigint_destroy(one);
	bigint_destroy(n_minus_two);
	bigint_destroy(n_minus_one);
}

bigint_t bigint_prime(size_t len) {
	bigint_t n = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);

	bigint_set_random_bytes(n, len);
	bigint_set_msb_and_lsb_to_one(n, len);

	bigint_t mod = bigint_clone(n);
	bigint_t borrow_sub = bigint_clone(n);
	bigint_t y_sub = bigint_clone(n);
	bigint_t d = bigint_clone(n);
	d.data[0] -= 1;
	uint32_t s = 0;
	while (!(d.data[0] & 1)) {
		bigint_bitwise_right_shift(d);
		s += 1;
	}

	bigint_t x = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
	bigint_t y = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
	bigint_t custom = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
	bigint_t custom2 = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);

	bigint_t two = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
	bigint_t one = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
	two.data[0] = 2;
	one.data[0] = 1;
	bigint_t n_minus_two = bigint_clone(n);
	bigint_t n_minus_one = bigint_clone(n);
	n_minus_two.data[0] -= 1;
	n_minus_one.data[0] -= 1;
	bigint_decrement(n_minus_two);
	bigint_t a = bigint_zero(RSA_BLOCK_SIZE / 8 / sizeof(uint32_t) * 2);
	for (uint32_t k = 0; k < 128; k++) {
		bigint_set_zeros(a);
		while (bigint_cmp(a, two) < 0 || bigint_cmp(a, n_minus_two) > 0) {
			bigint_set_random_bytes(a, len);
		}
		custom_bigint_pow_mod(a, d, n, x, custom, custom2, mod, borrow_sub, y_sub);
		for (uint32_t i = 0; i < s; i++) {
			custom_bigint_pow_mod(x, two, n, y, custom, custom2, mod, borrow_sub, y_sub);
			if (!bigint_dif(y, one) && bigint_dif(x, one) && bigint_dif(x, n_minus_one)) {
				bulk_destroy(x, y, n, d, two, one, n_minus_two, n_minus_one);
				bigint_destroy(custom);
				bigint_destroy(custom2);
				bigint_destroy(a);
				return bigint_prime(len);
			}
			bigint_destroy(x);
			x = bigint_clone(y);
		}
		if (bigint_dif(y, one)) {
			bulk_destroy(x, y, n, d, two, one, n_minus_two, n_minus_one);
			bigint_destroy(custom);
			bigint_destroy(custom2);
			bigint_destroy(a);
			return bigint_prime(len);
		}
	}
	bulk_destroy(x, y, custom, d, two, one, n_minus_two, n_minus_one);
	bigint_destroy(custom2);
	bigint_destroy(a);
	return n;
}