Coverage for agentos/tools/jwt.py: 24%

1""" 

2JWT — JSON Web Token encode/decode/verify (HS256/RS256/ES256). 

3 

4Supports: 

5 - HS256 (HMAC-SHA256), RS256 (RSA), ES256 (ECDSA) algorithms 

6 - Encode with claims (iss, sub, aud, exp, iat, nbf, jti, custom) 

7 - Decode with signature verification 

8 - Decode without verification (for inspection) 

9 - Token expiry checking 

10 - Claim validation 

11""" 

12 

13from __future__ import annotations 

14 

15import base64 

16import hashlib 

17import hmac 

18import json 

19import time 

20 

21# ============================================================================ 

22# JWTError 

23# ============================================================================ 

24 

25 

26class JWTError(Exception): 

27 pass 

28 

29 

30class ExpiredTokenError(JWTError): 

31 pass 

32 

33 

34class InvalidTokenError(JWTError): 

35 pass 

36 

37 

38# ============================================================================ 

39# Helpers 

40# ============================================================================ 

41 

42 

43def _b64url_encode(data: bytes) -> str: 

44 return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii") 

45 

46 

47def _b64url_decode(s: str) -> bytes: 

48 # Restore padding 

49 padding = 4 - len(s) % 4 

50 if padding != 4: 

51 s += "=" * padding 

52 return base64.urlsafe_b64decode(s) 

53 

54 

55def _json_b64_decode(s: str) -> dict: 

56 return json.loads(_b64url_decode(s)) 

57 

58 

59# ============================================================================ 

60# JWT 

61# ============================================================================ 

62 

63ALGORITHMS = frozenset( 

64 {"HS256", "HS384", "HS512", "RS256", "RS384", "RS512", "ES256", "ES384", "ES512"} 

65) 

66 

67 

68class JWT: 

69 """JSON Web Token encoder/decoder. 

70 

71 Usage: 

72 jwt = JWT(secret="my-secret") # for HS256 

73 

74 # Encode 

75 token = jwt.encode({"sub": "user123", "role": "admin"}, ttl=3600) 

76 

77 # Decode & verify 

78 payload = jwt.decode(token) 

79 

80 # Decode without verification (inspect only) 

81 payload = jwt.decode(token, verify=False) 

82 """ 

83 

84 def __init__( 

85 self, 

86 secret: str | None = None, 

87 private_key: str | None = None, 

88 public_key: str | None = None, 

89 algorithm: str = "HS256", 

90 ): 

91 if algorithm not in ALGORITHMS: 

92 raise ValueError(f"Unsupported algorithm: {algorithm}. Use one of {sorted(ALGORITHMS)}") 

93 

94 self._algorithm = algorithm 

95 self._hash_func = { 

96 "HS256": hashlib.sha256, 

97 "HS384": hashlib.sha384, 

98 "HS512": hashlib.sha512, 

99 } 

100 

101 if algorithm.startswith("HS"): 

102 if not secret: 

103 raise ValueError(f"{algorithm} requires a secret") 

104 self._secret = secret.encode("utf-8") 

105 elif algorithm.startswith("RS") or algorithm.startswith("ES"): 

106 if not private_key and not public_key: 

107 raise ValueError(f"{algorithm} requires at least one key") 

108 self._private_key = private_key 

109 self._public_key = public_key 

110 

111 # ---------- Encode ---------- 

112 

113 def encode( 

114 self, 

115 payload: dict, 

116 ttl: int | None = None, 

117 headers_extra: dict | None = None, 

118 ) -> str: 

119 """Encode a JWT token. 

120 

121 Args: 

122 payload: Claims to include 

123 ttl: Time-to-live in seconds (sets 'exp' claim) 

124 headers_extra: Additional header parameters 

125 """ 

126 header = {"alg": self._algorithm, "typ": "JWT"} 

127 if headers_extra: 

128 header.update(headers_extra) 

129 

130 claims = dict(payload) 

131 now = int(time.time()) 

132 

133 # Standard claims 

134 if "iat" not in claims: 

135 claims["iat"] = now 

136 if ttl is not None and "exp" not in claims: 

137 claims["exp"] = now + ttl 

138 

139 header_b64 = _b64url_encode(json.dumps(header).encode("utf-8")) 

