import os
import secrets
from numpy import random
from hashlib import sha256
from sage.all import prod, GF, EllipticCurve

flag = os.environ.get('FLAG','sctf{test_flag}')

def csidh(A, Q, priv, err=0):
    E = EllipticCurve(GF(p), [0, A, 0, 1, 0])
    assert E.is_supersingular()
    Q = E(Q)
    for e, ell in zip(priv, ells):
        if err: e = (e+int(round(random.poisson(1.2)))) % 4
        for i in range(e):
            while not (P := (p + 1) // ell * E.random_element()): pass
            Q = (phi := E.isogeny(P))(Q)
            E = phi.codomain()
    Q = E.isomorphism_to((E := E.montgomery_model()))(Q)
    return E.a2(), Q.xy()

ells = [3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193]
p = 84*prod(ells)-1
priv = [secrets.randbelow(4) for _ in ells]

notes = [{'password':sha256(str(priv).encode()).hexdigest(),'content':flag}]

while True:
    print('''Choose your option:
1. Get public key
2. Get shared key
3. Take note
4. Read note
5. Exit''')
    choice = int(input())
    match choice:
        case 1:
            A = 0
            E = EllipticCurve(GF(p), [0, A, 0, 1, 0])
            Q = 84*E.gen(0)
            pub = csidh(0,Q,priv)
            print('Public key:',pub)
        case 2:
            pub = int(input('Enter your public curve: '))
            Q = [int(x) for x in input('Enter your public point: ').split()][:2]
            shared = csidh(pub,Q,priv,1) # error for added security!
            print('Shared secret:',shared)
        case 3:
            password = input('Enter the note password: ')
            content = input('Enter the note content: ')
            notes.append({'password':password,'content':content})
            print('Your note is at index',len(notes)-1)
        case 4:
            idx = int(input('Enter the note index you wish to read: '))
            password = input('Enter the note password: ')
            if notes[idx]['password'] == password:
                print(notes[idx]['content'])
            else:
                print('Wrong password')
        case _:
            break
    print()