Diffie Hellman

import secrets

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

	#Generate Secret Exponent
	secret_exp = 0
	while secret_exp < 2:
		secret_exp = secrets.randbelow(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} 

if __name__ == '__main__':
	#Genearte User1 Keys
	user1_public_key, user1_private_key = generate_key_pair()

	#Generate User2 Keys
	user2_public_key, user2_private_key = generate_key_pair()

	#Generate Shared Secret from the other user's Public Key
	user1_shared_secert = pow(user2_public_key["transmit_key"], user1_private_key["secret_exp"], user1_private_key["prime_mod"])
	user2_shared_secert = pow(user1_public_key["transmit_key"], user2_private_key["secret_exp"], user2_private_key["prime_mod"])

	assert user1_shared_secert == user2_shared_secert

from collections.abc import Iterator
from cryptopals_lib import *
from sympy.ntheory import factorint
from sympy import isprime, gcd
import hashlib, random
from hmac import hmac
from functools import reduce

prime = 7199773997391911030609999317773941274322764333428698921736339643928346453700085358802973900485592910475480089726140708102474957429903531369589969318716771
generator = 4565356397095740655436854503483826832136106141639563487732438195343690437606117828318042418238184896212352329118608100083187535033402010599512641674644143
order = 236234353446506858198510045061214171961
j = 30477252323177606811760882179058908038824640750610513771646768011063128035873508507547741559514324673960576895059570

def extended_gcd(a: int, b: int) -> tuple[int, tuple[int, int]]:
	"""
	Extended Euclidean algorithm
	:return: ( 'gcd' - the resulting gcd,
			   'coeffs' - Bézout coefficients )
	"""
	old_r, r = a, b
	old_s, s = 1, 0
	old_t, t = 0, 1

	while r != 0:
		quotient = old_r // r
		old_r, r = r, old_r - quotient * r
		old_s, s = s, old_s - quotient * s
		old_t, t = t, old_t - quotient * t

	return old_r, (old_s, old_t)

def chinese_remainder(n_list: list[int], a_list: list[int]) -> int:
	"""
	Solution of the system:
	x = a1 (mod n1)
	x = a2 (mod n2)
	...
	x = ak (mod k)

	Such that 0 <= x < N,
	where N = n1 * n2 * ... * nk

	https://en.wikipedia.org/wiki/Chinese_remainder_theorem#Existence_(direct_construction)
	"""
	x = 0
	N = reduce(lambda a, b: a * b, n_list)
	for ni, ai in zip(n_list, a_list):
		Ni = N // ni
		_, (Mi, _) = extended_gcd(Ni, ni)
		x += ai*Mi*Ni

	return x % N


def trial_division(n: int) -> Iterator[int]:
	"""
	Basic Integer Factorization Algorithm.
	https://en.wikipedia.org/wiki/Trial_division
	"""
	while n % 2 == 0:
		yield 2
		n //= 2

	f = 3
	while f * f <= n:
		if n % f == 0:
			yield f
			n //= f
		else:
			f += 2

	if n != 1:
		yield n

def get_shared_key(transmit_key, user_private_key):
	user_shared_secret = pow(transmit_key, user_private_key, prime)
	m = b"crazy flamboyant for the rap enjoyment"
	digest = hmac(int_to_bytes(user_shared_secret), m, hashlib.sha256)
	return m, digest
	

