from PIL import Image
import PIL
import pyzbar.pyzbar as pyzbar
import otpauth_migration_decoder.src.decoder
import urllib.parse
import time
import qrcode

def decode_qr_file(img_file: str) -> str:
    img = Image.open(img_file)
    decoded = pyzbar.decode(img)
    return decoded[0].data.decode("utf8")

def decode_screenshot_img() -> str:
    from PIL import ImageGrab
    img = ImageGrab.grab()
    decoded = pyzbar.decode(img)
    return decoded[0].data.decode("utf8")

def decode_totp(totp_uri: str) -> list:
    issuer_start = totp_uri.split("&issuer=")[1]
    issuer = issuer_start.split("&")[0]

    secret_start = totp_uri.split("&secret=")[1]
    secret = secret_start.split("&")[0]

    algorithm_start = totp_uri.split("?algorithm=")[1]
    algorithm = algorithm_start.split("&")[0]

    account_start = totp_uri.split("totp/")[1]
    account = account_start.split("?")[0]

    digits_start = totp_uri.split("&digits=")[1]
    digits = digits_start.split("&")[0]

    issuer = urllib.parse.unquote_plus(issuer)
    secret = urllib.parse.unquote_plus(secret)
    algorithm = urllib.parse.unquote_plus(algorithm)
    account = urllib.parse.unquote_plus(account)
    digits = urllib.parse.unquote_plus(digits)
    
    return {"issuer": issuer, "account": account, "algorithm": algorithm, "secret": secret, "digits": digits}
    
def generate_qr_code(totp_details: dict):
    issuer = urllib.parse.quote(totp_details["issuer"])
    account = urllib.parse.quote(totp_details["account"])
    algorithm = urllib.parse.quote(totp_details["algorithm"])
    secret = urllib.parse.quote(totp_details["secret"])
    digits = urllib.parse.quote(totp_details["digits"])
    
    generated_uri = "otpauth://totp/%s?algorithm=%s&digits=%s&issuer=%s&secret=%s" % (
        account,
        algorithm,
        digits,
        issuer,
        secret
    )
    img = qrcode.make(generated_uri)
    print(generated_uri)
    img.show()


def main():
    print("Grabbing screenshot in 2 seconds...")
    time.sleep(2)
    decoded_data = decode_screenshot_img()
    print(decoded_data)
    
    totp_uri_list: list = []
    if "otpauth-migration" in decoded_data:
        totp_uri_list = otpauth_migration_decoder.src.decoder.start_decode_migration(decoded_data)
    elif "otpauth://totp/" in decoded_data:
        totp_uri_list.append(decoded_data)
    else:
        print("Error in QR code. Please check if it is proper Google Authenticator QR code.")
        return
    
    for uri in totp_uri_list:
        try:
            curr_totp_details = decode_totp(uri)
            print(curr_totp_details)
        except:
            print("Error in parsing the following URI: %s" % uri)

        # Checking for each item
        print()
        print()
        print("Do you want to change: (separator is \":::\")")
        print("%s:::%s   ...   " % (curr_totp_details["issuer"], curr_totp_details["account"]), end="")
        updated_data = input()
        updated_data_splitted = updated_data.split(":::")
        if updated_data != "":
            curr_totp_details["issuer"] = updated_data_splitted[0]
            curr_totp_details["account"] = updated_data_splitted[1]
            generate_qr_code(curr_totp_details)

if __name__ == "__main__":
    main()