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
« 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.
5Copyright 2025
6SPDX-License-Identifier: Apache-2.0
7Authors: Mihai Criveti
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.
13Quick usage
14-----------
15CLI (default secret, default payload):
16 $ python3 jwt_cli.py
18Library:
19```python
20from mcpgateway.utils.create_jwt_token import create_jwt_token, get_jwt_token
22# inside async context
23jwt = await create_jwt_token({"username": "alice"})
24```
25"""
27# Future
28from __future__ import annotations
30# Standard
31import argparse
32import asyncio
33import datetime as _dt
34import json
35import sys
36from typing import Any, Dict, List, Sequence
38# Third-Party
39import jwt # PyJWT
41# First-Party
42from mcpgateway.config import settings
44__all__: Sequence[str] = (
45 "create_jwt_token",
46 "get_jwt_token",
47 "_create_jwt_token",
48)
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
59# ---------------------------------------------------------------------------
60# Core sync helper (used by both CLI & async wrappers)
61# ---------------------------------------------------------------------------
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).
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.
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)
89# ---------------------------------------------------------------------------
90# **Async** wrappers for backward compatibility
91# ---------------------------------------------------------------------------
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.
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.
110 Returns:
111 The JWT token string.
112 """
113 return _create_jwt_token(data, expires_in_minutes, secret, algorithm)
116async def get_jwt_token() -> str:
117 """Return a token for ``{"username": "admin"}``, mirroring old behaviour.
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)
126# ---------------------------------------------------------------------------
127# **Decode** helper (non-verifying) - used by the CLI
128# ---------------------------------------------------------------------------
131def _decode_jwt_token(token: str, algorithms: List[str] | None = None) -> Dict[str, Any]:
132 """Decode *without* signature verification-handy for inspection.
134 Args:
135 token: JWT token string to decode.
136 algorithms: List of allowed algorithms for decoding. Defaults to [DEFAULT_ALGO].
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 )
149# ---------------------------------------------------------------------------
150# CLI Parsing & helpers
151# ---------------------------------------------------------------------------
154def _parse_args():
155 p = argparse.ArgumentParser(
156 description="Generate or inspect JSON Web Tokens.",
157 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
158 )
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).")
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.")
176 return p.parse_args()
179def _payload_from_cli(args) -> Dict[str, Any]:
180 if args.username is not None:
181 return {"username": args.username}
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
197 # Fallback default payload
198 return {"username": DEFAULT_USERNAME}
201# ---------------------------------------------------------------------------
202# Entry point for ``python jwt_cli.py``
203# ---------------------------------------------------------------------------
206def main() -> None: # pragma: no cover
207 args = _parse_args()
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
216 payload = _payload_from_cli(args)
218 if args.pretty:
219 print("Payload:")
220 print(json.dumps(payload, indent=2, default=str))
221 print("-")
223 token = _create_jwt_token(payload, args.exp, args.secret, args.algo)
224 print(token)
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)