#!/usr/bin/env python3
from ecdsa.ecdsa import generator_256
from hashlib import sha1
import random, json, sys

FLAG = "flag{debug_flag}"

G = generator_256
q = G.order()
d = random.randint(1, q-1)
P = d*G

def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    g, y, x = egcd(b % a, a)
    return (g, x - (b//a) * y, y)

def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        return 0
    return x % m

def ecdsa_sign(msg):
    h = int(sha1(msg).hexdigest(), 16)
    k = random.randint(1, q - 1)
    r = int((k * G).x())
    s = (h + d * r) * modinv(k, q) % q
    return {"s": hex(s), "r": hex(r)}

def ecdsa_verify(msg, s, r):
    h = int(sha1(msg).hexdigest(), 16)
    u1 = h * modinv(s, q) % q
    u2 = r * modinv(s, q) % q
    T = u1 * G + u2 * P
    print(T)
    x = int(T.x() or 0)
    if x == r:
        return True
    return False


def main():
    while True:
        print("[1] - Sign")
        print("[2] - Verify")
        print("[3] - Exit")
        c = str(input("> "))

        if c == "1":
            try:
                msg = bytes.fromhex(input("Insert your msg (in hex) : "))
            except:
                print("Invalid input!\n")
                continue

            if b"Winterfell" in msg:
                print("Sorry, I can't sign this message for you.")
                sys.exit()

            sig = ecdsa_sign(msg)
            print(f"Signature : {json.dumps(sig)}\n")

        elif c == "2":
            try:
                msg = bytes.fromhex(input("Insert your msg (in hex) : "))
                sig = json.loads(input("Insert your signature : "))
                s = int(sig["s"], 16) % q
                r = int(sig["r"], 16) % q

                if ecdsa_verify(msg, s, r):
                    if msg == b"Winterfell":
                        print(f"Welcome, Sir. Here is your flag {FLAG}\n")
                    else:
                        print("Nothing for you.\n")
                else:
                    print("Invalid signature!\n")

            except Exception:
                print("Invalid input!\n")
                continue

        elif c == "3":
            print("Bye :)")
            sys.exit()

        else:
            print("Invalid choice!\n")


if __name__ == "__main__":
    main()