1  from __future__ import annotations
   2  
   3  import base64
   4  import string
   5  import struct
   6  import typing as t
   7  
   8  from .exc import BadData
   9  
  10  
  11  def 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 s
  18  
  19  
  20  def base64_encode(string: str | bytes) -> bytes:
  21      """Base64 encode a string of bytes or text. The resulting bytes are
  22      safe to use in URLs.
  23      """
  24      string = want_bytes(string)
  25      return base64.urlsafe_b64encode(string).rstrip(b"=")
  26  
  27  
  28  def base64_decode(string: str | bytes) -> bytes:
  29      """Base64 decode a URL-safe string of bytes or text. The result is
  30      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 e
  39  
  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.pack
  46  _bytes_to_int = t.cast("t.Callable[[bytes], tuple[int]]", _int64_struct.unpack)
  47  
  48  
  49  def int_to_bytes(num: int) -> bytes:
  50      return _int_to_bytes(num).lstrip(b"\x00")
  51  
  52  
  53  def bytes_to_int(bytestr: bytes) -> int:
  54      return _bytes_to_int(bytestr.rjust(8, b"\x00"))[0]
  55