Add Token model and update endpoints to use it
Refactor token validation and saving endpoints to accept a Token model instead of a raw string. Set default values for token generation request fields. Update TokenService cleanup to use datetime.now().
This commit is contained in:
@@ -14,8 +14,8 @@ import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from models import (
|
||||
TokenGenerateRequest, TokenResponse, TokenValidationResponse,
|
||||
TokenListResponse, HealthResponse
|
||||
TokenGenerateRequest, TokenResponse, TokenValidationResponse,
|
||||
TokenListResponse, HealthResponse, Token
|
||||
)
|
||||
from database import connect_to_mongo, close_mongo_connection, get_database
|
||||
from token_service import TokenService
|
||||
@@ -32,9 +32,9 @@ async def lifespan(app: FastAPI):
|
||||
logger.info("Token Service starting up...")
|
||||
await connect_to_mongo()
|
||||
logger.info("Token Service startup complete")
|
||||
|
||||
|
||||
yield
|
||||
|
||||
|
||||
logger.info("Token Service shutting down...")
|
||||
await close_mongo_connection()
|
||||
logger.info("Token Service shutdown complete")
|
||||
@@ -64,7 +64,7 @@ async def health_check():
|
||||
try:
|
||||
db = await get_database()
|
||||
await db.command("ping")
|
||||
|
||||
|
||||
return HealthResponse(
|
||||
service="token-service",
|
||||
status="healthy",
|
||||
@@ -81,7 +81,7 @@ async def get_tokens(db=Depends(get_db)):
|
||||
try:
|
||||
token_service = TokenService(db)
|
||||
tokens = await token_service.get_tokens()
|
||||
|
||||
|
||||
return TokenListResponse(
|
||||
tokens=tokens,
|
||||
count=len(tokens)
|
||||
@@ -95,6 +95,16 @@ async def generate_token(request: TokenGenerateRequest, db=Depends(get_db)):
|
||||
"""Generate a new JWT token"""
|
||||
try:
|
||||
token_service = TokenService(db)
|
||||
|
||||
if not request.data_aggregation:
|
||||
request.data_aggregation = False
|
||||
if not request.time_aggregation:
|
||||
request.time_aggregation = False
|
||||
if not request.embargo:
|
||||
request.embargo = False
|
||||
if not request.exp_hours:
|
||||
request.exp_hours = 24
|
||||
|
||||
token = token_service.generate_token(
|
||||
name=request.name,
|
||||
list_of_resources=request.list_of_resources,
|
||||
@@ -103,23 +113,23 @@ async def generate_token(request: TokenGenerateRequest, db=Depends(get_db)):
|
||||
embargo=request.embargo,
|
||||
exp_hours=request.exp_hours
|
||||
)
|
||||
|
||||
|
||||
return TokenResponse(token=token)
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating token: {e}")
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@app.post("/tokens/validate", response_model=TokenValidationResponse)
|
||||
async def validate_token(token: str, db=Depends(get_db)):
|
||||
async def validate_token(token: Token, db=Depends(get_db)):
|
||||
"""Validate and decode a JWT token"""
|
||||
try:
|
||||
token_service = TokenService(db)
|
||||
is_valid = await token_service.is_token_valid(token)
|
||||
decoded = token_service.decode_token(token) if is_valid else None
|
||||
|
||||
is_valid = await token_service.is_token_valid(token.token)
|
||||
decoded = token_service.decode_token(token.token) if is_valid else None
|
||||
|
||||
return TokenValidationResponse(
|
||||
valid=is_valid,
|
||||
token=token,
|
||||
token=token.token,
|
||||
decoded=decoded if is_valid and "error" not in (decoded or {}) else None,
|
||||
error=decoded.get("error") if decoded and "error" in decoded else None
|
||||
)
|
||||
@@ -128,11 +138,11 @@ async def validate_token(token: str, db=Depends(get_db)):
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
|
||||
@app.post("/tokens/save")
|
||||
async def save_token(token: str, db=Depends(get_db)):
|
||||
async def save_token(token: Token, db=Depends(get_db)):
|
||||
"""Save a token to database"""
|
||||
try:
|
||||
token_service = TokenService(db)
|
||||
result = await token_service.insert_token(token)
|
||||
result = await token_service.insert_token(token.token)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
@@ -159,7 +169,7 @@ async def get_token_permissions(token: str, db=Depends(get_db)):
|
||||
try:
|
||||
token_service = TokenService(db)
|
||||
permissions = await token_service.get_token_permissions(token)
|
||||
|
||||
|
||||
if permissions:
|
||||
return {"permissions": permissions}
|
||||
else:
|
||||
@@ -176,7 +186,7 @@ async def cleanup_expired_tokens(db=Depends(get_db)):
|
||||
try:
|
||||
token_service = TokenService(db)
|
||||
expired_count = await token_service.cleanup_expired_tokens()
|
||||
|
||||
|
||||
return {
|
||||
"message": "Expired tokens cleaned up",
|
||||
"expired_tokens_removed": expired_count
|
||||
@@ -187,4 +197,4 @@ async def cleanup_expired_tokens(db=Depends(get_db)):
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8001)
|
||||
uvicorn.run(app, host="0.0.0.0", port=8001)
|
||||
|
||||
@@ -6,6 +6,9 @@ from pydantic import BaseModel, Field
|
||||
from typing import List, Optional, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
class Token(BaseModel):
|
||||
token: str = Field(..., description="JWT token")
|
||||
|
||||
class TokenGenerateRequest(BaseModel):
|
||||
"""Request model for token generation"""
|
||||
name: str = Field(..., description="Token owner name")
|
||||
@@ -48,8 +51,8 @@ class HealthResponse(BaseModel):
|
||||
status: str
|
||||
timestamp: datetime
|
||||
version: str
|
||||
|
||||
|
||||
class Config:
|
||||
json_encoders = {
|
||||
datetime: lambda v: v.isoformat()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,20 +11,20 @@ import os
|
||||
|
||||
class TokenService:
|
||||
"""Service for managing JWT tokens and authentication"""
|
||||
|
||||
|
||||
def __init__(self, db: AsyncIOMotorDatabase, secret_key: str = None):
|
||||
self.db = db
|
||||
self.secret_key = secret_key or os.getenv("JWT_SECRET_KEY", "energy-dashboard-secret-key")
|
||||
self.tokens_collection = db.tokens
|
||||
|
||||
def generate_token(self, name: str, list_of_resources: List[str],
|
||||
|
||||
def generate_token(self, name: str, list_of_resources: List[str],
|
||||
data_aggregation: bool = False, time_aggregation: bool = False,
|
||||
embargo: int = 0, exp_hours: int = 24) -> str:
|
||||
"""Generate a new JWT token with specified permissions"""
|
||||
|
||||
|
||||
# Calculate expiration time
|
||||
exp_timestamp = int((datetime.utcnow() + timedelta(hours=exp_hours)).timestamp())
|
||||
|
||||
|
||||
# Create token payload
|
||||
payload = {
|
||||
"name": name,
|
||||
@@ -36,11 +36,11 @@ class TokenService:
|
||||
"iat": int(datetime.utcnow().timestamp()),
|
||||
"jti": str(uuid.uuid4()) # unique token ID
|
||||
}
|
||||
|
||||
|
||||
# Generate JWT token
|
||||
token = jwt.encode(payload, self.secret_key, algorithm="HS256")
|
||||
return token
|
||||
|
||||
|
||||
def decode_token(self, token: str) -> Optional[Dict[str, Any]]:
|
||||
"""Decode and verify JWT token"""
|
||||
try:
|
||||
@@ -50,11 +50,11 @@ class TokenService:
|
||||
return {"error": "Token has expired"}
|
||||
except jwt.InvalidTokenError:
|
||||
return {"error": "Invalid token"}
|
||||
|
||||
|
||||
async def insert_token(self, token: str) -> Dict[str, Any]:
|
||||
"""Save token to database"""
|
||||
now = datetime.utcnow()
|
||||
|
||||
|
||||
# Decode token to verify it's valid
|
||||
decoded = self.decode_token(token)
|
||||
if decoded and "error" not in decoded:
|
||||
@@ -68,14 +68,14 @@ class TokenService:
|
||||
"resources": decoded.get("list_of_resources", []),
|
||||
"expires_at": datetime.fromtimestamp(decoded.get("exp", 0))
|
||||
}
|
||||
|
||||
|
||||
# Upsert token (update if exists, insert if not)
|
||||
await self.tokens_collection.replace_one(
|
||||
{"token": token},
|
||||
token_record,
|
||||
upsert=True
|
||||
)
|
||||
|
||||
|
||||
return {
|
||||
"token": token,
|
||||
"datetime": now.isoformat(),
|
||||
@@ -83,11 +83,11 @@ class TokenService:
|
||||
}
|
||||
else:
|
||||
raise ValueError("Invalid token cannot be saved")
|
||||
|
||||
|
||||
async def revoke_token(self, token: str) -> Dict[str, Any]:
|
||||
"""Revoke a token by marking it as inactive"""
|
||||
now = datetime.utcnow()
|
||||
|
||||
|
||||
result = await self.tokens_collection.update_one(
|
||||
{"token": token},
|
||||
{
|
||||
@@ -98,7 +98,7 @@ class TokenService:
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
if result.matched_count > 0:
|
||||
return {
|
||||
"token": token,
|
||||
@@ -107,51 +107,49 @@ class TokenService:
|
||||
}
|
||||
else:
|
||||
raise ValueError("Token not found")
|
||||
|
||||
|
||||
async def get_tokens(self) -> List[Dict[str, Any]]:
|
||||
"""Get all tokens from database"""
|
||||
cursor = self.tokens_collection.find({})
|
||||
tokens = []
|
||||
|
||||
|
||||
async for token_record in cursor:
|
||||
# Convert ObjectId to string and datetime to ISO format
|
||||
token_record["_id"] = str(token_record["_id"])
|
||||
for field in ["datetime", "created_at", "updated_at", "expires_at", "revoked_at"]:
|
||||
if field in token_record and token_record[field]:
|
||||
token_record[field] = token_record[field].isoformat()
|
||||
|
||||
|
||||
tokens.append(token_record)
|
||||
|
||||
|
||||
return tokens
|
||||
|
||||
|
||||
async def is_token_valid(self, token: str) -> bool:
|
||||
"""Check if token is valid and active"""
|
||||
# Check if token exists and is active in database
|
||||
token_record = await self.tokens_collection.find_one({
|
||||
"token": token,
|
||||
"active": True
|
||||
})
|
||||
|
||||
|
||||
if not token_record:
|
||||
return False
|
||||
|
||||
# Verify JWT signature and expiration
|
||||
|
||||
decoded = self.decode_token(token)
|
||||
return decoded is not None and "error" not in decoded
|
||||
|
||||
|
||||
async def get_token_permissions(self, token: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get permissions for a valid token"""
|
||||
if await self.is_token_valid(token):
|
||||
return self.decode_token(token)
|
||||
return None
|
||||
|
||||
|
||||
async def cleanup_expired_tokens(self) -> int:
|
||||
"""Remove expired tokens from database"""
|
||||
now = datetime.utcnow()
|
||||
|
||||
now = datetime.now()
|
||||
|
||||
# Delete tokens that have expired
|
||||
result = await self.tokens_collection.delete_many({
|
||||
"expires_at": {"$lt": now}
|
||||
})
|
||||
|
||||
return result.deleted_count
|
||||
|
||||
return result.deleted_count
|
||||
|
||||
Reference in New Issue
Block a user