diffiehellman.py

from Crypto.Util.number import *
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import hashlib
import os
flag="SECRETFLAGHEREGOFINDIT"

p=getStrongPrime(1024)
assert(isPrime(p))
g=3
farts=[]
temp=g
for i in range(2000000):
    farts.append(temp)
    temp = (temp*g) % p
assert(len(farts)  == len(set(farts)))

def pickSecretKey(p):
    a = getRandomRange(3,p-2)
    while GCD(a,p-1) != 1:
        a = getRandomRange(3,p-2)
    return a

a=pickSecretKey(p)
publicKey1 = pow(g,a,p)
b=pickSecretKey(p)
publicKey2 = pow(g,b,p)

secretNum = pow(publicKey1, b, p)
secretNum2 = pow(publicKey2, a, p)
assert(secretNum == secretNum2)

skh=hashlib.sha256(long_to_bytes(secretNum)).digest()
IV=os.urandom(16)

cipher = AES.new(skh, IV=IV, mode=AES.MODE_CBC)
ciphertext = cipher.encrypt(pad(flag,16))

testcipher=AES.new(skh, IV=IV, mode=AES.MODE_CBC)
assert(flag == unpad(testcipher.decrypt(ciphertext), 16))

print("p = %d" % p)
print("g = %d" % g)
print("b = %d" % b)
print("publicKey1 = %d" % publicKey1)
print("IV_hex = %s" % IV.encode('hex'))
print("ciphertext_hex = %s" % ciphertext.encode('hex'))

output.txt

p = 175760826546130146670478708940854358700963188586387497598428176061255235753328194210586479166483906068549887806516374247499364201774371713942125297267932513958332147170199712442590862762048271405557522879811154553625889909560001403118715972908844599080746999192112835772835919048243767310848383701980665816511
g = 3
b = 13184960970601525589288059743994324132933405978848238154954441584486593323716237825632302967388380716183715696730920935171552204483830481820157131475753665290567530523206036398608636547514374496894189724553613799389901328332084368199398585731691548843597507672649232596298654248967932616321402992184890777681
publicKey1 = 75144868973605077181985914328133353748540947222264151362542531471278925488240573655741644116843532041074627803713738975102109012431614885642472279368111040101108314554097616686117679261741778262270017685283044190799620881807833229337880744520936466493583960071893398457307472614480752064553599919949195933406
IV_hex = bytes.fromhex('83e45141558f0b2b64be4fc7cf789684')
ciphertext_hex = bytes.fromhex('09173df75ae4d7ebccbceed17389694b4528fecd861428547cf5ba15b5bda079b2290992c8c51200478477f784c28640')