import math
import hashlib

# According to NIST Special Publication 800-90A, Revision 1, this should be a cryptographically secure pseudo-random
# number generator, provided I've implemented it properly, which is of course very possible I haven't
class CSPRNG:
	def __init__(self, entropy: bytes, nonce: bytes=b'', personalization_string: bytes=b''):
		self.V = hash_df(entropy + nonce + personalization_string, 888).to_bytes(111)
		self.C = hash_df(int(0).to_bytes(0) + self.V, 888).to_bytes(111)
		self.reseed_counter = 1

	def hash_gen(self, requested_number_of_bits: int):
		m = int(math.ceil(requested_number_of_bits / 512))
		data = self.V
		w = b''
		for i in range(m):
			hasher = hashlib.sha512()
			hasher.update(data)
			w += hasher.digest()
			data = int.from_bytes(data)
			data = (data + 1) % 2 ** 888
			data = data.to_bytes(111)
		w = int.from_bytes(w)
		w = w >> (512 * m - requested_number_of_bits)
		return w

	def get_random_bytes(self, number_of_bytes: int):
		return_bytes = self.hash_gen(number_of_bytes * 8).to_bytes(number_of_bytes)
		hasher = hashlib.sha512()
		hasher.update(int(3).to_bytes(1) + self.V)
		h = hasher.digest()
		new_v = (int.from_bytes(self.V) + int.from_bytes(h) + int.from_bytes(self.C) + self.reseed_counter) % 2 ** 888
		self.V = new_v.to_bytes(111)
		self.reseed_counter += 1
		return return_bytes


# Hash derivation function as specified in section 10.3.1 of NIST Special Publication 800-90A, Revision 1
def hash_df(input_string: bytes, number_of_bits: int):
	temp = b''
	length = int(math.ceil(number_of_bits / 512))
	for i in range(length):
		hash_input = (i + 1).to_bytes(1) + number_of_bits.to_bytes(4) + input_string
		m = hashlib.sha512()
		m.update(hash_input)
		temp += m.digest()
	number = int.from_bytes(temp)
	number = number >> (512 * length - number_of_bits)
	return number