Secure Remote Password Protocol

from cryptopals_lib import *
import os, hashlib
import hmac

def generate_key_pair(prime=((1 << 1024) - 1093337), g=7):
	#generate a 1024 bit prime number for the Mod
	#prime_mod = generate_probable_prime(1024)
	prime_mod = prime

	#Genreate Secret Exponent
	secret_exp = secure_rand_between(2, prime_mod-2)
	generator = g

	#Generate the Transmit key A = G^x % P
	transmit_key = pow(generator, secret_exp, prime_mod)

	return {"generator":generator, "prime_mod":prime_mod, "transmit_key": transmit_key}, {"generator":generator, "prime_mod":prime_mod, "secret_exp": secret_exp} 


class SRPServer():
	def __init__(self, k, hash_obj=hashlib.sha256):
		self.salt_length = 16 
		self.k = k
		self.hash_obj = hash_obj

		self.database = {}

		self.auth_database = {}

		#Generate server Keys
		server_public_key, server_private_key = generate_key_pair()

	def register_user(self, username, salt, verifier):
		if username not in self.database:
			self.database[username] = {"salt": salt, "verifier": verifier}
		else:
			print("Error: Already in Database")

	def authenticate_init(self, username, client_public_key):
		if username not in self.database:
			print(f"Error: {username} is not registed")
			return
		else:
			salt, verifier = self.database[username]["salt"], self.database[username]["verifier"]

			#B = (k * v + pow(g, b, N)) % N
			session_public_key, session_private_key = generate_key_pair()
			tmp_server_key = ((k * verifier) + session_public_key["transmit_key"]) % prime

			self.auth_database[username] = {"client_public_key": client_public_key, "tmp_server_key": tmp_server_key, "session_private_key": session_private_key}

			return salt, tmp_server_key

	def authenticate_verify(self, username, client_proof):
		if username not in self.database:
			print(f"Error: {username} is not registed")
			return
		else:
			#Pull data from database
			salt, verifier = self.database[username]["salt"], self.database[username]["verifier"]
			client_public_key, tmp_server_key, session_private_key = self.auth_database[username]["client_public_key"], self.auth_database[username]["tmp_server_key"], self.auth_database[username]["session_private_key"]

			#Calculate scrambling paramater
			u = hash_obj(int_to_bytes(client_public_key) + b":" + int_to_bytes(tmp_server_key)).digest()


			#Compute Session Keys
			tmp = pow(verifier, bytes_to_int(u), prime)
			session_key_s = pow(client_public_key * tmp, session_private_key["secret_exp"], prime)
			session_key_k = hash_obj(int_to_bytes(session_key_s)).digest()

			print(f"Server: session_key_k: {session_key_k.hex()}")


			#Verify client_proof
			#M_c = H(H(N) ^ H(g), H(I), s, A, B, K_c)
			tmp = fixedlen_xor(hash_obj(int_to_bytes(prime)).digest(), hash_obj(int_to_bytes(g)).digest())
			client_proof_verify = hash_obj(tmp + b":" + hash_obj(username).digest() + b":" + salt + b":" 
				+ int_to_bytes(client_public_key) + b":" + int_to_bytes(tmp_server_key) + b":" + session_key_k).digest()

			print(f"client_proof:        {client_proof.hex()}")
			print(f"client_proof_verify: {client_proof_verify.hex()}")


			if client_proof == client_proof_verify:
				#Compute M_s
				M_s = hash_obj(int_to_bytes(client_public_key) + b":" + client_proof + b":" + session_key_k).digest()
				#print(f"Server: M_s: {M_s.hex()}")

				return M_s
			else:
				print("Error: Client Verification is invalid")
				return 



	def generate_salt(self):
		return os.urandom(self.salt_length)
		
		
#Globals
username = b"admin"
password = b"Password123"
k = 3
hash_obj = hashlib.sha256
prime=((1 << 1024) - 1093337)
g=7


if __name__ == '__main__':
	#Exchange Public Keys
	server = SRPServer(k)


	#### Register User
	print(f"Register User: {username}")
	salt = os.urandom(16)
	x = hash_obj(salt + b":" + username + b":" + password).digest()
	verifier = pow(g, bytes_to_int(x), prime)

	server.register_user(username, salt, verifier)
	print()


	#### Authenticate User
	print(f"Authenticate User: {username}")

	#Generate Temp Key
	client_public_key, client_private_key = generate_key_pair()
	#print(f"client_public_key: {client_public_key}")

	salt, tmp_server_key = server.authenticate_init(username, client_public_key["transmit_key"])

	#Calculate scrambling paramater
	u = hash_obj(int_to_bytes(client_public_key["transmit_key"]) + b":" + int_to_bytes(tmp_server_key)).digest()

	#Compute the Session Key
	x = bytes_to_int(hash_obj(salt + b":" + username + b":" + password).digest())
	session_key_s = pow(tmp_server_key - k * pow(g, x, prime), client_private_key["secret_exp"] + bytes_to_int(u) * x, prime)
	session_key_k = hash_obj(int_to_bytes(session_key_s)).digest()

	#print(f"session_key_s: {session_key_s}")
	print(f"Client: session_key_k: {session_key_k.hex()}")

	#H(H(N) ^ H(g), H(I), s, A, B, K_c)
	tmp = fixedlen_xor(hash_obj(int_to_bytes(prime)).digest(), hash_obj(int_to_bytes(g)).digest())
	client_proof = hash_obj(tmp + b":" + hash_obj(username).digest() + b":" + salt + b":" 
		+ int_to_bytes(client_public_key["transmit_key"]) + b":" + int_to_bytes(tmp_server_key) + b":" + session_key_k).digest()

	#Send Client Verification
	server_proof = server.authenticate_verify(username, client_proof)

	#Verify Server Verification
	#H(A, M_c, K_s)
	server_proof_verify = hash_obj(int_to_bytes(client_public_key["transmit_key"]) + b":" + client_proof + b":" + session_key_k ).digest()

	print(f"server_proof:        {server_proof.hex()}")
	print(f"server_proof_verify: {server_proof_verify.hex()}")


	#Verify Server Key