Coverage for an_website / utils / token.py: 92.405%

1# This program is free software: you can redistribute it and/or modify 

2# it under the terms of the GNU Affero General Public License as 

3# published by the Free Software Foundation, either version 3 of the 

4# License, or (at your option) any later version. 

5# 

6# This program is distributed in the hope that it will be useful, 

7# but WITHOUT ANY WARRANTY; without even the implied warranty of 

8# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

9# GNU Affero General Public License for more details. 

10# 

11# You should have received a copy of the GNU Affero General Public License 

12# along with this program. If not, see <https://www.gnu.org/licenses/>. 

13 

14"""A module providing special auth tokens.""" 

15 

16import hmac 

17import math 

18from base64 import b64decode, b64encode 

19from datetime import datetime 

20from hashlib import blake2b 

21from typing import ClassVar, Literal, NamedTuple, TypeAlias, TypeGuard, get_args 

22 

23from .utils import Permission 

24 

25TokenVersion: TypeAlias = Literal["0"] 

26SUPPORTED_TOKEN_VERSIONS: tuple[TokenVersion, ...] = get_args(TokenVersion) 

27 

28 

29class ParseResult(NamedTuple): 

30 """The class representing a token.""" 

31 

32 token: str 

33 permissions: Permission 

34 valid_until: datetime 

35 salt: bytes 

36 

37 

38class InvalidTokenError(Exception): 

39 """Exception thrown for invalid or expired tokens.""" 

40 

41 

42class InvalidTokenVersionError(InvalidTokenError): 

43 """Exception thrown when the token has an invalid version.""" 

44 

45 SUPPORTED_TOKEN_VERSIONS: ClassVar = SUPPORTED_TOKEN_VERSIONS 

46 

47 

48def is_supported_version(version: str) -> TypeGuard[TokenVersion]: 

49 """Check whether the argument is a supported token version.""" 

50 return version in SUPPORTED_TOKEN_VERSIONS 

51 

52 

53def _split_token(token: str) -> tuple[TokenVersion, str]: 

54 """Split a token into version and the body of the token.""" 

55 if not token: 

56 raise InvalidTokenError() 

57 

58 version = token[0] 

59 if is_supported_version(version): 

60 return version, token[1:] 

61 

62 raise InvalidTokenVersionError() 

63 

64 

65def parse_token( # pylint: disable=inconsistent-return-statements 

66 token: str, 

67 *, 

68 secret: bytes | str, 

69 verify_time: bool = True, 

70) -> ParseResult: 

71 """Parse an auth token.""" 

72 secret = secret.encode("UTF-8") if isinstance(secret, str) else secret 

73 version, token_body = _split_token(token) 

74 try: 

75 if version == "0": 

76 return _parse_token_v0(token_body, secret, verify_time=verify_time) 

77 except InvalidTokenError: 

78 raise 

79 except Exception as exc: 

80 raise InvalidTokenError from exc 

81 

82 

83def create_token( # pylint: disable=too-many-arguments 

84 permissions: Permission, 

85 *, 

86 secret: bytes | str, 

87 duration: int, 

88 start: None | datetime = None, 

89 salt: None | bytes | str = None, 

90 version: TokenVersion = SUPPORTED_TOKEN_VERSIONS[-1], 

91) -> ParseResult: 

92 """Create an auth token.""" 

93 secret = secret.encode("UTF-8") if isinstance(secret, str) else secret 

94 start = datetime.now() if start is None else start 

95 salt = salt.encode("UTF-8") if isinstance(salt, str) else salt or b"" 

96 token: str 

97 if version == "0": 

98 token = _create_token_body_v0( 

99 permissions, secret, duration, start, salt 

100 ) 

101 

102 return parse_token(version + token, secret=secret, verify_time=False) 

103 

104 

105def int_to_bytes(number: int, length: int, signed: bool = False) -> bytes: 

106 """Convert an int to bytes.""" 

107 return number.to_bytes(length, "big", signed=signed) 

108 

109 

110def bytes_to_int(bytes_: bytes, signed: bool = False) -> int: 

111 """Convert an int to bytes.""" 

112 if not bytes_: 

113 raise ValueError("Can't convert empty bytes to int") 

114 return int.from_bytes(bytes_, "big", signed=signed) 

115 

116 

117def _parse_token_v0( 

118 token_body: str, secret: bytes, *, verify_time: bool = True 

119) -> ParseResult: 

120 """Parse an auth token of version 0.""" 

121 data: bytes 

122 try: 

123 data = b64decode(token_body, validate=True) 

124 except ValueError as err: 

125 raise InvalidTokenError() from err 

126 data, hash_ = data[:-48], data[-48:] 

127 if not data or not hash_: 

128 raise InvalidTokenError() 

129 

130 if not hmac.compare_digest(hmac.digest(secret, data, "SHA3-384"), hash_): 

131 raise InvalidTokenError() 

132 

133 try: 

134 data, start = data[:-5], bytes_to_int(data[-5:]) 

135 data, duration = data[:-5], bytes_to_int(data[-5:]) 

136 permissions, salt = bytes_to_int(data[:-6]), data[-6:] 

137 except ValueError as err: 

138 raise InvalidTokenError() from err 

139 

140 now = int(datetime.now().timestamp()) 

141 if verify_time and (now < start or start + duration < now): 

142 raise InvalidTokenError() 

143 

144 return ParseResult( 

145 "0" + token_body, 

146 Permission(permissions), 

147 datetime.fromtimestamp(start + duration), 

148 salt, 

149 ) 

150 

151 

152def _create_token_body_v0( 

153 permissions: Permission, 

154 secret: bytes, 

155 duration: int, 

156 start: datetime, 

157 salt: bytes, 

158) -> str: 

159 """Create an auth token of version 0.""" 

160 if not salt: 

161 salt = blake2b( 

162 int_to_bytes(int(start.timestamp() - duration), 5), digest_size=6 

163 ).digest() 

164 elif len(salt) < 6: 

165 salt = b"U" * (6 - len(salt)) + salt 

166 elif len(salt) > 6: 

167 salt = salt[:6] 

168 

169 parts = ( 

170 int_to_bytes(permissions, math.ceil(len(Permission) / 8)), 

171 salt, 

172 int_to_bytes(duration, 5), 

173 int_to_bytes(int(start.timestamp()), 5), 

174 ) 

175 data: bytes = b"".join(parts) 

176 

177 len_token = len(data) + 384 // 8 

178 if len_token % 3: 

179 data = int_to_bytes(0, 3 - (len_token % 3)) + data 

180 

181 hash_ = hmac.digest(secret, data, "SHA3-384") 

182 return b64encode(data + hash_).decode("UTF-8")