def find_element_of_order(order: int, p: int) -> int:
	"""
	return element h of order r.
	h := rand(1, p)^((p-1)/r) mod p
	"""
	if (p-1) % order != 0:
		raise ValueError('r must divide p-1')

	while True:
		h = pow(random.randint(1, p-1), (p-1)//order, p)
		if h != 1:
			return h

if __name__ == '__main__':
	#Generate Bob Keys

	#Because the generator that we have picked is a small subgroup the order of the subgroup is also smaller. 
	#This means that the random even if chosen to be larger can be reduced to 
	bob_private_key = secure_rand_between(2, order-1)
	bob_private_key2 = secure_rand_between(2, prime-2)

	bob_public_key  = pow(generator, bob_private_key, prime)
	bob_public_key2 = pow(generator, bob_private_key2, prime)

	#The Public Key even though it is bigger than the order will be the same when it is computed
	bob_private_key3 = bob_private_key2 % order
	bob_public_key3 = pow(generator, bob_private_key3, prime)

	# print(f"bob_private_key: {bob_private_key}")
	# print(f"bob_private_key2: {bob_private_key2}")
	# print(f"bob_private_key3: {bob_private_key3}")
	# print(f"bob_public_key: {bob_public_key}")
	# print(f"bob_public_key2: {bob_public_key2}")
	# print(f"bob_public_key3: {bob_public_key3}")
	assert bob_public_key2 == bob_public_key3


	### Attack Start
	print(f"Attack: {bob_private_key}")


	#Find the order. Amusing the prime_mod is prime then the order is P-1
	#If it is not prime then P-1 = cofactor*order
	#So we factor P-1 and see if any of the factors are the order
	prime_order = prime - 1

	#order_factors = factorint(prime_order)
	count = 1
	order_factor_list = []
	message_factor_list = []


	# Use the prime_factors_iterator to factor prime_order step by step
	for factor in trial_division(prime_order):
		#Ignore Repeated factors
		if factor in order_factor_list:
			continue

		count *= factor

		# Perform the attack for each factor incrementally
		bad_key_candidate = find_element_of_order(factor, prime)
		print(f"Factor: {factor} bad_key_candidate: {bad_key_candidate}")

		# Get Hmac Message with shared key as HMAC Key
		message, hmac_data = get_shared_key(bad_key_candidate, bob_private_key)

		# Bruteforce Key
		for private_key_test in range(factor):

			possible_shared_secret = int_to_bytes(pow(bad_key_candidate, private_key_test, prime))

			digest = hmac(possible_shared_secret, message, hashlib.sha256)

			if digest == hmac_data:
				order_factor_list.append(factor)
				message_factor_list.append(private_key_test)
				#print("Found HMAC match")
				break

		if count > order:
			break

	# Once all factors are processed, use Chinese remainder theorem
	a = chinese_remainder(order_factor_list, message_factor_list)
	print(f"Recovered shared secret: {a}")

import secrets

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 = 0
	while secret_exp < 2:
		secret_exp = secrets.randbelow(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} 


def brute_dlp(g, y, p):
	mod_size = len(bin(p-1)[2:])

	print("[+] Using Brute Force algorithm to solve DLP")
	print("[+] Modulus size: " + str(mod_size) + ". Warning! Brute Force is not efficient\n")

	sol = pow(g, 2, p)
	if sol == y:
		return 2
	if y == 1:
		return p-1
	if y == g:
		return 1
	for i in range(3, p-1):
		sol = sol*g % p
		if sol == y:
			return i
	return None    

if __name__ == '__main__':
	#Genearte User1 Keys
	user1_public_key, user1_private_key = generate_key_pair(prime=3875972899)


	test = brute_dlp(user1_public_key["generator"], user1_public_key["transmit_key"], user1_public_key["prime_mod"])
	print(test)
#[+] Using Brute Force algorithm to solve DLP
#[+] Modulus size: 32. Warning! Brute Force is not efficient

#96057192
#[Finished in 13.8s]

import secrets
import math

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 = 0
	while secret_exp < 2:
		secret_exp = secrets.randbelow(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} 


def bsgs_dlp(g, y, p):
	mod_size = len(bin(p-1)[2:])

	print("[+] Using BSGS algorithm to solve DLP")
	print("[+] Modulus size: " + str(mod_size) + ". Warning! BSGS not space efficient\n")

	m = math.ceil(math.sqrt(p-1))

	# Make a lookup table for 1,j
	lookup_table = {pow(g, j, p): j for j in range(m)}

	# Precompute g^(-m) so it is easy to substitute in every lookup
	c = pow(g, m*(p-2), p)

	# Giant Steps
	for i in range(m):
		#Calculate y*g^(-m*i) = y * g^-m * g^i = y * c * g^i = y * c ^ i
		temp = (y*pow(c, i, p)) % p

		#Check if value is in lookup table and get the index
		if temp in lookup_table:
			# x found
			return i*m + lookup_table[temp]
	return None

  

if __name__ == '__main__':
	#Generate User1 Keys
	user1_public_key, user1_private_key = generate_key_pair(prime=3875972899)
	#user1_public_key, user1_private_key = generate_key_pair(prime=13091037018084742351)


	test = bsgs_dlp(user1_public_key["generator"], user1_public_key["transmit_key"], user1_public_key["prime_mod"])

	assert user1_private_key["secret_exp"] == test
	print(f"Found exponent: {test}")
#[+] Using BSGS algorithm to solve DLP
#[+] Modulus size: 32. Warning! BSGS not space efficient

#Found exponent: 3238504688
#[Finished in 269ms]

f"""
This is a simple, yet straight forward implementation of Pollard's rho algorithm for discrete logarithms
It computes X such that G^X = H mod P.

p does not need to be a safe prime.
"""
from gmpy2 import invert, gcd


def xab(x, a, b, params):
	"""
	Pollard Step
	:param x:
	:param a:
	:param b:
	:return:
	"""
	G, H, P, Q = params
	sub = x % 3 # Subsets

	if sub == 0:
		x = x*G % P
		a = (a+1) % Q
	elif sub == 1:
		x = x*H % P
		b = (b+1) % Q
	else:
		# sub == 2:
		x = x*x % P
		a = a*2 % Q
		b = b*2 % Q

	return x, a, b


def pollard(G, H, P):

	# P: prime
	# H:
	# G: generator
	Q = (P - 1)  # multiplicative sub group order

	x, a, b = 1, 0, 0
	X, A, B = x, a, b

	# for i in range(1, P):
	while True:
		# Hedgehog
		x, a, b = xab(x, a, b, (G, H, P, Q))

		# Hare
		X, A, B = xab(X, A, B, (G, H, P, Q))
		X, A, B = xab(X, A, B, (G, H, P, Q))

		if x == X:
			break

	num = a-A
	denum = B-b

	# It is necessary to compute the inverse to properly compute the fraction mod q
	# find gcd before solving linear equation denum*res = num mod Q
	gcdz = gcd(num, denum, Q)
	# print(f'gcds = {gcdz}')
	if gcdz == 1:  # standard solution
		res = invert(denum, Q)*num % Q  # divm(num, denum, Q)  # (inverse(denom, Q) * nom) % Q
	else:
		# needs a little bit more work
		# divide by gcd
		modz = Q//gcdz
		num = num//gcdz
		denum = denum//gcdz

		# baseline solution
		res0 = invert(denum, modz)*num % modz
		# check in solutions of the shape (denum/gcd)*res = (num/gcd) mod (Q/gcd) + k * (Q/gcd) (k in [0, gcd[) 
		for k in range(gcdz):
			res = res0+k*modz
			if verify(G, H, P, res):
				break
	return res


def verify(g, h, p, x):
	"""
	Verifies a given set of g, h, p and x
	:param g: Generator
	:param h:
	:param p: Prime
	:param x: Computed X
	:return:
	"""
	return pow(g, x, p) == h



import secrets

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

	#Generate Secret Exponent
	secret_exp = 0
	while secret_exp < 2:
		secret_exp = secrets.randbelow(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} 


if __name__ == '__main__':

	#user1_public_key, user1_private_key = generate_key_pair(prime=3875972899)
	user1_public_key, user1_private_key = generate_key_pair(prime=13091037018084742351)


	test = pollard(user1_public_key["generator"], user1_public_key["transmit_key"], user1_public_key["prime_mod"])
	print(f"Found Possible answer: {test}")

	print("Validates: ", verify(user1_public_key["generator"], user1_public_key["transmit_key"], user1_public_key["prime_mod"], test))

#Found Possible answer: 6244664414481744603
#Validates:  True
#[Finished in 1461.2s]

(g^a)^b mod p XOR nonce