import random
import hashlib
from sage.crypto.util import ascii_to_bin
from sage.arith.misc import random_prime
from sage.crypto.util import ascii_to_bin
from sage.crypto.util import bin_to_ascii

p = 0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f
n = 0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141

E = EllipticCurve(GF(p), [0, 7])

r = 0x
s = 0x
z = 0x


G = E.point( (0x79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798,0x483ada7726a3c4655da4fbfc0e1108a8fd17b448a68554199c47d08ffb10d4b8))   # Base point

def egcd(a, b):

    if a == 0:

        return (b, 0, 1)

    else:

        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:

        raise Exception('modular inverse does not exist')

    else:

        return x % m


def make_public(r, s, z):
    R = E.lift_x(Integer(r))
    w = int(modinv(s, n))
    u1 = int((z * w) % n)
    u2 = int((r * w) % n)
    u_n2 = modinv(u2, n) % n
    u_n1 = - u1 * modinv(u2, n) % n
   
    pub = u_n1 * G + u_n2 * R
    pub2 = u_n1 * G + u_n2 * (-R)
    return pub, pub2



def calc_u(r, s, z):
    mod_s = modinv(s, n) % n
    u1 = mod_s * z % n
    u2 = mod_s * r % n
    print("u1 =", hex(u1), "n-u1 =", hex(n - u1))
    print("u2 =", hex(u2), "n-u2 =", hex(n - u2))
    return u1, u2
u1, u2 = calc_u(r, s, z)

def verify(r, s, z, pub, k):
    w = int(modinv(s, n))
    u1 = int((z * w) % n)
    u2 = int((r * w) % n)
    D = u1 * G + u2 * pub
    x, y = D.xy()
    x = int(x)

    if (r % n) == (x % n):
        print(f"Signature k matches: {hex(k)} ")
        return True
    else:
        print(f"Signature k is invalid {hex(k)} ")
        return False

# Calculate the modular inverse of s (w = s^(-1) mod n)
g, x, y = egcd(s, n)
w = x % n

def find_k(r, s, z, pub):
    # Step 1: Calculate w = s^(-1) mod n
    w = int(modinv(s, n))
   
    # Step 2: Calculate u1 and u2
    u1 = int((z * w) % n)
    u2 = int((r * w) % n)
   
    # Step 3: Recover the elliptic curve point R
    R = u1 * G + u2 * pub
   
    # Extract x-coordinate from R
    x_R, _ = R.xy()
    x_R = int(x_R)
   
    # Calculate k
    k = (x_R - z) * modinv(w, n) % n
    print("k:",k)
   
    return k

# Call the function to find k
k = find_k(r, s, z, pub1)
print("Found k:", hex(k))

pub1, pub2 = make_public(r, s, z)
print("public_key1", pub1)
print("pub1_x=", hex(pub1.xy()[0]))
print("public_key2", pub2)
print("pub2_x=", hex(pub2.xy()[0]))
print("Step 1: Calculate g, x, y using extended Euclidean algorithm")
print("g:", hex(g))
print("x:", hex(x))
print("y:", hex(y))
print("Step 2: Calculate w (modular inverse of s)")
print("Calculated w:", hex(w))
print()


def find_k_bruteforce(r, s, z, pub):
    i = 1
    while True:
        k_candidate = (r * i + z) * modinv(s, n) % n
       
        # Verify the signature using the candidate k
        if verify(r, s, z, pub, k_candidate):
            print("Found k:", hex(k_candidate))
            break
        else:
            print(f"Attempt {i}: Incorrect k value {hex(k_candidate)}")
       
        i += 1  # Increment the attempt counter

# Call the function to find k using brute force
find_k_bruteforce(r, s, z, pub1)

def find_private_key(r, s, z, k, pub):
    # Calculate w = s^(-1) mod n
    w = int(modinv(s, n))
   
    # Calculate u1 and u2
    u1 = int((z * w) % n)
    u2 = int((r * w) % n)
   
    # Recover the elliptic curve point R
    R = u1 * G + u2 * pub
   
    # Extract x-coordinate from R
    x_R, _ = R.xy()
    x_R = int(x_R)
   
    # Calculate private key d
    d = (k * x_R - z) * modinv(r, n) % n
   
    return d

def find_private_key_bruteforce(r, s, z, k, pub, i):
    k_candidate = (r * i + z) * modinv(s, n) % n
   
    # Calculate private key d
    d = (k_candidate * r - z) * modinv(k_candidate, n) % n
   
    return d

# Call the function to find k using the formula-based approach
k_formula = find_k(r, s, z, pub1)
d_formula = find_private_key(r, s, z, k_formula, pub1)
print("Private key (formula-based):", hex(d_formula))

# Call the function to find k using brute force
k_bruteforce = find_k_bruteforce(r, s, z, pub1)
i_bruteforce = 1
d_bruteforce = find_private_key_bruteforce(r, s, z, k_bruteforce, pub1, i_bruteforce)
print(f"Private key (brute-force attempt {i_bruteforce}):", hex(d_bruteforce))

guessed_k_values = []
i = 1

while True:
    k = (r * i + z) * modinv(s, n) % n
    # Calculate R using the guessed k value
    guessed_R = k * G
    guessed_k_values.append((i, k, guessed_R))

    # Check if the guessed R value matches the original r value
    if guessed_R.xy()[0] == r:
        print(f"The correct k value ({k}) for iteration {i} is found.")
        # Calculate the private key using the correct k value
        d_bruteforce = find_private_key_bruteforce(r, s, z, k, pub1, i)
        print(f"Corresponding private key: {hex(d_bruteforce)}")
        break  # Exit the loop if the correct k is found

    i += 1  # Increment the iteration counter