From da1fb2a058f6b891932300930f2202d185bc1bec Mon Sep 17 00:00:00 2001 From: rafaeldpsilva Date: Wed, 8 Oct 2025 14:05:11 +0100 Subject: [PATCH] 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(). --- microservices/token-service/main.py | 44 +++++++++------ microservices/token-service/models.py | 7 ++- microservices/token-service/token_service.py | 56 ++++++++++---------- 3 files changed, 59 insertions(+), 48 deletions(-) diff --git a/microservices/token-service/main.py b/microservices/token-service/main.py index 37f72f3..276967a 100644 --- a/microservices/token-service/main.py +++ b/microservices/token-service/main.py @@ -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) \ No newline at end of file + uvicorn.run(app, host="0.0.0.0", port=8001) diff --git a/microservices/token-service/models.py b/microservices/token-service/models.py index e9d75db..78d8a16 100644 --- a/microservices/token-service/models.py +++ b/microservices/token-service/models.py @@ -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() - } \ No newline at end of file + } diff --git a/microservices/token-service/token_service.py b/microservices/token-service/token_service.py index 3ea291b..182dab8 100644 --- a/microservices/token-service/token_service.py +++ b/microservices/token-service/token_service.py @@ -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 \ No newline at end of file + + return result.deleted_count