Connection and messages are now correct and fluently pass between client and server.

common/utils/key_encryption.py

@@ -6,6 +6,7 @@ import random
 import time
 from flask import Flask
 import os
+import ast
 
 
 def generate_api_key(prefix="EveAI-Chat"):
@@ -72,6 +73,8 @@ class JosKMSClient(kms_v1.KeyManagementServiceClient):
 
     def decrypt_api_key(self, encrypted_data):
         """Decrypts the API key using the specified key version."""
+        if isinstance(encrypted_data, str):
+            encrypted_data = ast.literal_eval(encrypted_data)
         key_version = encrypted_data['key_version']
         key_name = self.key_name
         encrypted_dek = b64decode(encrypted_data['encrypted_dek'].encode('utf-8'))

eveai_chat/socket_handlers/chat_handler.py

@@ -1,27 +1,38 @@
-from flask_jwt_extended import verify_jwt_in_request, get_jwt_identity, verify_jwt_in_request, decode_token
+from flask_jwt_extended import create_access_token, get_jwt_identity, verify_jwt_in_request, decode_token
 from flask_socketio import emit, disconnect
 from flask import current_app, request
-from common.extensions import socketio
+
+from common.extensions import socketio, kms_client
+from common.models.user import Tenant
 
 
 @socketio.on('connect')
 def handle_connect():
     try:
-        # Extract token from the auth object
+        current_app.logger.debug(f'SocketIO: Connection handling started using {request.args}')
-        token = request.args.get('token')
+        tenant_id = request.args.get('tenantId')
-        if not token:
+        if not tenant_id:
-            raise Exception("Missing Authorization Token")
+            raise Exception("Missing Tenant ID")
-        current_app.logger.debug(f'SocketIO: Received token: {token}')
+        api_key = request.args.get('apiKey')
-        # Verify token
+        if not api_key:
-        decoded_token = decode_token(token.split(" ")[1])  # Split to remove "Bearer " prefix
+            raise Exception("Missing API Key")
-        tenant_id = decoded_token["identity"]["tenant_id"]
+        current_app.logger.info(f'SocketIO: Connection handling found Tenant {tenant_id} with API Key {api_key}')
-        current_app.logger.info(f'SocketIO: Tenant {decoded_token["identity"]["tenant_id"]} connected')
+
-        # communicate connection to client
+        if not validate_api_key(tenant_id, api_key):
+            raise Exception("Invalid tenant_id - api_key combination")
+
+        # Create JWT token
+        token = create_access_token(identity={"tenant_id": tenant_id, "api_key": api_key})
+        current_app.logger.debug(f'SocketIO: Connection handling created token: {token} for tenant {tenant_id}')
+
+        # Communicate connection to client
         emit('connect', {'status': 'Connected', 'tenant_id': tenant_id})
+        emit('authenticated', {'token': token})  # Emit custom event with the token
+        current_app.logger.debug(f'SocketIO: Connection handling sent token to client for tenant {tenant_id}')
     except Exception as e:
         current_app.logger.error(f'SocketIO: Connection failed: {e}')
         # communicate connection problem to client
-        emit('connect', {'status': 'Connection Failed'})
+        emit('connect_error', {'status': 'Connection Failed'})
         disconnect()
 
 
@@ -33,20 +44,46 @@ def handle_disconnect():
 @socketio.on('user_message')
 def handle_message(data):
     try:
-        current_app.logger.debug(f"SocketIO: Received message from tenant {data['tenantId']}: {data['message']}")
+        current_app.logger.debug(f"SocketIO: Message handling received message from tenant {data['tenantId']}: "
-        verify_jwt_in_request()
+                                 f"{data['message']} with token {data['token']}")
-        current_tenant = get_jwt_identity()
+        token = data.get('token')
-        print(f'Tenant {current_tenant["tenant_id"]} sent a message: {data}')
+        if not token:
+            raise Exception("Missing token")
+
+        # decoded_token = decode_token(token.split(" ")[1]) # remove "Bearer "
+        decoded_token = decode_token(token)
+        if not decoded_token:
+            raise Exception("Invalid token")
+        current_app.logger.debug(f"SocketIO: Message handling decoded token: {decoded_token}")
+
+        token_sub = decoded_token.get('sub')
+
+        current_tenant_id = token_sub.get('tenant_id')
+        if not current_tenant_id:
+            raise Exception("Missing tenant_id")
+
+        current_api_key = token_sub.get('api_key')
+        if not current_api_key:
+            raise Exception("Missing api_key")
+
         # Store interaction in the database
