from pwn import *
import numpy as np
from zlib import compress, decompress
from base64 import b64encode as b64e, b64decode as b64d
from itertools import product

q = 2 ** 11
n = 280
n_bar = 4

LAMBDA = 512

s = remote('challenges1.france-cybersecurity-challenge.fr', 2001)

msg = s.recvuntil(b'Possible actions')
s.recv(1024)

A = msg.split(b'A = ')[1].split(b'\n')[0]
B = msg.split(b'B = ')[1].split(b'\n')[0]

A = np.reshape(np.frombuffer(decompress(b64d(A)), dtype = np.int64), (n, n))
B = np.reshape(np.frombuffer(decompress(b64d(B)), dtype = np.int64), (n, n_bar))

__S_a = np.zeros((n, n_bar), dtype = np.int64)

s.send(b'1\n')

for k in range(n):
  U = np.zeros((n_bar, n), dtype = np.int64)
  C = np.zeros((n_bar, n_bar), dtype = np.int64)

  U[0][k] = LAMBDA

  U = b64e(compress(U.tobytes()))
  C = b64e(compress(C.tobytes()))

  for c in product([-1, 0, 1], repeat=n_bar):
    if not (c.count(0) == 2 and c.count(1) == 1 and c.count(-1) == 1):
      continue

    CMP = np.zeros((n_bar, n_bar), dtype = np.int64)
    for i in range(n_bar):
      CMP[0][i] = c[i]
    CMP = b64e(compress(CMP.tobytes()))

    s.recv(1024)
    s.send(U + b'\n')

    s.recv(1024)
    s.send(C + b'\n')

    s.recv(1024)
    s.send(CMP + b'\n')

    msg = s.recv(1024)
    s.send(b'1\n')

    if b'Success' in msg:
      break

  for i in range(n_bar):
    __S_a[k][i] = -c[i]

  print("[+] Ligne %s: %s" % (k, repr(__S_a[k])))

__E_a = np.mod(B - np.dot(A, __S_a), q)

def t(x):
  if x == q - 1:
    return -1
  return x

__S_a = b64e(compress(np.vectorize(t)(__S_a).tobytes()))
__E_a = b64e(compress(np.vectorize(t)(__E_a).tobytes()))

print(__S_a)
print(__E_a)

s.interactive()
s.close()