Coverage for mcpgateway/utils/create_jwt_token.py: 100%

54 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-22 16:17 +0100

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3"""jwt_cli.py - generate, inspect, **and be imported** for token helpers. 

4 

5Copyright 2025 

6SPDX-License-Identifier: Apache-2.0 

7Authors: Mihai Criveti 

8 

9* **Run as a script** - friendly CLI (works with *no* flags). 

10* **Import as a library** - drop-in async functions `create_jwt_token` & `get_jwt_token` 

11 kept for backward-compatibility, now delegating to the shared core helper. 

12 

13Quick usage 

14----------- 

15CLI (default secret, default payload): 

16 $ python3 jwt_cli.py 

17 

18Library: 

19```python 

20from mcpgateway.utils.create_jwt_token import create_jwt_token, get_jwt_token 

21 

22# inside async context 

23jwt = await create_jwt_token({"username": "alice"}) 

24``` 

25""" 

26 

27from __future__ import annotations 

28 

29import argparse 

30import asyncio 

31import datetime as _dt 

32import json 

33import sys 

34from typing import Any, Dict, List, Sequence 

35 

36import jwt # PyJWT 

37 

38from mcpgateway.config import settings 

39 

40__all__: Sequence[str] = ( 

41 "create_jwt_token", 

42 "get_jwt_token", 

43 "_create_jwt_token", 

44) 

45 

46# --------------------------------------------------------------------------- 

47# Defaults & constants 

48# --------------------------------------------------------------------------- 

49DEFAULT_SECRET: str = settings.jwt_secret_key 

50DEFAULT_ALGO: str = settings.jwt_algorithm 

51DEFAULT_EXP_MINUTES: int = settings.token_expiry # 7 days (in minutes) 

52DEFAULT_USERNAME: str = settings.basic_auth_user 

53 

54 

55# --------------------------------------------------------------------------- 

56# Core sync helper (used by both CLI & async wrappers) 

57# --------------------------------------------------------------------------- 

58 

59 

60def _create_jwt_token( 

61 data: Dict[str, Any], 

62 expires_in_minutes: int = DEFAULT_EXP_MINUTES, 

63 secret: str = DEFAULT_SECRET, 

64 algorithm: str = DEFAULT_ALGO, 

65) -> str: 

66 """Return a signed JWT string (synchronous, timezone-aware). 

67 

68 Args: 

69 data: Dictionary containing payload data to encode in the token. 

70 expires_in_minutes: Token expiration time in minutes. Default is 7 days. 

71 Set to 0 to disable expiration. 

72 secret: Secret key used for signing the token. 

73 algorithm: Signing algorithm to use. 

74 

75 Returns: 

76 The JWT token string. 

77 """ 

78 payload = data.copy() 

79 if expires_in_minutes > 0: 

80 expire = _dt.datetime.now(_dt.timezone.utc) + _dt.timedelta(minutes=expires_in_minutes) 

81 payload["exp"] = int(expire.timestamp()) 

82 return jwt.encode(payload, secret, algorithm=algorithm) 

83 

84 

85# --------------------------------------------------------------------------- 

86# **Async** wrappers for backward compatibility 

87# --------------------------------------------------------------------------- 

88 

89 

90async def create_jwt_token( 

91 data: Dict[str, Any], 

92 expires_in_minutes: int = DEFAULT_EXP_MINUTES, 

93 *, 

94 secret: str = DEFAULT_SECRET, 

95 algorithm: str = DEFAULT_ALGO, 

96) -> str: 

97 """Async facade for historic code. Internally synchronous—almost instant. 

98 

99 Args: 

100 data: Dictionary containing payload data to encode in the token. 

101 expires_in_minutes: Token expiration time in minutes. Default is 7 days. 

102 Set to 0 to disable expiration. 

103 secret: Secret key used for signing the token. 

104 algorithm: Signing algorithm to use. 

105 

106 Returns: 

107 The JWT token string. 

108 """ 

109 return _create_jwt_token(data, expires_in_minutes, secret, algorithm) 

110 

111 

112async def get_jwt_token() -> str: 

113 """Return a token for ``{"username": "admin"}``, mirroring old behaviour. 

114 

115 Returns: 

116 The JWT token string with default admin username. 

117 """ 

118 user_data = {"username": DEFAULT_USERNAME} 

119 return await create_jwt_token(user_data) 

120 

121 

