ElGamal

import secrets
from cryptopals_lib import bytes_to_int, int_to_bytes

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 gen_random(low, high):
	random = 0
	while random < low:
		random = secrets.randbelow(high)

	return random

def encrypt_message(to_public_key, from_private_key, message):
	message = bytes_to_int(message)

	#Generate One Time Key
	ephemeral_key = gen_random(1, to_public_key["prime_mod"]-2)

	#Generate Ciphered Key = G^y % P
	ciphertext_ephemeral_key = pow(to_public_key["generator"], ephemeral_key, to_public_key["prime_mod"])

	#Generate Secret S = (G^x)^y %p
	shared_secret = pow(to_public_key["transmit_key"], ephemeral_key, to_public_key["prime_mod"])

	#Generate Ciphertext message c2 = (m*s) % p
	cipheretext_message = (message * shared_secret) % to_public_key["prime_mod"]

	return (ciphertext_ephemeral_key, cipheretext_message)


def decrypt_message(ephemeral_key_ciphertext, message_ciphertext, to_private_key):
	#Generate Shared Secret
	shared_secret = pow(ephemeral_key_ciphertext, to_private_key["secret_exp"], to_private_key["prime_mod"])

	#Mod inverse
	shared_secret_inverse = pow(shared_secret, -1, to_private_key["prime_mod"])

	#Decrypt Message
	return (message_ciphertext * shared_secret_inverse) % to_private_key["prime_mod"]



if __name__ == '__main__':
	message = b"Hello World"

	#Generate User1 Ephemeral Key
	user1_public_key, user1_private_key = generate_key_pair()

	#Generate User2 Ephemeral Key
	user2_public_key, user2_private_key = generate_key_pair()

	#Message sent from User1 -> User2
	print("Message: {}".format(message))
	ephemeral_key_ciphertext, message_ciphertext = encrypt_message(user2_public_key, user1_private_key, message)
	print("Encrypted Message: {}".format(ephemeral_key_ciphertext, message_ciphertext))

	#Message decrypred by User2
	decrypted_int_message = decrypt_message(ephemeral_key_ciphertext, message_ciphertext, user2_private_key)
	decrypted_message =int_to_bytes(decrypted_int_message)
	print("Decrypted Message: {}".format(decrypted_message))

	assert(message == decrypted_message)

B = 1044388881413152506679602719846529545831269060992135009022588756444338172022322690710444046669809783930111585737890362691860127079270495454517218673016928427459146001866885779762982229321192368303346235204368051010309155674155697460347176946394076535157284994895284821633700921811716738972451834979455897010306333468590751358365138782250372269117968985194322444535687415522007151638638141456178420621277822674995027990278673458629544391736919766299005511505446177668154446234882665961680796576903199116089347634947187778906528008004756692571666922964122566174582776707332452371001272163776841229318324903125740713574141005124561965913888899753461735347970011693256316751660678950830027510255804846105583465055446615090444309583050775808509297040039680057435342253926566240898195863631588888936364129920059308455669454034010391478238784189888594672336242763795138176353222845524644040094258962433613354036104643881925238489224010194193088911666165584229424668165441688927790460608264864204237717002054744337988941974661214699689706521543006262604535890998125752275942608772174376107314217749233048217904944409836238235772306749874396760463376480215133461333478395682746608242585133953883882226786118030184028136755970045385534758453246
c = 626538023179183383530326775878874061243537575637678230775990839517480703838547661223043351239918394865532739593994849154430591305198322288328417436892902443505917379006443598061697092201021536603987285163299838045075649512028921956548205676208712226470308113290505059049053554206416767772651512039632936795232724801922443388952796545084060280910673215063573749033668717290815841834685051547255093401042248880082370281195663905731764841136214458769095907569163184276786105979509363239839613385785232307550226120692268670717158539363476974302593377514937470953649913863219081380729197866499068054753726732556988368232198510160183731571741889761632557457157261911256679728707526809880224074940827825962124087677609578885077944156589121586951372962667391498824853627326311588259929756093569512874696625507873168095797130866697040006079191565592914811736694534863664544646201900855732710994613768016495896341014229622129847800138058041138449691835652916517559125407914079964177419733079221280850917496784527992741048638603428068452948455069671414667008770873913118369984179054902139270844669700487715140145038503768704763641223894377872853332484296603268546135216483844722413025955229649646292980986542646384960063297213437076185709400450
p = 1044388881413152506679602719846529545831269060992135009022588756444338172022322690710444046669809783930111585737890362691860127079270495454517218673016928427459146001866885779762982229321192368303346235204368051010309155674155697460347176946394076535157284994895284821633700921811716738972451834979455897010306333468590751358365138782250372269117968985194322444535687415522007151638638141456178420621277822674995027990278673458629544391736919766299005511505446177668154446234882665961680796576903199116089347634947187778906528008004756692571666922964122566174582776707332452371001272163776841229318324903125740713574141005124561965913888899753461735347970011693256316751660678950830027510255804846105583465055446615090444309583050775808509297040039680057435342253926566240898195863631588888936364129920059308455669454034010391478238784189888594672336242763795138176353222845524644040094258962433613354036104643881925238489224010194193088911666165584229424668165441688927790460608264864204237717002054744337988941974661214699689706521543006262604535890998125752275942608772174376107314217749233048217904944409836238235772306749874396760463376480215133461333478395682746608242585133953883882226786118030184028136755970045385534758453247
g = 5
A = 1026312539297800437474663698165859314949881437729617621666434357798219198741950468733395500361477359726152747087790103309627020498122003777642051150130697457594304849673838709900017711265818285080832347734747895550397950729716624922572654209637755195129162139245110756558638081495998280747642920484467428206475906559638681536868548289456924005274209311355030582255692087426910634838198143851507435754029135363794578075936092774722678311786272841489629294721103591751528609388061794369341067986401129462942050916582521451289187645626081017578576190303952351748434876686541368607656026867091583868645619423975306245327421218767449273192101105293424028461698783545171866070124432565063559495566733441286372612161876492134408160732339966921175762866198980795890946054558528891296203285979664329713156129091098226212735763844909789916934266711879564086741733061623347281499025678164709559814150194881622611023214199434022258730549350019749882889143749385314934896284396513061241138504029046053964944026179039768718830854958634216721133676746317913559932277050177463811150719675119168868527853864167729206220819613297736800799391257602899169041109002518019207734013899840092155297682096290489330476118066934735827328128343402508975429994312

