from sage.all import *
from Crypto.Util.number import *
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from hashlib import md5
from secret import flag


class algebra(CommutativeRingElement):
    def __init__(self, parent, a1, a2, c=None):
        CommutativeRingElement.__init__(self, parent)
        self.a1 = a1
        self.a2 = a2
        self.c = c

    def __add__(self, other):
        return algebra(
            self.parent(), max(self.a1, other.a1), max(self.a2, other.a2), c=self.c
        )

    def __mul__(self, other):
        p = 0
        if self.c is not None:
            p += self.c
        first = max(p + self.a1 + other.a1, p + self.a2 + other.a2)
        second = max(p + self.a1 + other.a2, p + self.a2 + other.a1)
        return algebra(self.parent(), first, second, c=self.c)

    def __eq__(self, other):
        return (self.a1 == other.a1) and (self.a2 == other.a2)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __repr__(self):
        return f"T({self.a1}, {self.a2})"

    def _mul_(self, other):
        return self.__mul__(other)


class algebraring(UniqueRepresentation, Parent):
    def __init__(self, c=None):
        Parent.__init__(self, category=CommutativeRings())
        self.c = c

    def _element_constructor_(self, a1, a2):
        return algebra(self, a1, a2, self.c)

    def __repr__(self):
        return "algebraring"

    def zero(self):
        return self(float("-inf"), float("-inf"))

    def random_element(self, lower_bound=-10, upper_bound=10):
        from random import randint

        a1 = randint(lower_bound, upper_bound)
        a2 = randint(lower_bound, upper_bound)
        return self._element_constructor_(a1, a2)


def change_ring(matrix, ring):
    n = len(matrix[0, :].list())
    new_matrix = [
        [ring(matrix[i, j].a1, matrix[i, j].a2) for j in range(n)] for i in range(n)
    ]
    return Matrix(ring, new_matrix)


def shift(matrix, x):
    n = len(matrix[0, :].list())
    ring = matrix.base_ring()
    new_matrix = [
        [ring(matrix[i, j].a1 + x, matrix[i, j].a2 + x) for j in range(n)] for i in range(n)
    ]
    return Matrix(ring, new_matrix)


n = 3

c = getRandomNBitInteger(256)
k = getRandomNBitInteger(256)
l = getRandomNBitInteger(256)
p = getRandomNBitInteger(256)

d = getRandomNBitInteger(256)
r = getRandomNBitInteger(256)
s = getRandomNBitInteger(256)
q = getRandomNBitInteger(256)

T = algebraring()
TA = algebraring(c)
TB = algebraring(d)

X = random_matrix(T, n, lower_bound=0, upper_bound=2**256)
Y = random_matrix(T, n, lower_bound=0, upper_bound=2**256)

XA = change_ring(change_ring(X, TA) ** k, T)
YA = change_ring(change_ring(Y, TA) ** l, T)
A = XA * YA
A_p = shift(A, p)

XB = change_ring(change_ring(X, TB) ** r, T)
YB = change_ring(change_ring(Y, TB) ** s, T)
B = XB * YB
B_q = shift(B, q)

KA = shift(XA, p) * B_q * YA
KB = shift(XB, q) * A_p * YB
assert KA == KB

