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
« 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.
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"""
27from __future__ import annotations
29import argparse
30import asyncio
31import datetime as _dt
32import json
33import sys
34from typing import Any, Dict, List, Sequence
36import jwt # PyJWT
38from mcpgateway.config import settings
40__all__: Sequence[str] = (
41 "create_jwt_token",
42 "get_jwt_token",
43 "_create_jwt_token",
44)
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
55# ---------------------------------------------------------------------------
56# Core sync helper (used by both CLI & async wrappers)
57# ---------------------------------------------------------------------------
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).
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.
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)
85# ---------------------------------------------------------------------------
86# **Async** wrappers for backward compatibility
87# ---------------------------------------------------------------------------
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.
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.
106 Returns:
107 The JWT token string.
108 """
109 return _create_jwt_token(data, expires_in_minutes, secret, algorithm)
112async def get_jwt_token() -> str:
113 """Return a token for ``{"username": "admin"}``, mirroring old behaviour.
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)
122# ---------------------------------------------------------------------------
123# **Decode** helper (non-verifying) - used by the CLI
124# ---------------------------------------------------------------------------
127def _decode_jwt_token(token: str, algorithms: List[str] | None = None) -> Dict[str, Any]:
128 """Decode *without* signature verification—handy for inspection.
130 Args:
131 token: JWT token string to decode.
132 algorithms: List of allowed algorithms for decoding. Defaults to [DEFAULT_ALGO].
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 )
145# ---------------------------------------------------------------------------
146# CLI Parsing & helpers
147# ---------------------------------------------------------------------------
150def _parse_args():
151 p = argparse.ArgumentParser(
152 description="Generate or inspect JSON Web Tokens.",
153 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
154 )
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).")
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.")
172 return p.parse_args()
175def _payload_from_cli(args) -> Dict[str, Any]:
176 if args.username is not None:
177 return {"username": args.username}
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
193 # Fallback default payload
194 return {"username": DEFAULT_USERNAME}
197# ---------------------------------------------------------------------------
198# Entry point for ``python jwt_cli.py``
199# ---------------------------------------------------------------------------
202def main() -> None: # pragma: no cover
203 args = _parse_args()
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
212 payload = _payload_from_cli(args)
214 if args.pretty:
215 print("Payload:")
216 print(json.dumps(payload, indent=2, default=str))
217 print("-")
219 token = _create_jwt_token(payload, args.exp, args.secret, args.algo)
220 print(token)
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)