from Crypto.Util.number import getPrime, inverse
from math import gcd
from hashlib import sha256
from Crypto.Cipher import AES



flag = b"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"


class Alice:
    def __init__(self) -> None:
        p, q = getPrime(1024), getPrime(1024)
        N = p * q
        d = getPrime(int(2048 * 0.35))
        phi = (p - 1) * (q - 1)
        if gcd(d, phi) == 1:
            e = inverse(d, phi)

        self.sk = (p, q, d)
        self.pk = (N, e)

        self.partialkey = []

    def RSAdecrypt(self, c):
        _, _, d = self.sk
        N, _ = self.pk
        m = pow(c, d, N)
        return m


class Bob:
    def __init__(self) -> None:
        self.secret = getPrime(1024)
        self.partialkey = []

    def RSAencrypt(self, pk, m):
        N, e = pk
        c = pow(m, e, N)
        return c


A = Alice()
B = Bob()

print("A -> B:", A.pk)
pk = A.pk
N, e = pk

for _ in range(3):
    Bm = B.secret * getPrime(512)
    Bc = B.RSAencrypt(pk, Bm)
    print("B -> A:", Bc)

    Am = A.RSAdecrypt(Bc)
    A.partialkey.append(getPrime(511))
    Ac = A.RSAdecrypt(Am + A.partialkey[-1])
    print("A -> B:", Ac)

    B.partialkey.append((B.RSAencrypt(pk, Ac) - Bm) % N)


hfunc = sha256()
hfunc.update(str(B.partialkey).encode())
cipher = AES.new(hfunc.digest(), AES.MODE_ECB)
print("B -> A:", cipher.encrypt(flag))


"""
A -> B: (29287941964093188883554358759962324766413054938509973020469060324897666596790911417785791718225662909407046748730700393511727164992638859309051914817752161453430071520197806236982668426788996148542614675962835327600646675114044776422291762895540351295488463620830840260895942085590293936010897487443334022180921068879759894357182300863216381540594965567845230665432808496498721899010296008077888600276230436453346416223766922153731558028227295502011211796080344786948543217682991981791223266548648600316048178386585096118093900539239292482498040472725870838971971059537341275594528418355973441717474777281448026845373, 27955266925253148642253678752829183603219107778520606980764762752995571036134821430889656665717477157214005599893648302567419024006700075224145901058864025471376122721606447684082747992022267965241792216735472384680654771240354932987171309705823754809028693606076338867464980753571886581376406114561483162250351593164007074977896197934437555326889625656047853373889179731769563558844917159374790171076366851539889349713510477985229970467404417110370407654250262525961760288684308272656639979945733821040592567089478937570171101226406294152472692793873319781044539193653344814998127081830989348128648490881576398235955)
B -> A: 3205220094080592407592828125679505405222175350233807694327159593032969469579539804535020905896745919090265556507179288821213875752712326943938625398874083262608628666420762655582303641231925289465398823744770158764392477253388389480024069287363779774060962366992843006326211329071156343468269190434281421424221587984048575740789112158698868948466566947408953807967713643274889608573354692166180828636721319812054682845978388714126347955276139299995584288733190964976250581524403399583537399029174408442816250678567356095603425069263938872549246059702018186374731429179320660396979172676464090137034199876766815040549
A -> B: 4831644744350799393765016099396882177026534773569532962146766169148481243322219737516506695852464031849695952207227141415079187438443873384079566512604202776962147357699367571125364213415800582067021480708275284860444744873611933615510520333222572461931300221867821081702084458292525283508482399395833410218592334630954773959254262941692220532586083298023976096784214039362483222310058451221948873959013103817332959309430450412106921677559864731324809944745039703799490755791317245753122131983084501003067118380162680898017587067862940816260310005852375433008408683639088699648954723722337562683163853590166299207882
B -> A: 10223665720049500731665673770103917961793186693687648371672171897872352435447003168707921372495448171568308786361270844274614650508367480702414916432634701500942569012336282352568912820462024329259048798831366489477953163885962745990200563709834731868518742507641941828621899056783884909373008325439482906379351126150044191130243110199555815953195410189701922253411950464540623741254828469722487917152632951477621077376226540842930933084730079740736321515593713245055389579335748273149419361624409497735156695920380256756873465619561744004161764910755943356577122675253183293946422864924947394830590760908284180805653
A -> B: 6022461415472339281836498790278755583726889448623920530347868182630937987521253477126939602533575565331657993511358069473757046630594242611137100942575346645763052742511258016344480503774775023629872360358298995380098714687236835234204193892762002768854579522920204073999133902681913826901342657210771151320309334953071731535713048675285675974744410502359479970589333783354881282915026585051250730890093253822846141777793849657061271396501904548945234830806282646236893655892167888170712944257063558285805501798738399575903576977831532203975392444095534068494054836142511588829555267892072625075416002302243176635324
B -> A: 38295886401800770247937208725926281780443425977810381315056059604563712323892215369756043382931918604092036060639406709054651071621090310774115354952282617734839924719863747396452802997833397137971837803315063878192004294473556651925283669552605730995960381204942868510331336686033996535943375343952758789091959462318596236669589726125157492897857522812637457878020275467551375672236923393880594194964453929585538109959704846836439542984266313034272010738296189222642481431068045312778286597586289119135190665002893852930340193114169492386623439934950489147486565869065688389391266191081890052221792059016651759935
A -> B: 22997295428814514195091139337373650893365094613018428287375171979263198804955694038921673706847017725112033894298965577877631026760694604442860104769929651081470274609617397037067997884566485822781893833508224191144921037000725886575812219990814743777146135558370440033727198066101727185074410043409716527519221297212798655733539022858296903864363604189024082966485173370094663250147084643430665337682727021661753415715552161602283717204286568751928695603206357216967106479737282869745702423639348584040558580768510604253083698350240121502964676110415110162800693161878015666130229334069872276017309750136299165288207
B -> A: b'|\x9aN\x16puZq\xb8R\xce\x84G\x06\xb0H\x91\xb4,\xd2\xba|\xac\x91\xbf"\xa7C\n\x82\xad\xa26s- 2\xc9[\xd5>\xa3\x81"t\xd2\x00"'
"""