Added note on imagemagick

README.md

@@ -30,6 +30,13 @@ Mac
 homebrew install zbar
 ```
 
+`imagemagick` is also needed to display the image through `Image.show()`. This may be installed using:
+
+Ubuntu
+```
+sudo apt install imagemagick
+```
+
 ## References
 
 This project utilises code from digitalduke (https://github.com/digitalduke/otpauth-migration-decoder)

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:
-    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]
-
+def decode_totp(totp_uri: str) -> dict:
     account_start = totp_uri.split("totp/")[1]
     account = account_start.split("?")[0]
 
-    digits_start = totp_uri.split("&digits=")[1]
-    digits = digits_start.split("&")[0]
+    uri_params = account_start.split("?")[1]
+
+    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()

requirements.txt

@@ -2,3 +2,4 @@ pillow
 pyzbar
 qrcode
 protobuf
+pytest

tests/test_gauth_reencoder.py

@@ -0,0 +1,28 @@
+import pytest
+from gauth_reencoder import decode_totp
+from gauth_reencoder import generate_new_url
+
+@pytest.mark.parametrize("url, account, issuer, secret, algorithm, digits", [
+    ("otpauth://totp/DEF?algorithm=&digits=&issuer=ABC&secret=jidowio390u20", "DEF", "ABC", "jidowio390u20", "", ""),
+    ("otpauth://totp/Label?secret=jidowio390u20&issuer=Issuer", "Label", "Issuer", "jidowio390u20", "", ""),
+])
+def test_decode_totp(url: str, account: str, issuer: str, secret: str, algorithm: str, digits: str):
+    decoded = decode_totp(url)
+    assert decoded["account"] == account
+    assert decoded["issuer"] == issuer
+    assert decoded["secret"] == secret
+    assert decoded["algorithm"] == algorithm
+    assert decoded["digits"] == digits
+
+@pytest.mark.parametrize("url, account, issuer, secret, algorithm, digits", [
+    ("otpauth://totp/DEF?algorithm=&digits=&issuer=ABC&secret=jidowio390u20", "DEF", "ABC", "jidowio390u20", "", ""),
+])
+def test_generate_new_url(url: str, account: str, issuer: str, secret: str, algorithm: str, digits: str):
+    decoded = dict()
+    decoded["account"] = account
+    decoded["issuer"] = issuer
+    decoded["secret"] = secret
+    decoded["algorithm"] = algorithm
+    decoded["digits"] = digits
+    new_url = generate_new_url(decoded)
+    assert new_url == url