from sage.all import *
from socket import socket, AF_INET, SOCK_DGRAM
import uuid
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad
import hashlib

IP = "127.0.0.1"
flag_id = "b0e29b3a-f560-449c-ad59-1b760c4aab47"

PR = PolynomialRing(GF(2), "x")

def bits(n):
    return [int(c) for c in bin(n)[2:]][::-1]

def unbits(bs):
    return sum(int(x) << i for i, x in enumerate(list(bs)))

get_crypto_params_payload = bytes((
    0xF0,
    0x7D,
    0x53,
    0x58,
    0x05,
    0xF7
))

sock = socket(AF_INET, SOCK_DGRAM)
sock.bind(("0.0.0.0", 0))
sock.connect((IP, 5004))

sock.sendall(get_crypto_params_payload)
msg = sock.recv(1024)

modulus = PR(bits(int.from_bytes(msg[5:5 + 64], "big")))
base = PR(bits(int.from_bytes(msg[5 + 64:5 + 64 + 64], "big")))


get_crypto_params_payload = bytes((
    0xF0,
    0x7D,
    0x53,
    0x58,
    0x03,
    *uuid.UUID(flag_id).bytes,
    0xF7
))

sock.sendall(get_crypto_params_payload)
msg = sock.recv(1024)
h = PR(bits(int.from_bytes(msg[5:5 + 64], "big")))
c1 = PR(bits(int.from_bytes(msg[5 + 64:5 + 64*2], "big")))
iv = msg[5 + 64*2:5 + 64*2 + 16]
ct = msg[5 + 64*2 + 16:-1]

fact = factor(modulus)
assert len(fact) == 1
assert fact[0][1] == 5
D = fact[0][0].degree()

# part modulo p
R = GF(2 ** D, 'a', modulus=fact[0][0])
print(R)
r = R(c1).log(base)

# remaining part
R = PR.quotient(modulus)
newbase = R(base) ** (2 ** D - 1)
newres = R(c1) ** (2 ** D - 1)
for x in range(1, 8):
    if newbase ** x == newres:
        break
else:
    print("failed to find part")
r = crt([r, x], [2 ** D - 1, 8])
assert R(base) ** r == R(c1)

shared_key = int(unbits(list(R(h) ** r))).to_bytes(64, "big")
key = hashlib.sha256(shared_key).digest()[:16]
cipher = AES.new(key, AES.MODE_CBC, iv)
pt = unpad(cipher.decrypt(ct), 16)
print(pt)