Fixed some base encodings and performance issues

Author: dhondta

docs/pages/guessing.md

@@ -11,9 +11,12 @@ The following parameters are tunable:
 - `stop_func`: can be a function or a regular expression to be matched (automatically converted to a function that uses the `re` module) ; by default, checks if all input characters are printable.
 - `min_depth`: the minimum depth for the tree search (allows to avoid a bit of overhead while checking the current decoded output at a depth with the stop function when we are sure it should not be the right result) ; by default 0.
 - `max_depth`: the maximum depth for the tree search ; by default 5.
-- `codec_categories`: a string indicating a codec [category](features.md#list-codecs) or a list of [category](features.md#list-codecs) strings ; by default, `None`, meaning the whole [categories](features.md#list-codecs) (very slow).
+- `include`: the [categories](features.md#list-codecs), codecs or encodings to be considered, as a single name, a list of names or a dictionary of lists of names per depth (from 0, -1 meaning any depth) ; by default, `None`, meaning all of them (very slow).
+- `exclude`: the [categories](features.md#list-codecs), codecs or encodings not to be considered, in the same formats ; by default, `None`, meaning none of them.
 - `found`: a list or tuple of currently found encodings that can be used to save time if the first decoding steps are known ; by default, an empty tuple.
 
+The result is a dictionary of the decoded strings by tuple of encodings, in order of application.
+
 A simple example for a 1-stage base64-encoded string:
 
 ```python
@@ -25,7 +28,7 @@ An example of a 2-stages base64- then base62-encoded string:
 
 ```python
 >>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7")
-{('base62',): 'VGhpcyBpcyBhIHRlc3Q='}
+{('chaocipher',): 'NGH3Hb8tYyJWGEEseSOq28aTtS7'}
 ```
 
 In the second example, we can see that the given encoded string is not decoded as expected. This is the case because the (default) stop condition is too broad and stops if all the characters of the output are printable. If we have a prior knowledge on what we should expect, we can input a simple string or a regex:
@@ -43,37 +46,43 @@ In the second example, we can see that the given encoded string is not decoded a
 {('base62', 'base64'): 'This is a test'}
 ```
 
-In this example, the string "*test*" is converted to a function that uses this string as regular expression. Instead of a string, we can also pass a function. For this purpose, standard [stop functions](#available-stop-functions) are predefined. So, we can for instance use `stopfunc.lang_en` to stop when we find something that is English. Note that working this way gives lots of false positives if the text is very short like in the example case. That's why the `codec_categories` argument is used to only consider baseX codecs. This is also demonstrated in the next examples.
+In this example, the string "*test*" is converted to a function that uses this string as regular expression. Instead of a string, we can also pass a function. For this purpose, standard [stop functions](#available-stop-functions) are predefined. So, we can for instance use `stopfunc.lang_en` to stop when we find something that is English. Note that working this way gives lots of false positives if the text is very short like in the example case. That's why the `include` argument is used to only consider baseX codecs. This is also demonstrated in the next examples.
 
 ```python
 >>> codext.stopfunc._reload_lang("langdetect")
->>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", codext.stopfunc.lang_en, codec_categories="base")
-('This is a test', ('base62', 'base64'))
+>>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", codext.stopfunc.lang_en, include="base")
+{('base62', 'base64'): 'This is a test'}
 ```
 
 If we know the first encoding, we can set this in the `found` parameter to save time:
 
 ```python
 >>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "test", found=["base62"])
-('This is a test', ('base62', 'base64'))
+{('base62', 'base64'): 'This is a test'}
 ```
 
-If we are sure that only `base` (which is a valid [category](features.md#list-codecs)) encodings are used, we can restrict the tree search using the `codec_categories` parameter to save time:
+If we are sure that only `base` (which is a valid [category](features.md#list-codecs)) encodings are used, we can restrict the tree search using the `include` parameter to save time, or even give the encodings to be considered at each depth, while the `exclude` parameter removes some of them:
 
 ```python
->>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "test", codec_categories="base")
-('This is a test', ('base62', 'base64'))
+>>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "test", include="base")
+{('base62', 'base64'): 'This is a test'}
+>>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "test", include={0: ["base62"], 1: ["base64"]})
+{('base62', 'base64'): 'This is a test'}
+>>> codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "test", include="base", exclude="base62")
+{}
 ```
 
 Another example of 2-stages encoded string:
 
 ```python
 >>> codext.guess("LSAuLi4uIC4uIC4uLiAvIC4uIC4uLiAvIC4tIC8gLSAuIC4uLiAt", "test")
-('this is a test', ('base64', 'morse'))
->>> codext.guess("LSAuLi4uIC4uIC4uLiAvIC4uIC4uLiAvIC4tIC8gLSAuIC4uLiAt", "test", codec_categories=["base", "language"])
-('this is a test', ('base64', 'morse'))
+{('base64', 'morse'): 'this is a test'}
+>>> codext.guess("LSAuLi4uIC4uIC4uLiAvIC4uIC4uLiAvIC4tIC8gLSAuIC4uLiAt", "test", include=["base", "language"])
+{('base64', 'morse'): 'this is a test'}
 ```
 
+With the [CLI tool](cli.md), these parameters are given with the `-i`/`--include` and `-e`/`--exclude` options, the names being possibly prefixed with a depth (e.g. `codext guess -i 0:base62 1:base64`).
+
 When multiple results are expected, `stop` and `show` arguments can be used respectively to avoid stopping while finding a result and to display the intermediate result.
 
 !!! warning "Computation time"

src/codext/base/_base.py

@@ -11,9 +11,10 @@
  does).
 """
 from argparse import ArgumentParser, RawTextHelpFormatter
+from functools import lru_cache
 from math import log
 from string import ascii_lowercase as lower, ascii_uppercase as upper, digits, printable
-from sys import stdout
+from sys import stderr, stdout
 from textwrap import wrap as wraptext
 from types import FunctionType, MethodType
 
@@ -116,10 +117,57 @@ def _get_charset(charset, p=""):
         # or simply rely on key ''
         if default is not None:
             return default
-    raise ValueError("Bad charset descriptor ('%s')" % p)
+    raise ValueError(f"Bad charset descriptor ('{p}')")
 
 
 # generic base en/decoding functions
+# conversions between integers and digits in base N split the numbers with powers of N (divide and conquer), as repeated
+#  divisions or multiplications by N take a quadratic time with large numbers ; below this number of digits, the
+#  simple loops are faster
+_SMALL = 64
+
+
+@lru_cache(maxsize=128)
+def _power(n, k):
+    return n ** k
+
+
+def _to_digits(i, n, width=0):
+    """ Digits of the integer i in base N (as indices in the charset), left-padded with zeros to the given width (by
+         default, the smallest power of 2 holding i, the leading zeros being then removed). """
+    if (width or _SMALL + 1) <= _SMALL or not width and i < _power(n, _SMALL):
+        d = []
+        while i:
+            i, c = divmod(i, n)
+            d.append(c)
+        return [0] * (width - len(d)) + d[::-1]
+    top = not width
+    if top:
+        width = 2 * _SMALL
+        while _power(n, width) <= i:
+            width *= 2
+    half = width // 2
+    hi, lo = divmod(i, _power(n, half))
+    d = _to_digits(hi, n, width - half) + _to_digits(lo, n, half)
+    if top:
+        k = 0
+        while d[k] == 0:
+            k += 1
+        d = d[k:]
+    return d
+
+
+def _from_digits(d, n):
+    """ Integer of the digits in base N (as indices in the charset), combining halves as hi * N^len(lo) + lo. """
+    if len(d) <= _SMALL:
+        i = 0
+        for c in d:
+            i = i * n + c
+        return i
+    half = len(d) // 2
+    return _from_digits(d[:half], n) * _power(n, len(d) - half) + _from_digits(d[half:], n)
+
+
 def base_encode(input, charset, errors="strict", exc=BaseEncodeError):
     """ Base-10 to base-N encoding.
     
@@ -128,18 +176,16 @@ def base_encode(input, charset, errors="strict", exc=BaseEncodeError):
     :param errors:  errors handling marker
     :param exc:     exception to be raised in case of error
     """
-    i, n, r = input if isinstance(input, int) else s2i(input), len(charset), ""
+    i, n = input if isinstance(input, int) else s2i(input) if input else 0, len(charset)
     if n == 1:
         if i > SIZE_LIMIT:
             raise InputSizeLimitError("Input exceeded size limit")
         return i * charset[0]
     # keep this fast-path only for a non-standard 10-character charset ; the digits
-    # charset uses the generic loop below (leading zeros, bignums, no str/int round-trip)
+    # charset uses the generic conversion below (leading zeros, bignums, no str/int round-trip)
     if n == 10 and charset != digits:
         return "".join(charset[int(x)] for x in str(i))
-    while i > 0:
-        i, c = divmod(i, n)
-        r = charset[c] + r
+    r = "".join(charset[c] for c in _to_digits(i, n)) if i > 0 else ""
     # preserve leading zero bytes: big-integer bases such as Base58 map each
     # leading null byte of the input to a leading charset[0] character
     if not isinstance(input, int):
@@ -155,18 +201,23 @@ def base_decode(input, charset, errors="strict", exc=BaseDecodeError):
     :param errors:  errors handling marker
     :param exc:     exception to be raised in case of error
     """
-    i, n, dec = 0, len(charset), lambda n: base_encode(n, [chr(x) for x in range(256)], errors, exc)
+    n = len(charset)
     if n == 1:
         return i2s(len(input))
     if n == 10 and charset != digits:
         return "".join(str(charset.index(c)) for c in input)
+    # index of each character (the first one, if the charset holds a character more than once)
+    index, d = {c: k for k, c in reversed(list(enumerate(charset)))}, []
     for k, c in enumerate(input):
         try:
-            i = i * n + charset.index(c)
-        except ValueError:
-            handle_error("base", errors, exc, decode=True)(c, k, dec(i), "base%d" % n)
+            d.append(index[c])
+        except KeyError:
+            # no output decoded up to the error, as the digits before it are not the beginning of the number
+            handle_error("base", errors, exc, decode=True)(c, k, "", f"base{n}")
+    i = _from_digits(d, n)
     # restore the leading zero bytes encoded as leading charset[0] characters
-    return chr(0) * (len(input) - len(input.lstrip(charset[0]))) + dec(i)
+    return chr(0) * (len(input) - len(input.lstrip(charset[0]))) + i.to_bytes((i.bit_length() + 7) // 8,
+                                                                              "big").decode("latin-1")
 
 
 # base codec factory functions
@@ -231,7 +282,7 @@ def _decode(input, errors="strict"):
         return _decode
 
     add("base", encode, decode, r"^base[-_]?([2-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(?:[-_]generic)?$",
-        guess=["base%d-generic" % i for i in range(2, 255)], entropy=lambda e, n: log(int(n.split("-")[0][4:]), 2),
+        guess=[f"base{i}-generic" for i in range(2, 255)], entropy=lambda e, n: log(int(n.split("-")[0][4:]), 2),
         len_charset=lambda n: int(n.split("-")[0][4:]), printables_rate=1., category="base-generic", penalty=.4,
         expansion_factor=lambda f, n: (EXPANSION_FACTOR(int(n.split("-")[0][4:])), .05))
 
@@ -294,7 +345,10 @@ def _main():
             c = f(c, "base" + base + ["", "-inv"][getattr(args, "invert", False)],
                   ["strict", "ignore"][args.ignore_garbage])
         except Exception as err:
-            print("%sbase%s: invalid input" % (getattr(err, "output", ""), base))
+            # as GNU's base64, write the output decoded up to the error (if any), then report the error on stderr
+            stdout.buffer.write(b(getattr(err, "output", "") or ""))
+            stdout.flush()
+            print(f"base{base}: invalid input", file=stderr)
             return 1
         if args.decode:
             stdout.buffer.write(c)

src/codext/base/_base2n.py

@@ -69,7 +69,7 @@ def base2n_decode(string, charset, errors="strict"):
     :param charset: base-N characters set
     :param errors:  errors handling marker
     """
-    bs, r, n = "", "", len(charset)
+    bs, r, n, nbits = "", "", len(charset), 0
     # particular case: for hex, ensure the right case in the charset ; not that this way, if mixed cases are used, it
     #                   will trigger an error (this is the expected behavior)
     if n == 16:
@@ -81,16 +81,24 @@ def base2n_decode(string, charset, errors="strict"):
     # find the number of bits for the given character set and the number of padding characters
     nb_in = int(log(n, 2))
     n_pad = len(string) - len(string.rstrip("="))
+    def partial(pos):
+        """ Output decoded up to an error at the given position, as with GNU's base64 and base32: the whole bytes of the
+             bits of the characters before the error (excluding the padding characters), except that, with base32, a
+             group of 8 characters cut by the end of the input gives no output (as with gnulib's decode_8). """
+        out = (r + "".join(chr(int(bs[k:k+8], 2)) for k in range(0, len(bs) // 8 * 8, 8)))[:nbits // 8]
+        return out[:pos // 8 * 5] if n == 32 and len(string) - pos // 8 * 8 < 8 else out
     # iterate over the characters, mapping them to the character set and converting the resulting bits to 8-bits chars
     for i, c in enumerate(string):
         if c == "=":
             bs += "0" * nb_in
         else:
             try:
                 bs += ("{:0>%d}" % nb_in).format(_bin(charset.index(c))[2:])
+                nbits += nb_in
             except ValueError:
                 if errors == "strict":
                     e = Base2NDecodeError("'base%d' codec can't decode character '%s' in position %d" % (n, c, i))
+                    e.output = partial(i)
                     e.__cause__ = e  # block exceptions chaining
                     raise e
                 elif errors == "replace":
@@ -105,7 +113,9 @@ def base2n_decode(string, charset, errors="strict"):
     # if the number of bits is not multiple of 8 bits, it could mean a bad padding
     if len(bs) != 8:
         if errors == "strict":
-            raise Base2NDecodeError("Incorrect padding")
+            e = Base2NDecodeError("Incorrect padding")
+            e.output = partial(len(string))
+            raise e
         elif errors in ["replace", "ignore"]:
             pass
         else:

src/codext/base/base45.py

@@ -89,6 +89,9 @@ def decode(text, errors="strict"):
     return decode
 
 
-add("base45", base45_encode, base45_decode, r"^base[-_]?45(|[-_]inv(?:erted)?)$", expansion_factor=1.5)
+# the expansion factor is 1.5 for pairs of bytes, up to 2 for a single byte (hence the tolerance for short inputs),
+#  and the charset has 45 characters
+add("base45", base45_encode, base45_decode, r"^base[-_]?45(|[-_]inv(?:erted)?)$", expansion_factor=(1.5, .17),
+    len_charset=45)
 main = main(45, "<https://datatracker.ietf.org/doc/html/draft-faltstrom-base45-04.txt>")
 

tests/test_base.py

@@ -3,6 +3,7 @@
 """Base codecs tests.
 
 """
+import codext
 import sys
 from unittest import TestCase
 
@@ -288,6 +289,10 @@ def test_codec_base45(self):
         self.assertEqual(codecs.decode(b"OBQR0", "base45"), b"\xcf\xb1\x1b")
         for data in [b"\xff\xfe", b"hello", b"\x00", b"\x80\x81\x82\x83\x84"]:
             self.assertEqual(codecs.decode(codecs.encode(data, "base45"), "base45"), data)
+        # the ranking heuristic puts base45 first, also for a single byte (whose expansion factor is 2 instead of 1.5)
+        for s in ["test", "The quick brown fox jumps over the lazy dog"]:
+            self.assertTrue(codext.rank(codecs.encode(s, "base45"), limit=1)[0][1].startswith("base45"))
+        self.assertTrue(any(n.startswith("base45") for _, n in codext.rank(codecs.encode("a", "base45"), limit=50)))
 
     def test_codec_base_generic(self):
         for n in range(2, 255):
@@ -296,6 +301,17 @@ def test_codec_base_generic(self):
         self.assertRaises(LookupError, codecs.decode, "test", "base0-generic")
         self.assertRaises(LookupError, codecs.decode, "test", "base1-generic")
         self.assertRaises(LookupError, codecs.decode, "test", "base256-generic")
+        # empty inputs, a known vector, integers and large inputs (converted by splitting with powers of the base)
+        self.assertEqual(codecs.encode("", "base58-generic"), "")
+        self.assertEqual(codecs.encode("Hello World!", "base58"), "2NEpo7TZRRrLZSi2U")
+        from codext.base._base import base_decode, base_encode
+        cs = _generate_charset(58)
+        for i in [0, 57, 58, 58 ** 200 - 1, 58 ** 200]:
+            self.assertEqual(base_decode(base_encode(i, cs), cs), i2s(i) if i else "")
+        data = "\x00\x00" + "".join(chr(x % 256) for x in range(20000))
+        for n in [3, 58, 200]:
+            bn = f"base{n}-generic"
+            self.assertEqual(codecs.decode(codecs.encode(data, bn), bn), data)
 
     def test_codec_base1024(self):
         bad = codecs.encode("Hello!", "ecoji-v1") + "\U0001f170\u23f0\U0001f170\U0001f170"
@@ -307,6 +323,26 @@ def test_codec_base1024(self):
         bad = codecs.encode("Hello!", "ecoji").replace("\u2615", "!\u2615")
         self.assertEqual(codecs.decode(bad, "ecoji", "ignore"), "Hello!")
 
+    def test_base_decoding_partial_output(self):
+        import subprocess
+        # as GNU's base64 and base32, the output decoded up to an error is given (with the exception), base32 giving no
+        #  output for a group of 8 characters cut by the end of the input, and big-integer bases giving none at all
+        for s, enc, out in [("SGVsbG8gV29ybGQ$IQ==", "base64", "Hello World"),
+                            ("SGVsbG8gV29ybGQh$", "base64", "Hello World!"), ("SGVsbG8=$$$$", "base64", "Hello"),
+                            ("SGVsbG8", "base64", "Hello"), ("S$GVs", "base64", ""), ("JBSWY3D$", "base32", "Hell"),
+                            ("JBSWY3DPEB$", "base32", "Hello"), ("JBSWY3D", "base32", ""),
+                            ("JBSWY3DPEBLW6$4TM", "base32", "Hello Wo"), ("2NEpo7TZ0RRrLZSi2U", "base58", "")]:
+            with self.assertRaises(ValueError) as cm:
+                codecs.decode(s, enc)
+            self.assertEqual(getattr(cm.exception, "output", ""), out)
+        # the CLI writes this output to stdout and the error to stderr
+        src = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src")
+        env = {**os.environ, 'PYTHONPATH': src + os.pathsep + os.environ.get('PYTHONPATH', "")}
+        cmd = [sys.executable, "-c", "import sys ; from codext.base.baseN import main64 ; "
+               "sys.argv = ['base64', '-d'] ; sys.exit(main64())"]
+        r = subprocess.run(cmd, input=b"SGVsbG8gV29ybGQ$IQ==", capture_output=True, env=env)
+        self.assertEqual((r.stdout, r.stderr.strip(), r.returncode), (b"Hello World", b"base64: invalid input", 1))
+    
     def test_base_main(self):
         tmp = sys.argv[:]
         tfile = "test-base-main.txt"