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

54 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-09 11:03 +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 

27# Future 

28from __future__ import annotations 

29 

30# Standard 

31import argparse 

32import asyncio 

33import datetime as _dt 

34import json 

35import sys 

36from typing import Any, Dict, List, Sequence 

37 

38# Third-Party 

39import jwt # PyJWT 

40 

41# First-Party 

42from mcpgateway.config import settings 

43 

44__all__: Sequence[str] = ( 

45 "create_jwt_token", 

46 "get_jwt_token", 

47 "_create_jwt_token", 

48) 

49 

50# --------------------------------------------------------------------------- 

51# Defaults & constants 

52# --------------------------------------------------------------------------- 

53DEFAULT_SECRET: str = settings.jwt_secret_key 

54DEFAULT_ALGO: str = settings.jwt_algorithm 

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

56DEFAULT_USERNAME: str = settings.basic_auth_user 

57 

58 

59# --------------------------------------------------------------------------- 

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

61# --------------------------------------------------------------------------- 

62 

63 

64def _create_jwt_token( 

65 data: Dict[str, Any], 

66 expires_in_minutes: int = DEFAULT_EXP_MINUTES, 

67 secret: str = DEFAULT_SECRET, 

68 algorithm: str = DEFAULT_ALGO, 

69) -> str: 

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

71 

72 Args: 

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

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

75 Set to 0 to disable expiration. 

76 secret: Secret key used for signing the token. 

77 algorithm: Signing algorithm to use. 

78 

79 Returns: 

80 The JWT token string. 

81 """ 

82 payload = data.copy() 

83 if expires_in_minutes > 0: 

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

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

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

87 

88 

89# --------------------------------------------------------------------------- 

90# **Async** wrappers for backward compatibility 

91# --------------------------------------------------------------------------- 

92 

93 

94async def create_jwt_token( 

95 data: Dict[str, Any], 

96 expires_in_minutes: int = DEFAULT_EXP_MINUTES, 

97 *, 

98 secret: str = DEFAULT_SECRET, 

99 algorithm: str = DEFAULT_ALGO, 

100) -> str: 

101 """Async facade for historic code. Internally synchronous-almost instant. 

102 

103 Args: 

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

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

106 Set to 0 to disable expiration. 

107 secret: Secret key used for signing the token. 

108 algorithm: Signing algorithm to use. 

109 

110 Returns: 

111 The JWT token string. 

112 """ 

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

114 

115 

116async def get_jwt_token() -> str: 

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

118 

119 Returns: 

120 The JWT token string with default admin username. 

121 """ 

122 user_data = {"username": DEFAULT_USERNAME} 

123 return await create_jwt_token(user_data) 

124 

125 

126# --------------------------------------------------------------------------- 

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

128# --------------------------------------------------------------------------- 

129 

130 

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

132 """Decode *without* signature verification-handy for inspection. 

133 

134 Args: 

135 token: JWT token string to decode. 

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

137 

138 Returns: 

139 Dictionary containing the decoded payload. 

140 """ 

141 return jwt.decode( 

142 token, 

143 settings.jwt_secret_key, 

144 algorithms=algorithms or [DEFAULT_ALGO], 

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

146 ) 

147 

148 

149# --------------------------------------------------------------------------- 

150# CLI Parsing & helpers 

151# --------------------------------------------------------------------------- 

152 

153 

154def _parse_args(): 

155 p = argparse.ArgumentParser( 

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

157 formatter_class=argparse.ArgumentDefaultsHelpFormatter, 

158 ) 

159 

160 group = p.add_mutually_exclusive_group() 

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

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

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

164 

165 p.add_argument( 

166 "-e", 

167 "--exp", 

168 type=int, 

169 default=DEFAULT_EXP_MINUTES, 

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

171 ) 

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

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

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

175 

176 return p.parse_args() 

177 

178 

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

180 if args.username is not None: 

181 return {"username": args.username} 

182 

183 if args.data is not None: 

184 # Attempt JSON first 

185 try: 

186 return json.loads(args.data) 

187 except json.JSONDecodeError: 

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

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

190 for pair in pairs: 

191 if "=" not in pair: 

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

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

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

195 return payload 

196 

197 # Fallback default payload 

198 return {"username": DEFAULT_USERNAME} 

199 

200 

201# --------------------------------------------------------------------------- 

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

203# --------------------------------------------------------------------------- 

204 

205 

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

207 args = _parse_args() 

208 

209 # Decode mode takes precedence 

210 if args.decode: 

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

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

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

214 return 

215 

216 payload = _payload_from_cli(args) 

217 

218 if args.pretty: 

219 print("Payload:") 

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

221 print("-") 

222 

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

224 print(token) 

225 

226 

227if __name__ == "__main__": 

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

229 try: 

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

231 loop = asyncio.get_running_loop() 

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

233 except RuntimeError: 

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

235 main() 

236 else: 

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

238 loop.run_in_executor(None, main)