fix(ciphers): operate on raw bytes in base32 and base85

ciphers/base32.py

@@ -9,35 +9,48 @@
 
 def base32_encode(data: bytes) -> bytes:
     """
+    >>> base32_encode(b"")
+    b''
     >>> base32_encode(b"Hello World!")
     b'JBSWY3DPEBLW64TMMQQQ===='
     >>> base32_encode(b"123456")
     b'GEZDGNBVGY======'
     >>> base32_encode(b"some long complex string")
     b'ONXW2ZJANRXW4ZZAMNXW24DMMV4CA43UOJUW4ZY='
+    >>> base32_encode("héllo".encode())
+    b'NDB2S3DMN4======'
     """
-    binary_data = "".join(bin(ord(d))[2:].zfill(8) for d in data.decode("utf-8"))
-    binary_data = binary_data.ljust(5 * ((len(binary_data) // 5) + 1), "0")
-    b32_chunks = map("".join, zip(*[iter(binary_data)] * 5))
+    if not data:
+        return b""
+    binary_data = "".join(bin(byte)[2:].zfill(8) for byte in data)
+    remainder = len(binary_data) % 5
+    if remainder != 0:
+        binary_data = binary_data + "0" * (5 - remainder)
+    b32_chunks = [binary_data[i : i + 5] for i in range(0, len(binary_data), 5)]
     b32_result = "".join(B32_CHARSET[int(chunk, 2)] for chunk in b32_chunks)
-    return bytes(b32_result.ljust(8 * ((len(b32_result) // 8) + 1), "="), "utf-8")
+    pad_len = (8 - (len(b32_result) % 8)) % 8
+    return (b32_result + "=" * pad_len).encode("ascii")
 
 
 def base32_decode(data: bytes) -> bytes:
     """
+    >>> base32_decode(b"")
+    b''
     >>> base32_decode(b'JBSWY3DPEBLW64TMMQQQ====')
     b'Hello World!'
     >>> base32_decode(b'GEZDGNBVGY======')
     b'123456'
     >>> base32_decode(b'ONXW2ZJANRXW4ZZAMNXW24DMMV4CA43UOJUW4ZY=')
     b'some long complex string'
+    >>> base32_decode(b'NDB2S3DMN4======').decode()
+    'héllo'
     """
-    binary_chunks = "".join(
-        bin(B32_CHARSET.index(_d))[2:].zfill(5)
-        for _d in data.decode("utf-8").strip("=")
-    )
-    binary_data = list(map("".join, zip(*[iter(binary_chunks)] * 8)))
-    return bytes("".join([chr(int(_d, 2)) for _d in binary_data]), "utf-8")
+    if not data:
+        return b""
+    s = data.decode("ascii").strip("=")
+    binary_chunks = "".join(bin(B32_CHARSET.index(c))[2:].zfill(5) for c in s)
+    num_bytes = len(binary_chunks) // 8
+    return bytes(int(binary_chunks[i : i + 8], 2) for i in range(0, num_bytes * 8, 8))
 
 
 if __name__ == "__main__":

ciphers/base85.py

@@ -6,7 +6,11 @@
 
 
 def _base10_to_85(d: int) -> str:
-    return "".join(chr(d % 85 + 33)) + _base10_to_85(d // 85) if d > 0 else ""
+    res = []
+    for _ in range(5):
+        res.append(chr(d % 85 + 33))
+        d //= 85
+    return "".join(res)
 
 
 def _base85_to_10(digits: list) -> int:
@@ -21,13 +25,22 @@ def ascii85_encode(data: bytes) -> bytes:
     b'0etOA2#'
     >>> ascii85_encode(b"base 85")
     b'@UX=h+?24'
+    >>> ascii85_encode("héllo".encode())
+    b'BZ$fcCi9'
     """
-    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))]
+    if not data:
+        return b""
+    binary_data = "".join(bin(byte)[2:].zfill(8) for byte in data)
+    remainder = len(binary_data) % 32
+    null_bits = (32 - remainder) % 32
+    binary_data = binary_data + "0" * null_bits
+    null_bytes = null_bits // 8
+
+    b85_chunks = [
+        int(binary_data[i : i + 32], 2) for i in range(0, len(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")
+    return bytes(result[:-null_bytes] if null_bytes > 0 else result, "utf-8")
 
 
 def ascii85_decode(data: bytes) -> bytes:
@@ -38,18 +51,22 @@ def ascii85_decode(data: bytes) -> bytes:
     b'12345'
     >>> ascii85_decode(b"@UX=h+?24")
     b'base 85'
+    >>> ascii85_decode(b'BZ$fcCi9').decode()
+    'héllo'
     """
-    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
+    if not data:
+        return b""
+    s = data.decode("ascii")
+    null_values = (5 - (len(s) % 5)) % 5
+    binary_data = s + "u" * null_values
+    b85_chunks = [binary_data[i : i + 5] for i in range(0, len(binary_data), 5)]
+    b85_segments = [[ord(c) - 33 for c in chunk] for chunk in b85_chunks]
+    results = [bin(_base85_to_10(chunk))[2:].zfill(32) for chunk in b85_segments]
+    byte_chunks = [
+        [int(r[j : j + 8], 2) for j in range(0, 32, 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")
+    raw_bytes = bytes(b for chunk in byte_chunks for b in chunk)
+    return raw_bytes[:-null_values] if null_values > 0 else raw_bytes
 
 
 if __name__ == "__main__":