diff --git a/hdwallet/libs/base58.py b/hdwallet/libs/base58.py index 98f0a39b..c025f8b5 100644 --- a/hdwallet/libs/base58.py +++ b/hdwallet/libs/base58.py @@ -67,15 +67,18 @@ def decode(data): data = bytes(data, "ascii") val = 0 - for (i, c) in enumerate(data[::-1]): - val += __base58_alphabet_bytes.find(c) * (__base58_radix ** i) + prefix = 0 + for c in data: + val = (val * __base58_radix) + __base58_alphabet_bytes.find(c) + if val == 0: + prefix += 1 dec = bytearray() - while val >= 256: + while val > 0: val, mod = divmod(val, 256) dec.append(mod) - if val: - dec.append(val) + + dec.extend(bytearray(prefix)) return bytes(dec[::-1]) diff --git a/tests/test_base58.py b/tests/test_base58.py index 2b602c08..f321e765 100644 --- a/tests/test_base58.py +++ b/tests/test_base58.py @@ -7,7 +7,7 @@ import pytest from hdwallet.libs.base58 import ( - check_encode, check_decode, string_to_int + check_encode, check_decode, decode, encode, string_to_int ) @@ -29,3 +29,9 @@ def test_base58(): with pytest.raises(TypeError, match="string argument without an encoding"): assert string_to_int(str("meherett")) + + + assert decode("111233QC4") == b'\x00\x00\x00(\x7f\xb4\xcd' + + + assert encode(decode("111233QC4")) == "111233QC4"