from Crypto.Util.number import *
from sage.all import *
from re import findall
from subprocess import check_output
from tqdm import *
def flatter(M):
    # compile https://github.com/keeganryan/flatter and put it in $PATH
    z = "[[" + "]\n[".join(" ".join(map(str, row)) for row in M) + "]]"
    ret = check_output(["flatter"], input=z.encode())
    return matrix(M.nrows(), M.ncols(), map(int, findall(b"-?\\d+", ret)))
def decrypt(const, lwe_pubkey2, enc):
    return (int(enc)-int(const * vector(GF(lwe_ciphertext_modulus), lwe_pubkey2)))%lwe_plaintext_modulus

n = 512
m = n + 100
p = 257
q = next_prime(1048576)
c1,c2=load('/home/martin/code/task3/lwe_ciphertext.sobj')
k1,k2=load('/home/martin/code/task3/lwe_public_key.sobj')
pk_A=Matrix(GF(q),k1)
pk_b=Matrix(GF(q),k2)
encrypt_A=c1
encrypt_b=c2

def pk_Aexpress(pk_A):
    pkA_1 = pk_A[:512,:]
    pkA_2 = pk_A[512:,:]
    ks = []
    for row in pkA_2:
        ks.append(pkA_1.solve_left(row))
    return Matrix(ZZ,ks)

def fuck(A,pk_A):
    c_tmp = pk_A.solve_left(A)[:-100]
    tmpks = pk_Aexpress(pk_A)
    ks = tmpks[:,:100]
    ks = ks.stack(Matrix(ZZ,[c_tmp[:100]]))
    M = Matrix(ZZ,100 + 100 + 1,100 + 100 + 1)
    M[:101,:101] = identity_matrix(101)  
    M[:101,101:] = ks[:,:100]
    M[101:,101:] = q * identity_matrix(100)
    ML = flatter(M)
    rows = ML[0]
    c_new = [0 for i in range(612)]
    c_list = Matrix(ZZ,Matrix(GF(q),rows[:100]*tmpks) + Integer(rows[100]) * Matrix(GF(q),c_tmp))[0]
    for _ in range(512):
        if c_list[_] == q-1:
            c_new[_] = -1
        else:
            c_new[_] = int(c_list[_])
    for _ in range(100):
        if rows[_] == q-1:
            c_new[_+512] = -1
        else:
            c_new[_+512] = int(rows[_])

    return c_new

flag=''
for _ in trange(len(encrypt_A)):
    A = vector(GF(q),encrypt_A[_])
    b = encrypt_b[_]
    c_new = fuck(A,pk_A)

    c_first = c_new[:512]
    c_secon = c_new[512:]

    c = vector(ZZ, c_first+c_secon)
    if c*pk_A != A:
        c_first = [-i for i in c_first]
        c = vector(ZZ, c_first+c_secon)

    if c*pk_A != A:
        c_first = [-i for i in c_first]
        c_secon = [-i for i in c_secon]
        c = vector(ZZ, c_first+c_secon)

    if c*pk_A != A:
        c_first = [-i for i in c_first]
        c = vector(ZZ, c_first+c_secon)
    # print(c*pk_A == A)
    s=vector(GF(q),c)
    flag+=chr(decrypt(s,k2,encrypt_b[_]))
print(flag)
#a3bc5491-fa53-4f47

from Crypto.Util.number import *
from sage.all import *
from re import findall
from subprocess import check_output
def flatter(M):
    # compile https://github.com/keeganryan/flatter and put it in $PATH
    z = "[[" + "]\n[".join(" ".join(map(str, row)) for row in M) + "]]"
    ret = check_output(["flatter"], input=z.encode())
    return matrix(M.nrows(), M.ncols(), map(int, findall(b"-?\\d+", ret)))
