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

1#!/usr/bin/env python3

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

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

4

5 Copyright 2025

6 SPDX-License-Identifier: Apache-2.0

7 Authors: 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

13 Quick usage

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

15 CLI (default secret, default payload):

16 $ python3 jwt_cli.py

17

18 Library:

19```python

20 from mcpgateway.utils.create_jwt_token import create_jwt_token, get_jwt_token

21

22# inside async context

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

24```

25"""

26

27# Future

28 from __future__ import annotations

29

30# Standard

31 import argparse

32 import asyncio

33 import datetime as _dt

34 import json

35 import sys

36 from typing import Any, Dict, List, Sequence

37

38# Third-Party

39 import jwt # PyJWT

40

41# First-Party

42 from 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# ---------------------------------------------------------------------------

53 DEFAULT_SECRET: str = settings.jwt_secret_key

54 DEFAULT_ALGO: str = settings.jwt_algorithm

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

56 DEFAULT_USERNAME: str = settings.basic_auth_user

57

58

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

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

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

62

63

64 def _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

94 async 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

116 async 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

131 def _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

154 def _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

179 def _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

206 def 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

227 if __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)