key = md5(str(KA).encode()).digest()
aes = AES.new(key=key, mode=AES.MODE_ECB)
print(f"enc = {aes.encrypt(pad(flag, 16))}")
print(f"X = {X.list()}")
print(f"Y = {Y.list()}")
print(f"A_p = {A_p.list()}")
print(f"B_q = {B_q.list()}")
"""
enc = b'\x88I\xa6\xa5\x83\x8a\xd3$dk/\xf9\xdb\xc7\xec\xe8\xf6\x8c\x9d\xad\xe7\x8d\xb8\xe73~&\xb6O\xe5\x0f\x12\xf3?@\xd4\xc6I\xeel\x9d\xdd8\x92*gY`'
X = [T(59900590664029081588751597115949479657291730864508459758763562295986654520462, 68320552779283515276240499643319977619901440989200422494626425642962224959344), T(89698054093719949422762429034310548012077225043893131168337545047107450749496, 51479199576292239130066082976332639976062899059190752269878695108378488146679), T(62218736027330266192570144338062026977337231944985836220559539768109014406831, 61214354159525954062640070826510640683451706120539795727362849462686622564693), T(39805827211060089460158250456599968341463470401750393132792110953878249393569, 99103416592859133938491410425030428643259368893929508081326035011417306901272), T(20922493031536155240718266043337868355652138270147183521353456319726724009058, 31430964180770964194690164159235543367301572045989047438624310779165592317271), T(8959191626315499956806224120588498434035126032025230565571486221513367178374, 20067214390603904748675229391030612573671604470294660618614875919888930884780), T(66048220278298402781571427858585494416782927078968140383635650050886579125833, 18175597939833315588572188448012546734970985103503074208793106634216584036818), T(14432081053446841732493980558668616443341123394659689880088380115934894274635, 83557493069309650142798328460651561327339732728386566873314215777128303177667), T(34511162979554971448736265846077014796752072453004958739933710734947541396694, 77364571020765493103034943723333754233327577188690990300833856668531817680505)]
Y = [T(78533528545456238383274783853785358016819467733594018596743219756487720633746, 93976221569207299395273293543645974643384826524816615562497005292252667770601), T(103766862112463335445955325531890205089005117500922857426350491681602103969000, 47323529900779960229909476477918095608118443655903527805237066535366116839149), T(10190752330406968772077432866677538790850926811710030717810026338836331450472, 86375880251137784911361557173115785351126971447566503084261189167084220453288), T(13780196681979042044702022951434986810164135209191561823725939514845972817610, 55170363535880083300862447306063246181282182639154190917220001586753468629546), T(436029875577652990270691130860178026240806299188121622633421834221330373636, 110804259765120571425960874898903708961853698355143533817410855395636486531982), T(6397911849400997286703265695504239070876241554571709807303889197285939312681, 17360696619752702589750097897661044951586752285334140009669640422948343241568), T(58041034583523761593801980039841494595215125402194920275127841961766699686394, 54766044749028865088458074550257196739016298676359665986238379481142351990626), T(77530103367856397256662107816221151235947660504364282548276904992741575506849, 63854892725980295947169956482257581700294549443273322483764382131671051522332), T(47497441477952935509240863594022229337695084880724507207404445061842628668952, 56713418154978297086219120150091406983389112617302018465360827172758909144570)]
A_p = [T(40547797283575399311345964287271714617688023231464400205040285835460999315476731205228816544079212468645004319629302959473711018684702750872290449940972698, 40547797283575399311345964287271714617688023231464400205040285835460999315476698934497596879746169782553959201278357925243089469169482083236651301175074462), T(40547797283575399311345964287271714617688023231464400205040285835460999315476786839125045784567337567072597160092083530989427008027602941726099332958875134, 40547797283575399311345964287271714617688023231464400205040285835460999315476750004263963436463523550691311785519855896914393402679802307379634177780083118), T(40547797283575399311345964287271714617688023231464400205040285835460999315476706776849302561292697869327278531705692232746803441653969601206061897674894004, 40547797283575399311345964287271714617688023231464400205040285835460999315476682432777130064993198309463393760622192553532626436203592834759900982411655833), T(40547797283575399311345964287271714617688023231464400205040285835460999315476733573193663026027748192077028026006061293036706834385224678998540725414561492, 40547797283575399311345964287271714617688023231464400205040285835460999315476709827727502107645065946715613329058910783689656325976029039752886304715182546), T(40547797283575399311345964287271714617688023231464400205040285835460999315476789207089892266515873290504620866468841864552422823728124869852349608432463928, 40547797283575399311345964287271714617688023231464400205040285835460999315476765461623731348133191045143206169521691355205372315318929230606695187733084982), T(40547797283575399311345964287271714617688023231464400205040285835460999315476709144814149043241233592759302238082450566309799257354491529332312173148482798, 40547797283575399311345964287271714617688023231464400205040285835460999315476685399347988124858551347397887541135300056962748748945295890086657752449103852), T(40547797283575399311345964287271714617688023231464400205040285835460999315476692793936572469446889818453385542291673187750773962605187059907381322027502633, 40547797283575399311345964287271714617688023231464400205040285835460999315476725064667792133779932504544430660642618221981395512120407727543020470793400869), T(40547797283575399311345964287271714617688023231464400205040285835460999315476746626235376326170902419020010827637210112057551060804659922530107892689013547, 40547797283575399311345964287271714617688023231464400205040285835460999315476780698564021374268057602972023501105398793497111501463307918396829353811303305), T(40547797283575399311345964287271714617688023231464400205040285835460999315476676292216105654693918345362820101635507816040310929639297811430631003264084004, 40547797283575399311345964287271714617688023231464400205040285835460999315476700636288278150993417905226704872719007495254487935089674577876791918527322175)]
B_q = [T(44956584717324859629636950162153871886560739437959809729281253007078755840271381302732123072890496216398568268656559622476770594999251404506445788064037538, 44956584717324859629636950162153871886560739437959809729281253007078755840271413573463342737223538902489613387007504656707392144514472072142084936829935774), T(44956584717324859629636950162153871886560739437959809729281253007078755840271432372498489629607849984535920852898057594148074528509571628649428664669046194, 44956584717324859629636950162153871886560739437959809729281253007078755840271469207359571977711664000917206227470285228223108133857372262995893819847838210), T(44956584717324859629636950162153871886560739437959809729281253007078755840271364801011656258137524743308002828000394250766307562033362156029695469300618909, 44956584717324859629636950162153871886560739437959809729281253007078755840271389145083828754437024303171887599083893929980484567483738922475856384563857080), T(44956584717324859629636950162153871886560739437959809729281253007078755840271392195962028300789392380560222396437112480923337451805798361022680791604145622, 44956584717324859629636950162153871886560739437959809729281253007078755840271415941428189219172074625921637093384262990270387960214994000268335212303524568), T(44956584717324859629636950162153871886560739437959809729281253007078755840271447829858257541277517478987815236899893052439053441148698551876489674622048058, 44956584717324859629636950162153871886560739437959809729281253007078755840271471575324418459660199724349229933847043561786103949557894191122144095321427004), T(44956584717324859629636950162153871886560739437959809729281253007078755840271367767582514318002877781242496608513501754196429874775065211356452239338066928, 44956584717324859629636950162153871886560739437959809729281253007078755840271391513048675236385560026603911305460652263543480383184260850602106660037445874), T(44956584717324859629636950162153871886560739437959809729281253007078755840271407432902318326924258938389039728020819919215076637950177048812814957682363945, 44956584717324859629636950162153871886560739437959809729281253007078755840271375162171098662591216252297994609669874884984455088434956381177175808916465709), T(44956584717324859629636950162153871886560739437959809729281253007078755840271463066798547567412384036816632568483600490730792627293077239666623840700266381, 44956584717324859629636950162153871886560739437959809729281253007078755840271428994469902519315228852864619895015411809291232186634429243799902379577976623), T(44956584717324859629636950162153871886560739437959809729281253007078755840271383004522804344137744339071313940097209192488169060919443899146586405416285251, 44956584717324859629636950162153871886560739437959809729281253007078755840271358660450631847838244779207429169013709513273992055469067132700425490153047080)]
"""