import json
import os
import requests
from sage.all import *
from Crypto.Util.number import *
from Crypto.Util.strxor import strxor
from hashlib import sha512
from base64 import urlsafe_b64decode, urlsafe_b64encode
from pasta import PastaSigner, PastaVerifier

HOST = "10.100.101.102:13000"
# HOST = "0.0.0.0:8000"


def register(username):
    r = requests.post('http://{}/register'.format(HOST), json={'username': username, 'password': '123'})
    print(r.json())


def login(username):
    r = requests.post('http://{}/auth?version=4'.format(HOST), json={'username': username, 'password': '123'})
    token = r.json()['token']

    return token


def get_flag(token):
    r = requests.get('http://{}/flag'.format(HOST), headers={'Authorization': 'Bearer {}'.format(token)})
    return r.json()


secret = b'\x00' + os.urandom(65)
signer = PastaSigner(secret, 4)
verifier = PastaVerifier(secret)


sigs = []
n = 110
for i in range(n):
    username = "pasta-{}".format(i)
    register(username)
    token = login(username)
    sigs.append(token.encode())
    # sigs.append(signer.sign(json.dumps({"username": username, "role": "user"})).encode())


hs = []
rs = []
ss = []

for i in range(n):
    _, _, sig = sigs[i].split(b".")
    raw_data = urlsafe_b64decode(sig + (b'==' * 2))
    size = 132
    signature = raw_data[-size:]
    public_key = raw_data[-size * 2:-size]
    message = raw_data[:-size]

    r = bytes_to_long(signature[:66])
    s = bytes_to_long(signature[66:])

    hs.append(bytes_to_long(sha512(message).digest()))
    rs.append(r)
    ss.append(s)


h1 = int(hs[0])
r1 = int(rs[0])
s1 = int(ss[0])

captured = []
for i in range(len(hs)):
    captured.append((int(hs[i]), int(rs[i]), int(ss[i])))

order = 0x01fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffa51868783bf2f966b7fcc0148f709a5d03bb5c9b8899c47aebb6fb71e91386409

matrix = []
nonce_bit = 512  # bit size of nonce

i = 0
n = len(captured) + 2

max_nonce = 2**nonce_bit

for signature in captured:

    matrix.append([0] * n)
    matrix[i][i] = order

    i += 1

matrix.append([0] * n)
matrix.append([0] * n)

i = 0
for signature in captured:
    h, r, s = signature

    inv_s = inverse_mod(s, order)

    matrix[n - 2][i] = r * inv_s
    matrix[n - 1][i] = h * inv_s

    i += 1

matrix[n - 2][n - 2] = int(max_nonce) / order
matrix[n - 2][n - 1] = 0
matrix[n - 1][n - 2] = 0
matrix[n - 1][n - 1] = max_nonce

print("LLL")

B = Matrix(QQ, n, n, matrix)
L = B.LLL()

possible_d = []
for row in list(L):
    k1 = int(abs(row[0]))
    if k1 != 0 and k1 != max_nonce and k1 < max_nonce:
        d = (k1 * s1 - h1) * inverse_mod(r1, order) % order

        possible_d.append('00' + long_to_bytes(d).hex())

# assert secret.hex() in possible_d

for d in possible_d:
    fake_signer = PastaSigner(bytes.fromhex(d), 4)
    token = fake_signer.sign(json.dumps({"username": "pwned", "role": "admin"}))
    print(get_flag(token))