print("Check P - B = 1: {}".format(p-B))

#B = (p-1)
# (p - 1) = G^d mod p
# -1 = G^d mod p
# 1 = G^2d mod p

# G^(p-1) mod p = 1 mod p = G^2d mod p
# G^(p-1) mod p = G^2d mod p
# p-1 = 2d
# d = (p-1)/2

secret_exp = (p-1)/2

#Simple Decrypt with the secret_exp 
shared_secret = pow(A, secret_exp, p)
shared_secret_inv = inverse(shared_secret, p)
message = (shared_secret_inv*c) % p
print long_to_bytes(message)

if __name__ == '__main__':
	message = b"Hello World"
	message_int = bytes_to_int(message)

	#Generate User1 Ephemeral Key
	user1_public_key, user1_private_key = generate_key_pair()

	#Generate User2 Ephemeral Key
	user2_public_key, user2_private_key = generate_key_pair()

	#Forge Message 1
	message = decrypt_message(1, message_int, user2_private_key)
	print(int_to_bytes(message))
	#b'Hello World'

if __name__ == '__main__':
	message = b"ROOT USER"
	message_int = bytes_to_int(message)
	message_int_2 = message_int//2
	message_2 = int_to_bytes(message_int_2)

	#Generate User1 Ephemeral Key
	user1_public_key, user1_private_key = generate_key_pair()

	#Generate User2 Ephemeral Key
	user2_public_key, user2_private_key = generate_key_pair()

	#Message sent from User1 -> User2
	if message_2 != message:
		print("Message: {}".format(message_2))
		#Message: b")'\xa7\xaa\x10*\xa9\xa2\xa9"
		ephemeral_key_ciphertext, message_ciphertext = encrypt_message(user2_public_key, user1_private_key, message_2)
		print("Encrypted Message: {}".format(ephemeral_key_ciphertext, message_ciphertext))
		#Encrypted Message: 177913139231183757166602993688011837747475632893569829207588727849541580678770957146896394618561833477462244212416054898100733439176530010050589111378346021151466593709817552811046961391070465749142459198149578717258866028928562353077442201079640435402512552658507741723520448880299370579625268890621776663411

	#Message decrypted by User2
	decrypted_int_message = decrypt_message(ephemeral_key_ciphertext, message_ciphertext * 2, user2_private_key)
	decrypted_message = int_to_bytes(decrypted_int_message)
	print("Decrypted Message: {}".format(decrypted_message))
	#Decrypted Message: b'ROOT USER'

	#Forge Message 2
	message = decrypt_message(2, message_int, user2_private_key)
	print(int_to_bytes(message))