A Flask-based JWT token service designed for integration with APISIX gateway. This service provides JWT token generation with flexible authentication methods and support for dynamic claims based on API keys.
- Multiple Authentication Methods: Support for LDAP (Active Directory) and file-based authentication
- Configurable JWT Tokens: JWT tokens are signed with a configurable key that can be shared with the APISIX gateway
- Flexible API Key Configuration:
- File-based: Each API key has its own configuration file
- NEW: Inline payload: Pass API key configuration directly in the request
- Dynamic Claims: Support for generating claims dynamically via:
- Function calls (Python functions)
- API calls (External HTTP services)
- Token Decoding: Simple mechanism to decode and verify JWT tokens
- Control Tower Integration: Seamlessly use manifest configurations from DSP AI Control Tower
- Python 3.8+
- Virtual environment (recommended)
-
Clone the repository:
git clone https://github.com/your-org/dsp_ai_jwt.git cd dsp_ai_jwt -
Create and activate a virtual environment:
python -m venv venv venv\Scripts\activate -
Install dependencies:
pip install -r requirements.txt -
Copy the example environment file and update it:
copy .env.example .env -
Update the
.envfile with your configuration settings.
User credentials are stored in config/users.yaml. Example format:
username:
password: hashed_password # SHA-256 hash
name: User Full Name
email: [email protected]
groups:
- group1
- group2
roles:
- role1
- role2LDAP authentication requires the following settings in your .env file:
LDAP_SERVER=ldap://your-ldap-server:389
LDAP_BASE_DN=dc=example,dc=com
LDAP_USER_DN=cn=users,dc=example,dc=com
LDAP_ADMIN_DN=cn=admin,dc=example,dc=com
LDAP_ADMIN_PASSWORD=admin_password
Each API key has its own configuration file in the config/api_keys/ directory. The filename should match the API key value (e.g., api_key_openai_1234567890.yaml).
Example API key configuration:
id: service-id
owner: Team Name
provider_permissions:
- provider1
- provider2
endpoint_permissions:
- /v1/endpoint1
- /v1/endpoint2
claims:
static:
# Static claims that will always be included
tier: premium
models:
- model1
- model2
rate_limit: 100
dynamic:
# Dynamic claims that will be resolved at runtime
quota:
type: function
module: claims.quota
function: get_remaining_quota
args:
user_id: "{user_id}" # Placeholder that will be replaced
usage_stats:
type: api
url: "http://usage-service/api/stats/{api_key_id}"
headers:
Authorization: "Bearer {internal_token}"
method: GET
response_field: dataDynamic claims can be generated in two ways:
-
Function-based:
claim_name: type: function module: package.module function: function_name args: arg1: value1 arg2: "{placeholder}" # Will be replaced with context values
-
API-based:
claim_name: type: api url: "http://api-endpoint/{placeholder}" method: GET|POST|PUT|DELETE headers: Header-Name: "value" response_field: path.to.field # Optional dot notation to extract specific data
Available placeholders:
{user_id}: The authenticated user's ID{api_key}: The original API key{api_key_id}: The ID of the API key{team_id}: The team ID of the user{internal_token}: Token for internal service communication
POST /token
Content-Type: application/json
{
"username": "user",
"password": "password",
"api_key": "api_key_value"
}
POST /token
Content-Type: application/json
{
"username": "user",
"password": "password",
"api_key_config": {
"id": "custom-key-id",
"owner": "Team Name",
"claims": {
"static": {
"key": "consumer-key",
"tier": "premium",
"models": ["gpt-4", "gpt-3.5-turbo"],
"rate_limit": 100,
"exp_hours": 2
}
}
}
}
Note: If both api_key and api_key_config are provided, api_key_config takes precedence.
Response:
{
"access_token": "eyJ0eXAi...",
"refresh_token": "eyJ0eXAi..."
}For detailed documentation on the inline configuration feature, see API_KEY_CONFIG_PAYLOAD.md.
POST /refresh
Authorization: Bearer {refresh_token}
Response:
{
"access_token": "eyJ0eXAi..."
}POST /decode
Content-Type: application/json
{
"token": "eyJ0eXAi..."
}
Response:
{
"sub": "username",
"tier": "premium",
"provider_permissions": ["openai", "groq"],
...
}pytest
Configure APISIX to use the JWT token for authentication by setting up the JWT plugin with the same secret key.
Example APISIX configuration:
{
"plugins": {
"jwt-auth": {
"key": "user-key",
"secret": "your-jwt-secret-key"
}
}
}- Add a new module in the
claims/directory - Implement functions that accept parameters and return dictionaries of claim values
- Reference the new function in API key configurations
Testing the Application To run the tests, you can use pytest with your virtual environment:
.\venv\Scripts\activate
pytest
pytest --cov=.
pytest tests/test_authentication.py
pytest -v The tests will use a test configuration defined in tests/conftest.py, which creates temporary test files for users and API keys.
Manually Testing the API You can also test the API endpoints manually using tools like curl or Postman:
Generate a token: curl -X POST http://localhost:5000/token -H "Content-Type: application/json" -d "{"username":"admin","password":"password","api_key":"api_key_openai_1234567890"}"
Decode a token: curl -X POST http://localhost:5000/decode -H "Content-Type: application/json" -d "{"token":"your-jwt-token-here"}"
Access a protected endpoint: curl -X GET http://localhost:5000/protected -H "Authorization: Bearer your-jwt-token-here"
curl -X POST http://192.168.1.25:5000/token -H "Content-Type: application/json" -d "{"username":"admin","password":"password","api_key":"ak_tiered_model_exec"}"
TOKEN=$(curl -s -X POST http://192.168.1.25:5000/token -H "Content-Type: application/json" -d "{"username":"admin","password":"password","api_key":"ak_tiered_model_exec"}" | jq -r '.access_token') && curl -X POST http://192.168.1.25:5000/decode -H "Content-Type: application/json" -d "{"token":"$TOKEN"}"
tk