#####################################################################
#   Author: Chelsea Marie Hicks
#
#   Description: Program works with Google Authenticator 
#       and has the following capabilities:
#         -generates a jpg QR code that encodes the URI GA expects
#         -generates a OTP, which matches OTP generated by GA for 
#          a 30 second period and prints this OTP to the screen
#
#   Sources in addition to those cited: 
#   https://note.nkmk.me/en/python-pillow-qrcode/
#   https://github.com/google/google-authenticator/wiki/Key-Uri-Format
######################################################################

#Import essential libraries for QR code and OTP generation
import qrcode
import qrcode.image.svg
import pyotp
import time
#import sys

#Function to generate a QR Code accepts a secret code and the totp uri
def generate_qrcode(key, uri):
    #Create a qr code class to specify QR code parameters
    #Sources: https://www.geeksforgeeks.org/python-generate-qr-code-using-pyqrcode-module/?ref=lbp
    #https://pypi.org/project/qrcode/
    qr = qrcode.QRCode(
        version = 6,
        error_correction = qrcode.constants.ERROR_CORRECT_M,
        box_size = 2,
        border = 4
    )
    #New time-based OTP created every 30-seconds with key entered as argument
    totp = pyotp.TOTP(key)
    
    #Add uri data to QR code, generate QR code, and save image as pharaohcious_qrcode.jpg to program folder
    qr.add_data(uri)
    qr.make()
    #Could be fun to experiment with different colors
    img = qr.make_image(fill_color="black", back_color="white")
    img.save('pharaohcious_qrcode.jpg')
    return 0

#Function to generate OTP and print it to screen, accepts 
#Sources: https://pypi.org/project/pyotp/
def get_otp(key, uri):
    #Counter used for continuously printing OTP
    print_count = 0

    #New time-based OTP created every 30-seconds with key entered as argument from key generation
    totp = pyotp.TOTP(key)

    #Wile loop to continually print the current OTP every 30 seconds up to 10 times, or specified number
    while(print_count < 10):
        print_count += 1
        #Print current OTP to the screen
        print("Current OTP: ", totp.now())
        #Sleep for 30 seconds
        time.sleep(30)

    return 0


#Get command from user at program start
command = input()

#Generate secret key to be shared between authenticator and user
secret_key = "3A5QCZTNABFSY777"
#Utilize this to make random key: pyotp.random_base32()
#print('Secret Key: ', secret_key)

#String to encode URI in QR code from https://github.com/google/google-authenticator/wiki/Key-Uri-Format
totp_uri = pyotp.totp.TOTP(secret_key).provisioning_uri(
    "pharaohcious@egyptiancats.com",
    issuer_name="Secure Cat"
)
#print(totp_uri)

#If command matches acceptable options, complete that action, otherwise print error message
if(command == "--generate-qr" and len(command) < 20):
    #print("qr code will be generated")
    generate_qrcode(secret_key, totp_uri)
elif(command == "--get-otp" and len(command) < 10):
    #print("creating otp")
    get_otp(secret_key, totp_uri)
else:
    print("ERROR! Invalid command entered. Goodbye!")
    exit()