codekingpro/portable-devtools

1from __future__ import annotations2 3import base644import string5import struct6import typing as t7 8from .exc import BadData9 10 11def want_bytes(12    s: str | bytes, encoding: str = "utf-8", errors: str = "strict"13) -> bytes:14    if isinstance(s, str):15        s = s.encode(encoding, errors)16 17    return s18 19 20def base64_encode(string: str | bytes) -> bytes:21    """Base64 encode a string of bytes or text. The resulting bytes are22    safe to use in URLs.23    """24    string = want_bytes(string)25    return base64.urlsafe_b64encode(string).rstrip(b"=")26 27 28def base64_decode(string: str | bytes) -> bytes:29    """Base64 decode a URL-safe string of bytes or text. The result is30    bytes.31    """32    string = want_bytes(string, encoding="ascii", errors="ignore")33    string += b"=" * (-len(string) % 4)34 35    try:36        return base64.urlsafe_b64decode(string)37    except (TypeError, ValueError) as e:38        raise BadData("Invalid base64-encoded data") from e39 40 41# The alphabet used by base64.urlsafe_*42_base64_alphabet = f"{string.ascii_letters}{string.digits}-_=".encode("ascii")43 44_int64_struct = struct.Struct(">Q")45_int_to_bytes = _int64_struct.pack46_bytes_to_int = t.cast("t.Callable[[bytes], tuple[int]]", _int64_struct.unpack)47 48 49def int_to_bytes(num: int) -> bytes:50    return _int_to_bytes(num).lstrip(b"\x00")51 52 53def bytes_to_int(bytestr: bytes) -> int:54    return _bytes_to_int(bytestr.rjust(8, b"\x00"))[0]55