fix(ciphers): operate on raw bytes in base85 to fix corruption and decode crashes

ciphers/base85.py

@@ -5,51 +5,87 @@
 """
 
 
-def _base10_to_85(d: int) -> str:
-    return "".join(chr(d % 85 + 33)) + _base10_to_85(d // 85) if d > 0 else ""
-
-
-def _base85_to_10(digits: list) -> int:
-    return sum(char * 85**i for i, char in enumerate(reversed(digits)))
-
-
 def ascii85_encode(data: bytes) -> bytes:
     """
+    Encode bytes into Ascii85 (Base85) bytes.
+
     >>> ascii85_encode(b"")
     b''
     >>> ascii85_encode(b"12345")
     b'0etOA2#'
     >>> ascii85_encode(b"base 85")
     b'@UX=h+?24'
+    >>> ascii85_encode(b"\\xff\\xfe\\xfd")
+    b's8Mr'
+    >>> import base64
+    >>> all(
+    ...     ascii85_encode(bytes([i])) == base64.a85encode(bytes([i]))
+    ...     for i in range(256)
+    ... )
+    True
     """
-    binary_data = "".join(bin(ord(d))[2:].zfill(8) for d in data.decode("utf-8"))
-    null_values = (32 * ((len(binary_data) // 32) + 1) - len(binary_data)) // 8
-    binary_data = binary_data.ljust(32 * ((len(binary_data) // 32) + 1), "0")
-    b85_chunks = [int(_s, 2) for _s in map("".join, zip(*[iter(binary_data)] * 32))]
-    result = "".join(_base10_to_85(chunk)[::-1] for chunk in b85_chunks)
-    return bytes(result[:-null_values] if null_values % 4 != 0 else result, "utf-8")
+    if not data:
+        return b""
+    result = bytearray()
+    for i in range(0, len(data), 4):
+        chunk = data[i : i + 4]
+        pad = 4 - len(chunk)
+        chunk_padded = chunk + b"\x00" * pad
+        val = int.from_bytes(chunk_padded, "big")
+        chars = bytearray(5)
+        for j in range(4, -1, -1):
+            chars[j] = (val % 85) + 33
+            val //= 85
+        if pad:
+            result.extend(chars[: 5 - pad])
+        else:
+            result.extend(chars)
+    return bytes(result)
 
 
 def ascii85_decode(data: bytes) -> bytes:
     """
+    Decode Ascii85 (Base85) bytes back to original bytes.
+
     >>> ascii85_decode(b"")
     b''
     >>> ascii85_decode(b"0etOA2#")
     b'12345'
     >>> ascii85_decode(b"@UX=h+?24")
     b'base 85'
+    >>> ascii85_decode(b"s8Mr")
+    b'\\xff\\xfe\\xfd'
+    >>> ascii85_decode(b"!")
+    Traceback (most recent call last):
+        ...
+    ValueError: Invalid Ascii85 length
+    >>> ascii85_decode(b"~~")
+    Traceback (most recent call last):
+        ...
+    ValueError: Character '~' out of Ascii85 range
     """
-    null_values = 5 * ((len(data) // 5) + 1) - len(data)
-    binary_data = data.decode("utf-8") + "u" * null_values
-    b85_chunks = map("".join, zip(*[iter(binary_data)] * 5))
-    b85_segments = [[ord(_s) - 33 for _s in chunk] for chunk in b85_chunks]
-    results = [bin(_base85_to_10(chunk))[2::].zfill(32) for chunk in b85_segments]
-    char_chunks = [
-        [chr(int(_s, 2)) for _s in map("".join, zip(*[iter(r)] * 8))] for r in results
-    ]
-    result = "".join("".join(char) for char in char_chunks)
-    offset = int(null_values % 5 == 0)
-    return bytes(result[: offset - null_values], "utf-8")
+    if not data:
+        return b""
+    rem = len(data) % 5
+    pad = (5 - rem) % 5
+    if rem == 1:
+        raise ValueError("Invalid Ascii85 length")
+    padded = data + b"u" * pad
+    result = bytearray()
+    for i in range(0, len(padded), 5):
+        chunk = padded[i : i + 5]
+        val = 0
+        for b in chunk:
+            if not (33 <= b <= 117):
+                msg = f"Character {chr(b)!r} out of Ascii85 range"
+                raise ValueError(msg)
+            val = val * 85 + (b - 33)
+        if val > 0xFFFFFFFF:
+            raise ValueError("Ascii85 chunk value overflow")
+        result.extend(val.to_bytes(4, "big"))
+    if pad:
+        return bytes(result[:-pad])
+    return bytes(result)
 
 
 if __name__ == "__main__":