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
+      - API_USER=admin@example.com
+      - API_PASSWORD=Admin!123
+      - API_AUTH_TYPE=jwt
       - 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 }}
+    API_AUTH_TYPE: {{ .Values.chatbot.config.apiAuthType | quote }}

deploy/helm/values.yaml

@@ -153,6 +153,9 @@ chatbot:
     secretKey: crapi
     defaultModel: gpt-4o-mini
     chromaPersistDirectory: /app/vectorstore
+    apiUser: admin@example.com
+    apiPassword: Admin!123
+    apiAuthType: jwt
   storage:
     # type: "manual"
     # pv:

services/chatbot/src/mcpserver/config.py

@@ -0,0 +1,11 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+class Config:
+    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")
+    TLS_ENABLED = os.getenv("TLS_ENABLED", "false").lower() in ("true", "1", "yes")
+    API_AUTH_TYPE = os.getenv("API_AUTH_TYPE", "jwt") # Toggle between "apikey" and "jwt" based on your auth type

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,58 +16,52 @@
 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}"
+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}"
+AUTH_TYPE_MAPPING = {
+    "apikey": "apiKey",
+    "jwt": "token"
+}
+CLIENT_AUTH = None
 
-API_USER = os.environ.get("API_USER", "admin@example.com")
-API_PASSWORD = os.environ.get("API_PASSWORD", "Admin!123")
-API_URL = f"{'https' if TLS_ENABLED else 'http'}://{WEB_SERVICE}"
-
-API_KEY = None
-API_AUTH_TYPE = "ApiKey"
-
-def get_api_key():
-    global API_KEY
-    # Try 5 times to get API key
+def get_client_auth():
+    global CLIENT_AUTH
+    # 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"
+        logger.info(f"Attempt {i+1} to get client auth...")
+        if CLIENT_AUTH is None:
+            login_body = {"email": Config.API_USER, "password": Config.API_PASSWORD}
+            auth_url = f"{BASE_IDENTITY_URL}/identity/management/user/{Config.API_AUTH_TYPE}"
             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}")
-                        raise Exception(f"Failed to get API key after {i+1} attempts: {response.status_code} {response.text}")
-                    logger.error(f"Failed to get API key in attempt {i+1}: {response.status_code} {response.text}. Sleeping for {i} seconds...")
+                        logger.error(f"Failed to get client auth after {i+1} attempts: {response.status_code} {response.text}")
+                        raise Exception(f"Failed to get client auth after {i+1} attempts: {response.status_code} {response.text}")
+                    logger.error(f"Failed to get client auth in attempt {i+1}: {response.status_code} {response.text}. Sleeping for {i} seconds...")
                     time.sleep(i)
                 response_json = response.json()
                 logger.info(f"Response: {response_json}")
-                API_KEY = response_json.get("apiKey")
-                logger.info(f"Chatbot API Key: {API_KEY}")
-                return API_KEY
-    return API_KEY
-
+                CLIENT_AUTH = f"{response_json.get('type')} {response_json.get(AUTH_TYPE_MAPPING[Config.API_AUTH_TYPE])}"
+                logger.info(f"MCP Server API Auth: {CLIENT_AUTH}")
+                return CLIENT_AUTH
+    return CLIENT_AUTH
 
 # 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(),
+        "Authorization": get_client_auth(),
     }
     return httpx.AsyncClient(
-        base_url=API_URL,
+        base_url=BASE_URL,
         headers=headers,
     )
 

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/constant/UserMessage.java

@@ -23,6 +23,9 @@ public class UserMessage {
       "Api Key generated successfully. Use it in authorization header with ApiKey prefix.";
   public static final String API_KEY_GENERATION_FAILED =
       "Api Key generation failed! Only permitted for admin users.";
+  public static final String JWT_TOKEN_GENERATED_MESSAGE =
+      "JWT Token generated successfully. Use it in authorization header with Bearer prefix.";
+  public static final String JWT_TOKEN_GENERATION_FAILED = "JWT Token generation failed!";
   public static final String ACCOUNT_LOCK_MESSAGE = "User account has been locked.";
   public static final String ACCOUNT_LOCKED_MESSAGE =
       "User account is locked. Retry login with MFA to unlock.";

services/identity/src/main/java/com/crapi/controller/ManagementControlller.java

@@ -60,4 +60,12 @@ public ResponseEntity<ApiKeyResponse> generateApiKey(
     ApiKeyResponse response = userService.generateApiKey(request, loginForm);
     return ResponseEntity.status(HttpStatus.OK).body(response);
   }
+
+  @PostMapping("/user/jwt")
+  public ResponseEntity<JwtResponse> generateJwt(
+      @RequestBody LoginForm loginForm, HttpServletRequest request)
+      throws UnsupportedEncodingException {
+    JwtResponse response = userService.generateJwtToken(request, loginForm);
+    return ResponseEntity.status(HttpStatus.OK).body(response);
+  }
 }

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

@@ -469,19 +469,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,14 +491,47 @@ 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);
+    log.debug("Api Key for user {}: {}", user.getEmail(), apiKey);
     user.setApiKey(apiKey);
     userRepository.save(user);
     return new ApiKeyResponse(user.getApiKey(), UserMessage.API_KEY_GENERATED_MESSAGE);
   }
 
+  /**
+   * @param request None
+   * @param loginForm LoginForm with user email and password
+   * @return JwtResponse with generated JWT token
+   */
+  @Override
+  @Transactional
+  public JwtResponse generateJwtToken(HttpServletRequest request, LoginForm loginForm) {
+    Authentication authentication =
+        authenticationManager.authenticate(
+            new UsernamePasswordAuthenticationToken(loginForm.getEmail(), loginForm.getPassword()));
+    if (authentication == null) {
+      return new JwtResponse(null, UserMessage.INVALID_CREDENTIALS);
+    }
+    log.info("Generate JWT token for user: {}", loginForm.getEmail());
+    User user;
+    if (request == null || jwtAuthTokenFilter.getToken(request) == null) {
+      user = userRepository.findByEmail(loginForm.getEmail());
+    } else {
+      user = getUserFromToken(request);
+    }
+    if (user == null) {
+      log.debug("User not found to generate JWT token");
+      return new JwtResponse(null, UserMessage.INVALID_CREDENTIALS);
+    }
+    String jwt = jwtProvider.generateJwtToken(user);
+    log.debug("JWT token for user {}: {}", user.getEmail(), jwt);
+    if (jwt != null) {
+      return new JwtResponse(jwt, UserMessage.JWT_TOKEN_GENERATED_MESSAGE);
+    } else {
+      return new JwtResponse(null, UserMessage.JWT_TOKEN_GENERATION_FAILED);
+    }
+  }
+
   /**
    * @param changePhoneForm contains old phone number and new phone number, api will send otp to
    *     change number to email address.

services/identity/src/main/java/com/crapi/service/UserService.java

@@ -57,4 +57,6 @@ ResponseEntity<JwtResponse> authenticateUserLogin(LoginForm loginForm)
   CRAPIResponse lockAccount(HttpServletRequest request, LockAccountForm lockAccountForm);
 
   ApiKeyResponse generateApiKey(HttpServletRequest request, LoginForm loginForm);
+
+  JwtResponse generateJwtToken(HttpServletRequest request, LoginForm loginForm);
 }