encoding.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. import base64
  2. import string
  3. import struct
  4. import typing as _t
  5. from .exc import BadData
  6. _t_str_bytes = _t.Union[str, bytes]
  7. def want_bytes(
  8. s: _t_str_bytes, encoding: str = "utf-8", errors: str = "strict"
  9. ) -> bytes:
  10. if isinstance(s, str):
  11. s = s.encode(encoding, errors)
  12. return s
  13. def base64_encode(string: _t_str_bytes) -> bytes:
  14. """Base64 encode a string of bytes or text. The resulting bytes are
  15. safe to use in URLs.
  16. """
  17. string = want_bytes(string)
  18. return base64.urlsafe_b64encode(string).rstrip(b"=")
  19. def base64_decode(string: _t_str_bytes) -> bytes:
  20. """Base64 decode a URL-safe string of bytes or text. The result is
  21. bytes.
  22. """
  23. string = want_bytes(string, encoding="ascii", errors="ignore")
  24. string += b"=" * (-len(string) % 4)
  25. try:
  26. return base64.urlsafe_b64decode(string)
  27. except (TypeError, ValueError):
  28. raise BadData("Invalid base64-encoded data")
  29. # The alphabet used by base64.urlsafe_*
  30. _base64_alphabet = f"{string.ascii_letters}{string.digits}-_=".encode("ascii")
  31. _int64_struct = struct.Struct(">Q")
  32. _int_to_bytes = _int64_struct.pack
  33. _bytes_to_int = _t.cast("_t.Callable[[bytes], _t.Tuple[int]]", _int64_struct.unpack)
  34. def int_to_bytes(num: int) -> bytes:
  35. return _int_to_bytes(num).lstrip(b"\x00")
  36. def bytes_to_int(bytestr: bytes) -> int:
  37. return _bytes_to_int(bytestr.rjust(8, b"\x00"))[0]