122# --------------------------------------------------------------------------- 

123# **Decode** helper (non-verifying) - used by the CLI 

124# --------------------------------------------------------------------------- 

125 

126 

127def _decode_jwt_token(token: str, algorithms: List[str] | None = None) -> Dict[str, Any]: 

128 """Decode *without* signature verification—handy for inspection. 

129 

130 Args: 

131 token: JWT token string to decode. 

132 algorithms: List of allowed algorithms for decoding. Defaults to [DEFAULT_ALGO]. 

133 

134 Returns: 

135 Dictionary containing the decoded payload. 

136 """ 

137 return jwt.decode( 

138 token, 

139 settings.jwt_secret_key, 

140 algorithms=algorithms or [DEFAULT_ALGO], 

141 # options={"require": ["exp"]}, # Require expiration 

142 ) 

143 

144 

145# --------------------------------------------------------------------------- 

146# CLI Parsing & helpers 

147# --------------------------------------------------------------------------- 

148 

149 

150def _parse_args(): 

151 p = argparse.ArgumentParser( 

152 description="Generate or inspect JSON Web Tokens.", 

153 formatter_class=argparse.ArgumentDefaultsHelpFormatter, 

154 ) 

155 

156 group = p.add_mutually_exclusive_group() 

157 group.add_argument("-u", "--username", help="Add username=<value> to the payload.") 

158 group.add_argument("-d", "--data", help="Raw JSON payload or comma-separated key=value pairs.") 

159 group.add_argument("--decode", metavar="TOKEN", help="Token string to decode (no verification).") 

160 

161 p.add_argument( 

162 "-e", 

163 "--exp", 

164 type=int, 

165 default=DEFAULT_EXP_MINUTES, 

166 help="Expiration in minutes (0 disables the exp claim).", 

167 ) 

168 p.add_argument("-s", "--secret", default=DEFAULT_SECRET, help="Secret key for signing.") 

169 p.add_argument("--algo", default=DEFAULT_ALGO, help="Signing algorithm to use.") 

170 p.add_argument("--pretty", action="store_true", help="Pretty-print payload before encoding.") 

171 

172 return p.parse_args() 

173 

174 

175def _payload_from_cli(args) -> Dict[str, Any]: 

176 if args.username is not None: 

177 return {"username": args.username} 

178 

179 if args.data is not None: 

180 # Attempt JSON first 

181 try: 

182 return json.loads(args.data) 

183 except json.JSONDecodeError: 

184 pairs = [kv.strip() for kv in args.data.split(",") if kv.strip()] 

185 payload: Dict[str, Any] = {} 

186 for pair in pairs: 

187 if "=" not in pair: 

188 raise ValueError(f"Invalid key=value pair: '{pair}'") 

189 k, v = pair.split("=", 1) 

190 payload[k.strip()] = v.strip() 

191 return payload 

192 

193 # Fallback default payload 

194 return {"username": DEFAULT_USERNAME} 

195 

196 

197# --------------------------------------------------------------------------- 

198# Entry point for ``python jwt_cli.py`` 

199# --------------------------------------------------------------------------- 

200 

201 

202def main() -> None: # pragma: no cover 

203 args = _parse_args() 

204 

205 # Decode mode takes precedence 

206 if args.decode: 

207 decoded = _decode_jwt_token(args.decode, algorithms=[args.algo]) 

208 json.dump(decoded, sys.stdout, indent=2, default=str) 

209 sys.stdout.write("\n") 

210 return 

211 

212 payload = _payload_from_cli(args) 

213 

214 if args.pretty: 

215 print("Payload:") 

216 print(json.dumps(payload, indent=2, default=str)) 

217 print("-") 

218 

219 token = _create_jwt_token(payload, args.exp, args.secret, args.algo) 

220 print(token) 

221 

222 

223if __name__ == "__main__": 

224 # Support being run via ``python3 -m mcpgateway.utils.create_jwt_token`` too 

225 try: 

226 # Respect existing asyncio loop if present (e.g. inside uvicorn dev server) 

227 loop = asyncio.get_running_loop() 

228 loop.run_until_complete(asyncio.sleep(0)) # no-op to ensure loop alive 

229 except RuntimeError: 

230 # No loop; we're just a simple CLI call - run main synchronously 

231 main() 

232 else: 

233 # We're inside an active asyncio program - delegate to executor to avoid blocking 

234 loop.run_in_executor(None, main)