Merge pull request #57 from LuminousXLB/autofix

Autofix for similar letters

.gitignore

@@ -12,3 +12,4 @@ __pycache__/
 
 # Macos folder settings file
 .DS_Store
+.coverage

base58/__init__.py

@@ -65,21 +65,32 @@ def b58encode(
 
 
 @lru_cache()
-def _get_base58_decode_map(alphabet: bytes) -> Mapping[int, int]:
-    map = {char: index for index, char in enumerate(alphabet)}
-    return map
+def _get_base58_decode_map(alphabet: bytes,
+                           autofix: bool) -> Mapping[int, int]:
+    invmap = {char: index for index, char in enumerate(alphabet)}
+
+    if autofix:
+        groups = [b'0Oo', b'Il1']
+        for group in groups:
+            pivots = [c for c in group if c in invmap]
+            if len(pivots) == 1:
+                for alternative in group:
+                    invmap[alternative] = invmap[pivots[0]]
+
+    return invmap
 
 
 def b58decode_int(
-    v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET
+    v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET, *,
+    autofix: bool = False
 ) -> int:
     """
     Decode a Base58 encoded string as an integer
     """
     v = v.rstrip()
     v = scrub_input(v)
 
-    map = _get_base58_decode_map(alphabet)
+    map = _get_base58_decode_map(alphabet, autofix=autofix)
 
     decimal = 0
 
@@ -93,7 +104,8 @@ def b58decode_int(
 
 
 def b58decode(
-    v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET
+    v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET, *,
+    autofix: bool = False
 ) -> bytes:
     """
     Decode a Base58 encoded string
@@ -105,7 +117,7 @@ def b58decode(
     v = v.lstrip(alphabet[0:1])
     newlen = len(v)
 
-    acc = b58decode_int(v, alphabet=alphabet)
+    acc = b58decode_int(v, alphabet=alphabet, autofix=autofix)
 
     result = []
     while acc > 0:
@@ -128,11 +140,12 @@ def b58encode_check(
 
 
 def b58decode_check(
-    v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET
+    v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET, *,
+    autofix: bool = False
 ) -> bytes:
     '''Decode and verify the checksum of a Base58 encoded string'''
 
-    result = b58decode(v, alphabet=alphabet)
+    result = b58decode(v, alphabet=alphabet, autofix=autofix)
     result, check = result[:-4], result[-4:]
     digest = sha256(sha256(result).digest()).digest()
 

base58/__main__.py

@@ -1,8 +1,16 @@
 import argparse
 import sys
+from typing import Callable, Dict, Tuple
 
 from base58 import b58decode, b58decode_check, b58encode, b58encode_check
 
+_fmap = {
+    (False, False): b58encode,
+    (False, True): b58encode_check,
+    (True, False): b58decode,
+    (True, True): b58decode_check
+}  # type: Dict[Tuple[bool, bool], Callable[[bytes], bytes]]
+
 
 def main() -> None:
     '''Base58 encode or decode FILE, or standard input, to standard output.'''
@@ -26,12 +34,7 @@ def main() -> None:
         help='append a checksum before encoding')
 
     args = parser.parse_args()
-    fun = {
-        (False, False): b58encode,
-        (False, True): b58encode_check,
-        (True, False): b58decode,
-        (True, True): b58decode_check
-    }[(args.decode, args.check)]
+    fun = _fmap[(args.decode, args.check)]
 
     data = args.file.buffer.read()
 

test_base58.py

@@ -32,6 +32,11 @@ def test_simple_decode_bytes():
     assert_that(data, equal_to(b'hello world'))
 
 
+def test_autofix_decode_bytes():
+    data = b58decode(b'StVlDL6CwTryKyV', autofix=True)
+    assert_that(data, equal_to(b'hello world'))
+
+
 def test_leadingz_decode():
     data = b58decode('11StV1DL6CwTryKyV')
     assert_that(data, equal_to(b'\0\0hello world'))
@@ -66,6 +71,20 @@ def test_check_str():
     assert_that(back, equal_to(b'hello world'))
 
 
+def test_autofix_check_str():
+    data = '3vQB7B6MrGQZaxCuFg4Oh'
+    back = b58decode_check(data, autofix=True)
+    assert_that(back, equal_to(b'hello world'))
+
+
+def test_autofix_not_applicable_check_str():
+    charset = BITCOIN_ALPHABET.replace(b'x', b'l')
+    msg = b'hello world'
+    enc = b58encode_check(msg).replace(b'x', b'l').replace(b'o', b'0')
+    back = b58decode_check(enc, alphabet=charset, autofix=True)
+    assert_that(back, equal_to(msg))
+
+
 def test_check_failure():
     data = '3vQB7B6MrGQZaxCuFg4oH'
     assert_that(calling(b58decode_check).with_args(data), raises(ValueError))