Handle different URL param location

gauth_reencoder.py

@@ -17,21 +17,35 @@ def decode_screenshot_img() -> str:
     decoded = pyzbar.decode(img)
     return decoded[0].data.decode("utf8")
 
-def decode_totp(totp_uri: str) -> list:
+def decode_totp(totp_uri: str) -> dict:
-    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]
+    uri_params = account_start.split("?")[1]
-    digits = digits_start.split("&")[0]
+
+    try:
+        issuer_start = uri_params.split("issuer=")[1]
+        issuer = issuer_start.split("&")[0]
+    except:
+        issuer = ""
+
+    try:
+        secret_start = uri_params.split("secret=")[1]
+        secret = secret_start.split("&")[0]
+    except:
+        secret = ""
+
+    try:
+        algorithm_start = uri_params.split("algorithm=")[1]
+        algorithm = algorithm_start.split("&")[0]
+    except:
+        algorithm = ""
+
+    try:
+        digits_start = uri_params.split("digits=")[1]
+        digits = digits_start.split("&")[0]
+    except:
+        digits = ""
 
     issuer = urllib.parse.unquote_plus(issuer)
     secret = urllib.parse.unquote_plus(secret)
@@ -40,8 +54,8 @@ def decode_totp(totp_uri: str) -> list:
     digits = urllib.parse.unquote_plus(digits)
     
     return {"issuer": issuer, "account": account, "algorithm": algorithm, "secret": secret, "digits": digits}
-    
+
-def generate_qr_code(totp_details: dict):
+def generate_new_url(totp_details: dict) -> str:
     issuer = urllib.parse.quote(totp_details["issuer"])
     account = urllib.parse.quote(totp_details["account"])
     algorithm = urllib.parse.quote(totp_details["algorithm"])
@@ -55,6 +69,10 @@ def generate_qr_code(totp_details: dict):
         issuer,
         secret
     )
+    return generated_uri
+    
+def generate_qr_code(totp_details: dict):
+    generated_uri = generate_new_url(totp_details)
     img = qrcode.make(generated_uri)
     print(generated_uri)
     img.show()