+
         response = {
             'tenantId': data['tenantId'],
-            'message': 'This is a bot response. Actual implementation still required.',
+            'message': f'This is a bot response. Responding to message {data['message']} '
+                       f'from tenant {current_tenant_id}',
             'messageId': 'bot-message-id',
             'algorithm': 'alg1'
         }
-        current_app.logger.debug(f"SocketIO: Bot response: {response}")
+        current_app.logger.debug(f"SocketIO: Message handling sent bot response: {response}")
         emit('bot_response', response, broadcast=True)
     except Exception as e:
         current_app.logger.error(f'SocketIO: Message handling failed: {e}')
         disconnect()
 
+
+def validate_api_key(tenant_id, api_key):
+    tenant = Tenant.query.get_or_404(tenant_id)
+    decrypted_api_key = kms_client.decrypt_api_key(tenant.encrypted_chat_api_key)
+
+    return decrypted_api_key == api_key

public/chat.html

@@ -12,7 +12,7 @@
 <body>
   <div id="chat-container"></div>
   <script>
-    document.addEventListener('DOMContentLoaded', () => {
+    document.addEventListener('DOMContentLoaded', function() {
       const eveAI = new EveAI(
         '1',
         'EveAI-CHAT-8553-7987-2800-9115-6454',

static/js/eveai-chat-widget.js

@@ -1,4 +1,3 @@
-// static/js/eveai-chat-widget.js
 class EveAIChatWidget extends HTMLElement {
   static get observedAttributes() {
     return ['tenant-id', 'api-key', 'domain'];
@@ -8,6 +7,7 @@ class EveAIChatWidget extends HTMLElement {
     super();
     this.socket = null; // Initialize socket to null
     this.attributesSet = false; // Flag to check if all attributes are set
+    this.jwtToken = null; // Initialize jwtToken to null
     console.log('EveAIChatWidget constructor called');
   }
 
@@ -76,31 +76,36 @@ class EveAIChatWidget extends HTMLElement {
     }
     console.log(`Initializing socket connection to ${this.domain}`);
 
-    const token = 'Bearer ' + this.apiKey
+    // Ensure apiKey is passed in the query parameters
-
-    // Include tenantId in query parameters
     this.socket = io(this.domain, {
       path: '/chat/socket.io/',
       transports: ['websocket', 'polling'],
-      auth: {
-        token: token  // Add the token to the authentication object
-      },
       query: {
         tenantId: this.tenantId,
-        // apiKey: this.apiKey
+        apiKey: this.apiKey // Ensure apiKey is included here
       },
+      auth: {
+        token: 'Bearer ' + this.apiKey // Ensure token is included here
+      }
     });
 
-    this.socket.on('connect', () => {
+    this.socket.on('connect', (data) => {
       console.log('Socket connected');
     });
 
+    this.socket.on('authenticated', (data) => {
+      console.log('Authenticated event received: ', data);
+      if (data.token) {
+        this.jwtToken = data.token; // Store the JWT token received from the server
+      }
+    });
+
     this.socket.on('connect_error', (err) => {
       console.error('Socket connection error:', err);
     });
 
     this.socket.on('connect_timeout', () => {
-      console.error('Socket connection timeout')
+      console.error('Socket connection timeout');
     });
 
     this.socket.on('disconnect', () => {
@@ -158,8 +163,12 @@ class EveAIChatWidget extends HTMLElement {
       console.error('Socket is not initialized');
       return;
     }
+    if (!this.jwtToken) {
+      console.error('JWT token is not available');
+      return;
+    }
     console.log('Sending message to backend');
-    this.socket.emit('user_message', { tenantId: this.tenantId, apiKey: this.apiKey, message });
+    this.socket.emit('user_message', { tenantId: this.tenantId, token: this.jwtToken, message });
   }
 }
 

static/js/eveai-sdk.js

@@ -4,6 +4,7 @@ class EveAI {
     this.tenantId = tenantId;
     this.apiKey = apiKey;
     this.domain = domain;
+
     console.log('EveAI constructor:', { tenantId, apiKey, domain });
   }