a, b, f, rlwe_modulus=load('/home/martin/code/task3/rlwe_ciphertext.sobj')
n = 64
A = a
B = b.list()
p = rlwe_modulus
def construct_poly_mul_mat(n,v,b):
    assert v[-1] == 1
    mat1 = Matrix(ZZ,n,2*n-1)
    for i in range(n):
        for j in range(n):
            mat1[i,j+i] = b[j]
    mat2 = Matrix(ZZ,2*n-1,n)
    for i in range(n):
        mat2[i,i] = 1
    for i in range(n,2*n-1):
        for j in range(i-n,n):
            mat2[i,j] = -v[j-(i-n)]

        init_row = vector(ZZ,n*[0])
        for j in range(i-n):
            temp = -v[n-1-j]*vector(ZZ,mat2[i-j-1])
            init_row += temp
        for j in range(n):
            mat2[i,j] += init_row[j]
    return(mat1*mat2)

v = f.list()
poly_mul_mat = construct_poly_mul_mat(n,v,A)

I = identity_matrix(n)
B_mat = Matrix(ZZ,B)
O = diagonal_matrix([0]*n)
O_vec = Matrix(ZZ,1,n)
O_vec_T = Matrix(ZZ,n,1)
L = block_matrix(ZZ,[[p*I,O,0],[poly_mul_mat,I,O_vec_T],[B_mat,O_vec,1]])

res = flatter[L](0)

s = res[n:2*n]
for i in s:
    print(chr(abs(i)),end = "")
# -8819-856a8fe5ada0
# flag{a3bc5491-fa53-4f47-8819-856a8fe5ada0}

from random import getrandbits
from sage.all import *
from pwn import *
from gmpy2 import iroot
from extend_mt19937_predictor import ExtendMT19937Predictor
from re import findall
from subprocess import check_output
def flatter(M):
    # compile https://github.com/keeganryan/flatter and put it in $PATH
    z = "[[" + "]\n[".join(" ".join(map(str, row)) for row in M) + "]]"
    ret = check_output(["flatter"], input=z.encode())
    return matrix(M.nrows(), M.ncols(), map(int, findall(b"-?\\d+", ret)))
while(1):
    r = remote('', '')
    s=r.recvuntil(b'N = ')
    N=eval(r.recvuntil(b']').decode())
    res=[]
    for i in range(1,26):
        res.append(N[i]-N[0])
    m=matrix(ZZ,25,25)
    for i in range(1,25):
        m[0,i]=res[i]
        m[i,i]=-res[0]
    m[0,0]=2**(704*2)
    H=flatter(m)
    for i in H:
        if(i[0]%2**(704*2)==0):
            d=i[0]//(2**(704*2))
            if(res[0]%d==0):
                d=abs(d)
                pq=(res[0]//d-d)
                break
    p,q=var('p,q')
    f=p*q-N[0]
    g=p+q-pq
    solution=solve([f,g],[p,q])
    s=str(solution[0]).split('==')
    p=eval(s[1].split[','](0))
    q=eval(s[2].split[']'](0))
    p=p>>282
    q=q>>282
    d=[]
    print('getpq!!')
    for i in res:
        d.append(iroot[i%pq,2](0))
    rc = ExtendMT19937Predictor()
    sub=''
    sub+=bin[p](2:).zfill(1184)
    sub+=bin[q](2:).zfill(1184)
    for i in d:
        sub+=bin[i](2:).zfill(704)
    for i in range(78):
        c=int(sub[i:i+256],2)
        rc.setrandbits(c,256)
    for _ in range(78):
        s=rc.backtrack_getrandbits(256)
    m=rc.backtrack_getrandbits(256)
    s=r.recvuntil(b'm: ')
    r.sendline(str(m).encode())
    s=r.recvuntil(b'sinature = ')
    sig=r.recvuntil[b'\n'](:-1)
    s=r.recvuntil(b'signatures: ')
    sig=b"['"+sig+b"']"
    r.sendline(sig)
    flag=r.recv()
    if(b'}' in flag):
        print(flag)
        break
    r.close()