1from __future__ import annotations
2
3import base64
4import string
5import struct
6import typing as t
7
8from .exc import BadData
9
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 s
18
19
20def 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
28def 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
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]