JWT auth added for MCP server's api calls (#314)

* JWT auth added for MCP server's api calls

Author: keyurdoshi03

deploy/docker/docker-compose.yml

@@ -156,6 +156,8 @@ services:
     environment:
       - TLS_ENABLED=${TLS_ENABLED:-false}
       - SERVER_PORT=${CHATBOT_SERVER_PORT:-5002}
+      - WEB_SERVICE=crapi-web
+      - IDENTITY_SERVICE=crapi-identity:${IDENTITY_SERVER_PORT:-8080}
       - DB_NAME=crapi
       - DB_USER=admin
       - DB_PASSWORD=crapisecretpassword
@@ -166,6 +168,9 @@ services:
       - MONGO_DB_USER=admin
       - MONGO_DB_PASSWORD=crapisecretpassword
       - MONGO_DB_NAME=crapi
+      - [email protected]
+      - API_PASSWORD=Admin!123
+      - OPENAPI_SPEC=/app/resources/crapi-openapi-spec.json
       - DEFAULT_MODEL=gpt-4o-mini
       - CHROMA_PERSIST_DIRECTORY=/app/vectorstore
       # - CHATBOT_OPENAI_API_KEY=

deploy/helm/templates/chatbot/config.yaml

@@ -8,7 +8,7 @@ metadata:
 data:
     SERVER_PORT: {{ .Values.chatbot.port | quote }}
     IDENTITY_SERVICE: {{ .Values.identity.service.name }}:{{ .Values.identity.port }}
-    WEB_SERVICE: {{ .Values.web.service.name }}:{{ .Values.web.port }}
+    WEB_SERVICE: {{ .Values.web.service.name }}
     TLS_ENABLED: {{ .Values.tlsEnabled | quote }}
     DB_HOST: {{ .Values.postgresdb.service.name }}
     DB_USER: {{ .Values.postgresdb.config.postgresUser }}
@@ -23,3 +23,6 @@ data:
     CHATBOT_OPENAI_API_KEY: {{ .Values.openAIApiKey }}
     DEFAULT_MODEL: {{ .Values.chatbot.config.defaultModel | quote }}
     CHROMA_PERSIST_DIRECTORY: {{ .Values.chatbot.config.chromaPersistDirectory | quote }}
+    API_USER: {{ .Values.chatbot.config.apiUser | quote }}
+    API_PASSWORD: {{ .Values.chatbot.config.apiPassword | quote }}
+    OPENAPI_SPEC: {{ .Values.chatbot.config.openapiSpec | quote }}

deploy/helm/values.yaml

@@ -153,6 +153,9 @@ chatbot:
     secretKey: crapi
     defaultModel: gpt-4o-mini
     chromaPersistDirectory: /app/vectorstore
+    apiUser: [email protected]
+    apiPassword: Admin!123
+    openapiSpec: /app/resources/crapi-openapi-spec.json
   storage:
     # type: "manual"
     # pv:

services/chatbot/src/mcpserver/config.py

@@ -0,0 +1,13 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+class Config:
+    TLS_ENABLED = os.getenv("TLS_ENABLED", "false").lower() in ("true", "1", "yes")
+    WEB_SERVICE = os.getenv("WEB_SERVICE", "crapi-web")
+    IDENTITY_SERVICE = os.getenv("IDENTITY_SERVICE", "crapi-identity:8080")
+    CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY", "/app/vectorstore")
+    OPENAPI_SPEC = os.getenv("OPENAPI_SPEC", "/app/resources/crapi-openapi-spec.json")
+    API_USER = os.getenv("API_USER", "[email protected]")
+    API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123")

services/chatbot/src/mcpserver/server.py

@@ -2,6 +2,7 @@
 from fastmcp import FastMCP, settings
 import json
 import os
