import nclib
import ecdsa
import hashlib
import binascii
from Crypto.Util.number import *

def _hex(x):
    return binascii.hexlify(x).decode()

def _hash(msg):
    return bytes_to_long(hashlib.sha1(msg).digest())

conn = nclib.Netcat(("crypto.csaw.io", 5002))
# conn = nclib.Netcat(("localhost", 1337))

N = 6
h = []
sgns = []
msgs = []

def get_messages(count):
    all_msgs = []

    conn.recv_until(b": ")
    conn.send_line(b"2")

    for i in range(count):
        conn.recv_until(f"Block {i}\n".encode())
        msg_line = conn.recv_line()
        print(msg_line)
        msg = binascii.unhexlify(msg_line.strip().split(b" ")[1])
        sig_line = conn.recv_line()
        sig_rs = sig_line.decode().split("(")[1].split(")")[0]
        r, s = sig_rs.split(", ")
        r = int(r)
        s = int(s)
        all_msgs.append((msg, r, s))

    return all_msgs

for i in range(N):
    msg = f"meow{i}".encode()
    msgs.append(msg)
    conn.recv_until(b": ")
    conn.send_line(b"1")
    conn.send_line(_hex(msg))

all_msgs = get_messages(N+1)

for i in range(N):
    (_, r, s) = all_msgs[i+1]
    prev_msg = all_msgs[i][0]
    real_hash = _hash(prev_msg + msgs[i])
    h.append(real_hash)
    sgns.append((r,s))

order=115792089237316195423570985008687907852837564279074904382605163141518161494337

s_inv = []
s = []
r = []
for i in range(N):
    s.append(sgns[i][1])
    r.append(sgns[i][0])
    s_inv.append(ecdsa.numbertheory.inverse_mod(s[i], order))

Z = GF(order)
R = PolynomialRing(Z, names=('dd',))
(dd,) = R._first_ngens(1)

# the polynomial we construct will have degree 1 + Sum_(i=1)^(i=N-3)i in dd
# our task here is to compute this polynomial in a constructive way starting from the N signatures in the given list order
# the generic formula will be given in terms of differences of nonces, i.e. k_ij = k_i - k_j where i and j are the signature indexes
# each k_ij is a first-degree polynomial in dd
# this function has the goal of returning it given i and j
def k_ij_poly(i, j):
	hi = Z(h[i])
	hj = Z(h[j])
	s_invi = Z(s_inv[i])
	s_invj = Z(s_inv[j])
	ri = Z(r[i])
	rj = Z(r[j])
	poly = dd*(ri*s_invi - rj*s_invj) + hi*s_invi - hj*s_invj
	return poly

# the idea is to compute the polynomial recursively from the given degree down to 0
# the algorithm is as follows:
# for 4 signatures the second degree polynomial is:
# k_12*k_12 - k_23*k_01
# so we can compute its coefficients.
# the polynomial for N signatures has degree 1 + Sum_(i=1)^(i=N-3)i and can be derived from the one for N-1 signatures

# let's define dpoly(i, j) recursively as the dpoly of degree i starting with index j

def dpoly(n, i, j):
	if i == 0:
		return (k_ij_poly(j+1, j+2))*(k_ij_poly(j+1, j+2)) - (k_ij_poly(j+2, j+3))*(k_ij_poly(j+0, j+1))
	else:
		left = dpoly(n, i-1, j)
		for m in range(1,i+2):
			left = left*(k_ij_poly(j+m, j+i+2))
		right = dpoly(n, i-1, j+1)
		for m in range(1,i+2):
			right = right*(k_ij_poly(j, j+m))
		return (left - right)

def print_dpoly(n, i, j):
	if i == 0:
		print('(k', j+1, j+2, '*k', j+1, j+2, '-k', j+2, j+3, '*k', j+0, j+1, ')', sep='', end='')
	else:
		print('(', sep='', end='')
		print_dpoly(n, i-1, j)
		for m in range(1,i+2):
			print('*k', j+m, j+i+2, sep='', end='')
		print('-', sep='', end='')
		print_dpoly(n, i-1, j+1)
		for m in range(1,i+2):
			print('*k', j, j+m, sep='', end='')
		print(')', sep='', end='')

print_dpoly(N-4, N-4, 0)
poly_target = dpoly(N-4, N-4, 0)
d_guesses = poly_target.roots()
print("results!!!!!\n\n\n")
print(d_guesses)

if len(d_guesses) == 1:
    conn.recv_until(b": ")
    conn.send_line(b"4")
    res = str(d_guesses[0][0]).encode()
    print("SENDING", res)
    conn.send_line(res)
    print(conn.recv_all(timeout=5))