Finishing encode/decode test in Python3

module.c

@@ -294,7 +294,9 @@ static PyObject *c_decodetotpkeywithalg(C_RAW_DATA_OBJ *self, PyObject *args, Py
    size_t out_sz;
    PyObject *ret;
 
-   if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#l", kwlist, &totpkey, &totpkey_sz, &alg))
+   alg = ALG_SHA1;
+
+   if (!PyArg_ParseTupleAndKeywords(args, kwds, "s#|l", kwlist, &totpkey, &totpkey_sz, &alg))
       PANEL_ERROR("c_decodetotpkeywithalg: Can't parse totpkey to decode", NULL)
 
    if ((err=decode_totp_key_with_alg_check_dynamic(&out, &out_sz, (int)alg, (const char *)totpkey, (ssize_t)totpkey_sz)))
@@ -320,7 +322,9 @@ static PyObject *c_encodetotpkeywithalg(C_RAW_DATA_OBJ *self, PyObject *args, Py
    size_t out_sz;
    PyObject *ret;
 
-   if (!PyArg_ParseTupleAndKeywords(args, kwds, "y#l", kwlist, &in, &in_sz, &alg))
+   alg = ALG_SHA1;
+
+   if (!PyArg_ParseTupleAndKeywords(args, kwds, "y#|l", kwlist, &in, &in_sz, &alg))
       PANEL_ERROR("c_encodetotpkeywithalg: Can't parse value to decode", NULL)
 
    if ((err=encode_totp_key_with_alg_check_dynamic(&out, &out_sz, (int)alg, (const uint8_t *)in, (size_t)in_sz)))

test/test_main.py

@@ -368,3 +368,56 @@ def test_generate_totp_key_sha1(caplog) -> None:
 
     info(genkey3)
 
+def test_encodeTOTPKeyWithAlg(caplog) -> None:
+    caplog.set_level(logging.INFO)
+
+    encoded = p.encodeTOTPKeyWithAlg(b"12345678901234567890")
+
+    assert encoded != None
+    assert len(encoded) == 32
+    assert encoded == "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"
+
+    encoded = p.encodeTOTPKeyWithAlg(b"12345678901234567890", ALG_SHA1)
+
+    assert encoded != None
+    assert len(encoded) == 32
+    assert encoded == "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ"
+
+    encoded = p.encodeTOTPKeyWithAlg(b"12345678901234567890123456789012", ALG_SHA256)
+
+    assert encoded != None
+    assert len(encoded) == 56
+    assert encoded == "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZA===="
+
+    encoded = p.encodeTOTPKeyWithAlg(b"1234567890123456789012345678901212345678901234567890123456789012", ALG_SHA512)
+
+    assert encoded != None
+    assert len(encoded) == 104
+    assert encoded == "GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZDCMRTGQ2TMNZYHEYDCMRTGQ2TMNZYHEYDCMRTGQ2TMNZYHEYDCMQ="
+
+def test_decodeTOTPKeyWithAlg(caplog) -> None:
+    caplog.set_level(logging.INFO)
+
+    decoded = p.decodeTOTPKeyWithAlg("GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ")
+
+    assert decoded != None
+    assert len(decoded) == 20
+    assert decoded == b"12345678901234567890"
+
+    decoded = p.decodeTOTPKeyWithAlg("GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQ", ALG_SHA1)
+
+    assert decoded != None
+    assert len(decoded) == 20
+    assert decoded == b"12345678901234567890"
+
+    decoded = p.decodeTOTPKeyWithAlg("GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZA====", ALG_SHA256)
+
+    assert decoded != None
+    assert len(decoded) == 32
+    assert decoded == b"12345678901234567890123456789012"
+
+    decoded = p.decodeTOTPKeyWithAlg("GEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZDGNBVGY3TQOJQGEZDCMRTGQ2TMNZYHEYDCMRTGQ2TMNZYHEYDCMRTGQ2TMNZYHEYDCMQ=", ALG_SHA512)
+
+    assert decoded != None
+    assert len(decoded) == 64
+    assert decoded == b"1234567890123456789012345678901212345678901234567890123456789012"