# SETUP begin...



# install necessary libraries first

# if requirements not already satisfied

    # pip install pyprimes

    # pip install pycryptodome



# import necessary modules 

import random

import pyprimes

import warnings

from Crypto.Hash import SHA3_256 # for hash

import os.path # for GenerateorRead

import string # for generating random string in fixed range

import sympy # check if the number is prime

import sys # for path 



# will be utilized in the implementation, taken from homeworks provided before

def egcd(a, b):

    x,y, u,v = 0,1, 1,0

    while a != 0:

        q, r = b//a, b%a

        m, n = x-u*q, y-v*q

        b,a, x,y, u,v = a,r, u,v, m,n

    gcd = b

    return gcd, x, y



def modinv(a, m):

    if a < 0:

        a = a+m

    gcd, x, y = egcd(a, m)

    if gcd != 1:

        return None  # modular inverse does not exist

    else:

      return x % m



def random_prime(bitsize):

    warnings.simplefilter('ignore')

    chck = False

    while chck == False:

        p = random.randrange(2**(bitsize-1), 2**bitsize-1)

        chck = pyprimes.isprime(p)

    warnings.simplefilter('default')    

    return p



def large_DL_Prime(q, bitsize):

    warnings.simplefilter('ignore')

    chck = False

    while chck == False:

        k = random.randrange(2**(bitsize-1), 2**bitsize-1)

        p = k*q+1

        chck = pyprimes.isprime(p) and (p.bit_length() == 2048)

    warnings.simplefilter('default')    

    return p

# will be utilized in the implementation, taken from homeworks provided before



# generates public parameters q,p,g

# generate random primes such that q|p-1 

# algorithm for choosing the generator:

  # choose a random alpha 

  # try a random alpha raised to the power ((p-1)) % q mod(p)

  # if anything other than 1, it's the generator of Z*p with q elements



def Param_Generator(sizeq,sizep): # qsize = 224-bit, psize = 2048-bit

  q = random_prime(sizeq)

  diff = sizep - sizeq # determine the range for generation of p 

  p = large_DL_Prime(q, diff) # discrete log prime 

  hold = (p)//q

  g = 1

  while g == 1:

      alpha1 = random.randrange(1, p)

      g = pow(alpha1, hold, p)

  return q, p, g



# generate public and private keys for the users 

# generates private key -> secret k integer such that 0 < k < q-1 

# computes public key by taking modulo p of g^k which is beta 



def KeyGen(q,p,g): 

  alpha = random.randint(1,q-2) # private

  beta = pow(g,alpha,p) # public 

  return alpha, beta



# signature generation 

# signature = (s,r)

  # hash function is 256-bit SHA 

    # take modulo q after hashing -> get h value 

  # select j (random integer) such that range = [1,q-2]

  # compute r such that (g^j(mod(p))) mod(q) is r  -> signature part1

  # compute s such that k*r - j*h (mod(q)) is s -> signature part2 (remember that k is the secret key)

  # plaintext = m



def SignGen(m,q,p,g,alpha): 

  hashval = SHA3_256.new() # create new SHA object here

  hashval.update(m) # hash here

  hold = hashval.digest() 

  h = int.from_bytes(hold, byteorder='big') # conversion from bytes to integer

  k = random.randint(1,q-2)

  r = pow(g,k,p) % q  

  s = ((alpha*r)-(k*h)) % q 

  return s,r # returning as a tuple or by comma is essentially the same thing post-function call



# signature verification 

  # now (s,r) is the signature for m 

  # carry out the same procedure as above with SHA encryption

  # take modulo inverse of h in modulo q and equate this to v

  # compute z1 such that s*v mod(q) is z1

  # compute z2 such that r*v mod(q) is z2

  # finally, compute (((g^-z1)*(beta^z2)) mod(p)) mod(q) and equate this to u



  # only verify the signature if u = r



# Signature verification

def SignVer(m,s,r,q,p,g,beta):

  hashval = SHA3_256.new() # create new SHA object here

  hashval.update(m) # hash here

  hold = hashval.digest()

  h = int.from_bytes(hold, byteorder='big') % q # conversion from bytes to integer

  v = modinv(h,q)

  z1 = (s*v) % q

  z2 = (r*v) % q

  u = ((modinv(pow(g,z1,p),p)*pow(beta,z2,p)) % p) % q

  if(u == r):

    return 0

  else:

    return random.randint(1,42) # return an arbitrary value for failure


# read the file

# if exists, read public parameters from the file

# else, create public parameters and write them to pubparams.txt

def GenerateOrRead(filename):

    if os.path.isfile(filename) == True:   #if we can reach the file open it 

            file = open(filename, 'r+')
    
            if os.stat(filename).st_size == 0:      #if file is empty generate new parameters

                q, p, g = Param_Generator(224, 2048)

                file.write(str(q) + '\n' + str(p) + '\n' + str(g))  #write the new  parameters to file

                file.close()

                return int(q), int(p), int(g)


            else:     #if file is not empty
                
                file = open(filename, 'r+')   #open the file

                txt_info = file.readlines()  #read each line

                j = 0

                while(j != len(txt_info)):    #set each line to the parameters with respect to their order

                    q = int(txt_info[j])

                    p = int(txt_info[j+1])

                    g = int(txt_info[j+2])

                    return q, p, g

                
    else:       #if we can't reach the file open new files for reading and writing

        file = open(filename,'w+')

        q, p, g = Param_Generator(224, 2048)

        file.write(str(q) + '\n' + str(p) + '\n' + str(g))  #write the new  parameters to file

        file.close()

        return int(q), int(p), int(g)
    


# return a random combination of lowercase letters, uppercase letters and digits

def random_string(str_size):

    randomized = "".join(random.choices(string.ascii_lowercase + string.digits + string.ascii_uppercase, k = str_size))

    return randomized



# SETUP end...

