Coverage for src/taipanstack/security/jwt.py: 100%

1""" 

2Secure JWT Utility module. 

3 

4Provides explicitly secure wrappers around PyJWT encoding and decoding, 

5enforcing strict validation of algorithms, expiration, and audience claims. 

6All operations return ``Result`` types. 

7""" 

8 

9import secrets 

10from collections.abc import Iterable 

11from typing import TYPE_CHECKING, TypeAlias 

12 

13if TYPE_CHECKING: 

14 import jwt 

15 

16import jwt 

17from jwt.exceptions import PyJWTError 

18 

19from taipanstack.core.result import safe_from 

20 

21__all__ = ["decode_jwt", "encode_jwt"] 

22 

23JWTPayload: TypeAlias = dict[str, object] 

24 

25MAX_JWT_ALGORITHMS = 10 

26MAX_JWT_ALGORITHM_LENGTH = 50 

27MAX_JWT_AUDIENCES = 10 

28MAX_JWT_AUDIENCE_LENGTH = 1024 

29 

30 

31def _validate_encode_payload(payload: object) -> None: 

32 if not isinstance(payload, dict): 

33 raise TypeError("Payload must be a dictionary") 

34 

35 

36def _validate_encode_secret(secret_key: object) -> None: 

37 if not isinstance(secret_key, str): 

38 raise TypeError("Secret must be a string") 

39 

40 

41def _validate_encode_algorithm(algorithm: object) -> None: 

42 if not isinstance(algorithm, str): 

43 raise TypeError("Algorithm must be a string") 

44 

45 if len(algorithm) > MAX_JWT_ALGORITHM_LENGTH: 

46 raise ValueError("Algorithm string is too long") 

47 

48 alg_lower = algorithm.strip().lower() 

49 if alg_lower.isascii() and secrets.compare_digest(alg_lower, "none"): 

50 raise ValueError('Algorithm "none" is explicitly disallowed.') 

51 

52 

53@safe_from( 

54 PyJWTError, 

55 ValueError, 

56 TypeError, 

57 NotImplementedError, 

58 KeyError, 

59 AttributeError, 

60 Exception, 

61) 

62def encode_jwt( 

63 payload: JWTPayload, 

64 secret_key: str, 

65 algorithm: str = "HS256", 

66) -> str: 

67 """Encode a payload into a JWT securely. 

68 

69 Explicitly rejects the "none" algorithm to prevent bypass vulnerabilities. 

70 

71 Args: 

72 payload: Dictionary containing the JWT claims. 

73 secret_key: The secret key for signing the token. 

74 algorithm: The signing algorithm (default "HS256"). 

75 

76 Returns: 

77 The encoded JWT string. 

78 

79 Raises: 

80 ValueError: If the "none" algorithm is specified. 

81 PyJWTError: If encoding fails. 

82 

83 """ 

84 _validate_encode_payload(payload) 

85 _validate_encode_secret(secret_key) 

86 _validate_encode_algorithm(algorithm) 

87 

88 return jwt.encode(payload, secret_key, algorithm=algorithm) # nosem 

89 

90 

91def _check_disallowed_algorithm(alg: str) -> None: 

92 if not isinstance(alg, str): 

93 raise TypeError("Algorithm must be a string") 

94 

95 if len(alg) > MAX_JWT_ALGORITHM_LENGTH: 

96 raise ValueError("Algorithm string is too long") 

97 

98 alg_lower = alg.strip().lower() 

99 if alg_lower.isascii() and secrets.compare_digest(alg_lower, "none"): 

100 msg = 'Algorithm "none" is explicitly disallowed for decoding.' 

101 raise ValueError(msg) 

102 

103 

104def _validate_jwt_algorithms(algorithms: list[str]) -> None: 

105 if not isinstance(algorithms, list): 

106 raise TypeError("Algorithms must be a list of strings") 

107 

108 if len(algorithms) > MAX_JWT_ALGORITHMS: 

109 raise ValueError("Too many algorithms provided") 

110 

111 for alg in algorithms: 

112 _check_disallowed_algorithm(alg) 

113 

114 

115def _validate_jwt_audience_item(aud: object) -> None: 

116 if not isinstance(aud, str): 

117 raise TypeError("Audience items must be strings") 

118 if len(aud) > MAX_JWT_AUDIENCE_LENGTH: 

119 raise ValueError("Audience item string is too long") 

120 

121 

122def _validate_jwt_audience_collection( 

123 audience: list[str] | tuple[str, ...] | set[str], 

124) -> None: 

125 if len(audience) > MAX_JWT_AUDIENCES: 

126 raise ValueError("Too many audience items provided") 

127 

128 for aud in audience: 

129 _validate_jwt_audience_item(aud) 

130 

131 

132def _validate_jwt_audience(audience: str | Iterable[str]) -> None: 

133 if isinstance(audience, str): 

134 if len(audience) > MAX_JWT_AUDIENCE_LENGTH: 

135 raise ValueError("Audience string is too long") 

136 return 

137 

138 if isinstance(audience, (list, tuple, set)): 

139 _validate_jwt_audience_collection(audience) 

140 return 

141 

142 raise TypeError("Audience must be a string or iterable of strings") 

143 

144 

145@safe_from( 

146 PyJWTError, 

147 ValueError, 

148 TypeError, 

149 AttributeError, 

150 NotImplementedError, 

151 KeyError, 

152 Exception, 

153) 

154def decode_jwt( 

155 token: str, 

156 secret_key: str, 

157 algorithms: list[str], 

158 audience: str | Iterable[str], 

159) -> JWTPayload: 

160 """Decode a JWT securely with strict claim validation. 

161 

162 Enforces that 'exp' (expiration) and 'aud' (audience) claims are present 

163 and validated. Explicitly rejects the "none" algorithm. 

164 

165 Args: 

166 token: The encoded JWT string. 

167 secret_key: The secret key for verifying the signature. 

168 algorithms: List of exactly accepted algorithms. 

169 audience: The expected audience(s). 

170 

171 Returns: 

172 The decoded payload dictionary. 

173 

174 Raises: 

175 ValueError: If the "none" algorithm is present in the `algorithms` list. 

176 PyJWTError: If the token is invalid, expired, or has incorrect claims. 

177 

178 """ 

179 if not isinstance(token, str): 

180 raise TypeError("Token must be a string") 

181 

182 if not isinstance(secret_key, str): 

183 raise TypeError("Secret must be a string") 

184 

185 _validate_jwt_algorithms(algorithms) 

186 _validate_jwt_audience(audience) 

187 

188 return jwt.decode( # type: ignore[misc] 

189 token, 

190 secret_key, 

191 algorithms=algorithms, 

192 audience=audience, 

193 options={ 

194 "require": ["exp", "aud"], 

195 "verify_signature": True, 

196 "verify_exp": True, 

197 "verify_aud": True, 

198 }, 

199 )