sac4cps-backend/microservices/token-service/main.py

"""
Token Management Microservice
Handles JWT authentication, token generation, validation, and resource access control.
Port: 8001
"""

import asyncio
from datetime import datetime
from fastapi import FastAPI, HTTPException, Depends, Security
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
import logging
from typing import List, Optional

from models import (
    TokenGenerateRequest, TokenResponse, TokenValidationResponse, 
    TokenListResponse, HealthResponse
)
from database import connect_to_mongo, close_mongo_connection, get_database
from token_service import TokenService

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

security = HTTPBearer()

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager"""
    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")

app = FastAPI(
    title="Token Management Service",
    description="JWT authentication and token management microservice",
    version="1.0.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Dependency for database
async def get_db():
    return await get_database()

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Health check endpoint"""
    try:
        db = await get_database()
        await db.command("ping")
        
        return HealthResponse(
            service="token-service",
            status="healthy",
            timestamp=datetime.utcnow(),
            version="1.0.0"
        )
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        raise HTTPException(status_code=503, detail="Service Unavailable")

@app.get("/tokens", response_model=TokenListResponse)
async def get_tokens(db=Depends(get_db)):
    """Get all tokens"""
    try:
        token_service = TokenService(db)
        tokens = await token_service.get_tokens()
        
        return TokenListResponse(
            tokens=tokens,
            count=len(tokens)
        )
    except Exception as e:
        logger.error(f"Error getting tokens: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@app.post("/tokens/generate", response_model=TokenResponse)
async def generate_token(request: TokenGenerateRequest, db=Depends(get_db)):
    """Generate a new JWT token"""
    try:
        token_service = TokenService(db)
        token = token_service.generate_token(
            name=request.name,
            list_of_resources=request.list_of_resources,
            data_aggregation=request.data_aggregation,
            time_aggregation=request.time_aggregation,
            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)):
    """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
        
        return TokenValidationResponse(
            valid=is_valid,
            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
        )
    except Exception as e:
        logger.error(f"Error validating token: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@app.post("/tokens/save")
async def save_token(token: str, db=Depends(get_db)):
    """Save a token to database"""
    try:
        token_service = TokenService(db)
        result = await token_service.insert_token(token)
        return result
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.error(f"Error saving token: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@app.post("/tokens/revoke")
async def revoke_token(token: str, db=Depends(get_db)):
    """Revoke a token"""
    try:
        token_service = TokenService(db)
        result = await token_service.revoke_token(token)
        return result
    except ValueError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        logger.error(f"Error revoking token: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@app.get("/tokens/{token}/permissions")
async def get_token_permissions(token: str, db=Depends(get_db)):
    """Get permissions for a specific token"""
    try:
        token_service = TokenService(db)
        permissions = await token_service.get_token_permissions(token)
        
        if permissions:
            return {"permissions": permissions}
        else:
            raise HTTPException(status_code=401, detail="Invalid or expired token")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error getting token permissions: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

@app.delete("/tokens/cleanup")
async def cleanup_expired_tokens(db=Depends(get_db)):
    """Clean up expired tokens"""
    try:
        token_service = TokenService(db)
        expired_count = await token_service.cleanup_expired_tokens()
        
        return {
            "message": "Expired tokens cleaned up",
            "expired_tokens_removed": expired_count
        }
    except Exception as e:
        logger.error(f"Error cleaning up tokens: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)