+from .config import Config
 import logging
 import time
 from .tool_helpers import (
@@ -15,36 +16,27 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
-WEB_SERVICE = os.environ.get("WEB_SERVICE", "crapi-web")
-IDENTITY_SERVICE = os.environ.get("IDENTITY_SERVICE", "crapi-identity:8080")
-TLS_ENABLED = os.environ.get("TLS_ENABLED", "false").lower() in ("true", "1", "yes")
-BASE_URL = f"{'https' if TLS_ENABLED else 'http'}://{WEB_SERVICE}"
-BASE_IDENTITY_URL = f"{'https' if TLS_ENABLED else 'http'}://{IDENTITY_SERVICE}"
-
-API_USER = os.environ.get("API_USER", "[email protected]")
-API_PASSWORD = os.environ.get("API_PASSWORD", "Admin!123")
-API_URL = f"{'https' if TLS_ENABLED else 'http'}://{WEB_SERVICE}"
-
+BASE_URL = f"{'https' if Config.TLS_ENABLED else 'http'}://{Config.WEB_SERVICE}"
+BASE_IDENTITY_URL = f"{'https' if Config.TLS_ENABLED else 'http'}://{Config.IDENTITY_SERVICE}"
 API_KEY = None
-API_AUTH_TYPE = "ApiKey"
 
 def get_api_key():
     global API_KEY
-    # Try 5 times to get API key
+    # Try 5 times to get client auth
     MAX_ATTEMPTS = 5
     for i in range(MAX_ATTEMPTS):
         logger.info(f"Attempt {i+1} to get API key...")
         if API_KEY is None:
-            login_body = {"email": API_USER, "password": API_PASSWORD}
-            apikey_url = f"{BASE_IDENTITY_URL}/identity/management/user/apikey"
+            login_body = {"email": Config.API_USER, "password": Config.API_PASSWORD}
+            auth_url = f"{BASE_IDENTITY_URL}/identity/management/user/apikey"
             headers = {
                 "Content-Type": "application/json",
             }
             with httpx.Client(
-                base_url=API_URL,
+                base_url=BASE_URL,
                 headers=headers,
             ) as client:
-                response = client.post(apikey_url, json=login_body)
+                response = client.post(auth_url, json=login_body)
                 if response.status_code != 200:
                     if i == MAX_ATTEMPTS - 1:
                         logger.error(f"Failed to get API key after {i+1} attempts: {response.status_code} {response.text}")
@@ -54,24 +46,23 @@ def get_api_key():
                 response_json = response.json()
                 logger.info(f"Response: {response_json}")
                 API_KEY = response_json.get("apiKey")
-                logger.info(f"Chatbot API Key: {API_KEY}")
+                logger.info(f"MCP Server API Key: {API_KEY}")
                 return API_KEY
     return API_KEY
 
-
 # Async HTTP client for API calls
 def get_http_client():
     """Create and configure the HTTP client with appropriate authentication."""
     headers = {
         "Authorization": "ApiKey " + get_api_key(),
     }
     return httpx.AsyncClient(
-        base_url=API_URL,
+        base_url=BASE_URL,
         headers=headers,
     )
 
 # Load your OpenAPI spec 
-with open("/app/resources/crapi-openapi-spec.json", "r") as f:
+with open(Config.OPENAPI_SPEC, "r") as f:
     openapi_spec = json.load(f)
 
 # Create the MCP server

services/chatbot/src/mcpserver/tool_helpers.py

@@ -3,12 +3,10 @@
 from langchain_community.vectorstores import Chroma
 from langchain.prompts import PromptTemplate
 from chatbot.extensions import db
-from chatbot.config import Config
+from .config import Config
 from langchain.chains import RetrievalQA
 from langchain_openai import ChatOpenAI
 
-retrieval_index_path = "/app/resources/chat_index"
-
 async def get_any_api_key():
     if os.environ.get("CHATBOT_OPENAI_API_KEY"):
         return os.environ.get("CHATBOT_OPENAI_API_KEY")

services/identity/src/main/java/com/crapi/config/JwtAuthTokenFilter.java

@@ -121,7 +121,7 @@ public String getUserFromToken(HttpServletRequest request) throws ParseException
     if (token != null) {
       if (apiType == ApiType.APIKEY) {
         log.debug("Token is api token");
-        username = tokenProvider.getUserNameFromApiToken(token);
+        username = tokenProvider.getUserNameFromJwtToken(token);
       } else {
         log.debug("Token is jwt token");
         if (tokenProvider.validateJwtToken(token)) {

services/identity/src/main/java/com/crapi/config/JwtProvider.java

@@ -103,27 +103,27 @@ public String generateJwtToken(User user) {
   }
 
   /**
-   * @param token
-   * @return username from JWT Token
+   * @param user
+   * @return generated apikey token without expiry date
    */
-  public String getUserNameFromJwtToken(String token) throws ParseException {
-    // Parse without verifying token signature
-    return JWTParser.parse(token).getJWTClaimsSet().getSubject();
+  public String generateApiKey(User user) {
+    JwtBuilder builder =
+        Jwts.builder()
+            .subject(user.getEmail())
+            .issuedAt(new Date())
+            .claim("role", user.getRole().getName())
+            .signWith(this.keyPair.getPrivate());
+    String jwt = builder.compact();
+    return jwt;
   }
 
   /**
    * @param token
    * @return username from JWT Token
    */
-  public String getUserNameFromApiToken(String token) throws ParseException {
+  public String getUserNameFromJwtToken(String token) throws ParseException {
     // Parse without verifying token signature
-    if (token != null) {
-      User user = userRepository.findByApiKey(token);
-      if (user != null) {
-        return user.getEmail();
-      }
-    }
-    return null;
+    return JWTParser.parse(token).getJWTClaimsSet().getSubject();
   }
 
   // Load RSA Public Key for JKU header if present

services/identity/src/main/java/com/crapi/service/Impl/UserServiceImpl.java

@@ -25,7 +25,6 @@
 import com.crapi.model.*;
 import com.crapi.repository.*;
 import com.crapi.service.UserService;
-import com.crapi.utils.ApiKeyGenerator;
 import com.crapi.utils.EmailTokenGenerator;
 import com.crapi.utils.MailBody;
 import com.crapi.utils.OTPGenerator;
@@ -469,19 +468,17 @@ public JwtResponse unlockAccount(
   @Override
   @Transactional
   public ApiKeyResponse generateApiKey(HttpServletRequest request, LoginForm loginForm) {
-    // if user is unauthenticated, use loginForm else user token to authenticate
+    Authentication authentication =
+        authenticationManager.authenticate(
+            new UsernamePasswordAuthenticationToken(loginForm.getEmail(), loginForm.getPassword()));
+    if (authentication == null) {
+      return new ApiKeyResponse(null, UserMessage.INVALID_CREDENTIALS);
+    }
+    log.info("Generate Api Key for user: {}", loginForm.getEmail());
     User user;
     if (request == null || jwtAuthTokenFilter.getToken(request) == null) {
       user = userRepository.findByEmail(loginForm.getEmail());
     } else {
-      log.info("Generate Api Key for user: {}", loginForm.getEmail());
-      Authentication authentication =
-          authenticationManager.authenticate(
-              new UsernamePasswordAuthenticationToken(
-                  loginForm.getEmail(), loginForm.getPassword()));
-      if (authentication == null) {
-        return new ApiKeyResponse(null, UserMessage.INVALID_CREDENTIALS);
-      }
       user = getUserFromToken(request);
     }
     if (user == null) {
@@ -493,11 +490,13 @@ public ApiKeyResponse generateApiKey(HttpServletRequest request, LoginForm login
       log.debug("Api Key already generated for user: {}", user.getEmail());
       return new ApiKeyResponse(user.getApiKey());
     }
-    log.info("Generate Api Key for user in token: {}", user.getEmail());
-    String apiKey = ApiKeyGenerator.generateRandom(512);
-    log.debug("Api Key for user in token {}: {}", user.getEmail(), apiKey);
+    String apiKey = jwtProvider.generateApiKey(user);
+    log.debug("Api Key for user {}: {}", user.getEmail(), apiKey);
+    if (apiKey == null) {
+      return new ApiKeyResponse(null, UserMessage.API_KEY_GENERATION_FAILED);
+    }
     user.setApiKey(apiKey);
-    userRepository.save(user);
+    userRepository.saveAndFlush(user);
     return new ApiKeyResponse(user.getApiKey(), UserMessage.API_KEY_GENERATED_MESSAGE);
   }