140 payload_b64 = _b64url_encode(json.dumps(claims).encode("utf-8")) 

141 signing_input = f"{header_b64}.{payload_b64}" 

142 

143 signature = self._sign(signing_input) 

144 return f"{signing_input}.{signature}" 

145 

146 # ---------- Decode ---------- 

147 

148 def decode( 

149 self, 

150 token: str, 

151 verify: bool = True, 

152 audience: str | list[str] | None = None, 

153 issuer: str | None = None, 

154 ) -> dict: 

155 """Decode and optionally verify a JWT token. 

156 

157 Args: 

158 token: The JWT string 

159 verify: Whether to verify the signature (default True) 

160 audience: Expected audience (if present, validates 'aud' claim) 

161 issuer: Expected issuer (if present, validates 'iss' claim) 

162 """ 

163 parts = token.split(".") 

164 if len(parts) != 3: 

165 raise InvalidTokenError("JWT must have 3 parts (header.payload.signature)") 

166 

167 header_b64, payload_b64, signature_b64 = parts 

168 

169 # Decode header and payload (always safe) 

170 header = _json_b64_decode(header_b64) 

171 payload = _json_b64_decode(payload_b64) 

172 

173 # Verify algorithm 

174 alg = header.get("alg") 

175 if verify and alg != self._algorithm: 

176 raise InvalidTokenError(f"Algorithm mismatch: expected {self._algorithm}, got {alg}") 

177 

178 # Verify signature 

179 if verify: 

180 signing_input = f"{header_b64}.{payload_b64}" 

181 if not self._verify(signing_input, signature_b64): 

182 raise InvalidTokenError("Invalid signature") 

183 

184 # Check expiry 

185 exp = payload.get("exp") 

186 if exp and int(exp) < time.time(): 

187 raise ExpiredTokenError(f"Token expired at {exp}") 

188 

189 # Check not-before 

190 nbf = payload.get("nbf") 

191 if nbf and int(nbf) > time.time(): 

192 raise InvalidTokenError(f"Token not valid before {nbf}") 

193 

194 # Check audience 

195 if audience is not None: 

196 aud = payload.get("aud") 

197 if aud is None: 

198 raise InvalidTokenError("Token missing 'aud' claim") 

199 expected = [audience] if isinstance(audience, str) else audience 

200 if isinstance(aud, str): 

201 aud = [aud] 

202 if not set(expected) & set(aud): 

203 raise InvalidTokenError("Audience mismatch") 

204 

205 # Check issuer 

206 if issuer is not None: 

207 iss = payload.get("iss") 

208 if iss != issuer: 

209 raise InvalidTokenError(f"Issuer mismatch: expected {issuer}, got {iss}") 

210 

211 return payload 

212 

213 # ---------- Signature ---------- 

214 

215 def _sign(self, data: str) -> str: 

216 if self._algorithm.startswith("HS"): 

217 d = hmac.new( 

218 self._secret, data.encode("utf-8"), self._hash_func[self._algorithm] 

219 ).digest() 

220 return _b64url_encode(d) 

221 else: 

222 raise NotImplementedError( 

223 f"Signing with {self._algorithm} requires cryptographic libraries (cryptography)" 

224 ) 

225 

226 def _verify(self, data: str, signature_b64: str) -> bool: 

227 if self._algorithm.startswith("HS"): 

228 expected = self._sign(data) 

229 return hmac.compare_digest(expected, signature_b64) 

230 else: 

231 raise NotImplementedError( 

232 f"Verification with {self._algorithm} requires cryptographic libraries (cryptography)" 

233 ) 

234 

235 # ---------- Static helpers ---------- 

236 

237 @staticmethod 

238 def decode_unverified(token: str) -> dict: 

239 """Decode JWT without verifying signature (inspect only).""" 

240 parts = token.split(".") 

241 if len(parts) != 3: 

242 raise InvalidTokenError("JWT must have 3 parts") 

243 return _json_b64_decode(parts[1]) 

244 

245 @staticmethod 

246 def get_header(token: str) -> dict: 

247 """Extract JWT header without verification.""" 

248 parts = token.split(".") 

249 if len(parts) != 3: 

250 raise InvalidTokenError("JWT must have 3 parts") 

251 return _json_b64_decode(parts[0])