Implement chat API key generation, and create a tenant_overview

This commit is contained in:
Josako
2024-05-16 23:22:26 +02:00
parent 8c6d9bf5ca
commit 883988dbab
12 changed files with 387 additions and 52 deletions

View File

@@ -7,6 +7,7 @@ from flask_login import LoginManager
from flask_cors import CORS
from flask_socketio import SocketIO
from .utils.key_encryption import JosKMSClient
# Create extensions
db = SQLAlchemy()
@@ -17,3 +18,4 @@ mail = Mail()
login_manager = LoginManager()
cors = CORS()
socketio = SocketIO()
kms_client = JosKMSClient()

View File

@@ -37,7 +37,7 @@ class Tenant(db.Model):
license_start_date = db.Column(db.Date, nullable=True)
license_end_date = db.Column(db.Date, nullable=True)
allowed_monthly_interactions = db.Column(db.Integer, nullable=True)
encrypted_api_key = db.Column(db.String(500), nullable=True)
encrypted_chat_api_key = db.Column(db.String(500), nullable=True)
# Relations
users = db.relationship('User', backref='tenant')

View File

@@ -2,50 +2,68 @@ from google.cloud import kms
from base64 import b64encode, b64decode
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from flask import current_app
client = kms.KeyManagementServiceClient()
key_name = client.crypto_key_path('your-project-id', 'your-key-ring', 'your-crypto-key')
import random
from flask import Flask
def encrypt_api_key(api_key):
"""Encrypts the API key using the latest version of the KEK."""
dek = get_random_bytes(32) # AES 256-bit key
cipher = AES.new(dek, AES.MODE_GCM)
ciphertext, tag = cipher.encrypt_and_digest(api_key.encode())
# Encrypt the DEK using the latest version of the Google Cloud KMS key
encrypt_response = client.encrypt(
request={'name': key_name, 'plaintext': dek}
)
encrypted_dek = encrypt_response.ciphertext
# Store the version of the key used
key_version = encrypt_response.name
return {
'key_version': key_version,
'encrypted_dek': b64encode(encrypted_dek).decode('utf-8'),
'nonce': b64encode(cipher.nonce).decode('utf-8'),
'tag': b64encode(tag).decode('utf-8'),
'ciphertext': b64encode(ciphertext).decode('utf-8')
}
def generate_api_key(prefix="EveAI-Chat"):
parts = [str(random.randint(1000, 9999)) for _ in range(5)]
return f"{prefix}-{'-'.join(parts)}"
def decrypt_api_key(encrypted_data):
"""Decrypts the API key using the specified key version."""
key_version = encrypted_data['key_version']
encrypted_dek = b64decode(encrypted_data['encrypted_dek'])
nonce = b64decode(encrypted_data['nonce'])
tag = b64decode(encrypted_data['tag'])
ciphertext = b64decode(encrypted_data['ciphertext'])
class JosKMSClient(kms.KeyManagementServiceClient):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.key_name = None
self.crypto_key = None
self.key_ring = None
self.location = None
self.project_id = None
# Decrypt the DEK using the specified version of the Google Cloud KMS key
decrypt_response = client.decrypt(
request={'name': key_version, 'ciphertext': encrypted_dek}
)
dek = decrypt_response.plaintext
def init_app(self, app: Flask):
self.project_id = app.config.get('GC_PROJECT_NAME')
self.location = app.config.get('GC_LOCATION')
self.key_ring = app.config.get('GC_KEY_RING')
self.crypto_key = app.config.get('GC_CRYPTO_KEY')
self.key_name = self.crypto_key_path(self.project_id, self.location, self.key_ring, self.crypto_key)
cipher = AES.new(dek, AES.MODE_GCM, nonce=nonce)
api_key = cipher.decrypt_and_verify(ciphertext, tag)
return api_key.decode()
def encrypt_api_key(self, api_key):
"""Encrypts the API key using the latest version of the KEK."""
dek = get_random_bytes(32) # AES 256-bit key
cipher = AES.new(dek, AES.MODE_GCM)
ciphertext, tag = cipher.encrypt_and_digest(api_key.encode())
# Encrypt the DEK using the latest version of the Google Cloud KMS key
encrypt_response = self.encrypt(
request={'name': self.key_name, 'plaintext': dek}
)
encrypted_dek = encrypt_response.ciphertext
# Store the version of the key used
key_version = encrypt_response.name
return {
'key_version': key_version,
'encrypted_dek': b64encode(encrypted_dek).decode('utf-8'),
'nonce': b64encode(cipher.nonce).decode('utf-8'),
'tag': b64encode(tag).decode('utf-8'),
'ciphertext': b64encode(ciphertext).decode('utf-8')
}
def decrypt_api_key(self, encrypted_data):
"""Decrypts the API key using the specified key version."""
key_version = encrypted_data['key_version']
encrypted_dek = b64decode(encrypted_data['encrypted_dek'])
nonce = b64decode(encrypted_data['nonce'])
tag = b64decode(encrypted_data['tag'])
ciphertext = b64decode(encrypted_data['ciphertext'])
# Decrypt the DEK using the specified version of the Google Cloud KMS key
decrypt_response = self.decrypt(
request={'name': key_version, 'ciphertext': encrypted_dek}
)
dek = decrypt_response.plaintext
cipher = AES.new(dek, AES.MODE_GCM, nonce=nonce)
api_key = cipher.decrypt_and_verify(ciphertext, tag)
return api_key.decode()

View File

@@ -20,7 +20,7 @@ class Config(object):
SECURITY_CONFIRMABLE = True
SECURITY_TRACKABLE = True
SECURITY_PASSWORD_COMPLEXITY_CHECKER = 'zxcvbn'
SECURITY_POST_LOGIN_VIEW = '/user/tenant'
SECURITY_POST_LOGIN_VIEW = '/user/tenant_overview'
SECURITY_RECOVERABLE = True
SECURITY_EMAIL_SENDER = "eveai_super@flow-it.net"
PERMANENT_SESSION_LIFETIME = timedelta(minutes=60)
@@ -102,7 +102,8 @@ class DevConfig(Config):
SOCKETIO_ENGINEIO_LOGGER = True
# Google Cloud settings
GC_PROJECT_NAME = 'EveAI'
GC_PROJECT_NAME = 'eveai-420711'
GC_LOCATION = 'europe-west1'
GC_KEY_RING = 'eveai-chat'
GC_CRYPTO_KEY = 'envelope-encryption-key'

View File

@@ -6,7 +6,7 @@ from flask_security.signals import user_authenticated
from werkzeug.middleware.proxy_fix import ProxyFix
import logging.config
from common.extensions import db, migrate, bootstrap, security, mail, login_manager, cors
from common.extensions import db, migrate, bootstrap, security, mail, login_manager, cors, kms_client
from common.models.user import User, Role, Tenant, TenantDomain
from config.logging_config import LOGGING
from common.utils.security import set_tenant_session_data
@@ -79,6 +79,7 @@ def register_extensions(app):
mail.init_app(app)
login_manager.init_app(app)
cors.init_app(app)
kms_client.init_app(app)
# Register Blueprints

View File

@@ -51,5 +51,6 @@
<hr>
{% include 'footer.html' %}
{% include 'scripts.html' %}
{% block scripts %}{% endblock %}
</body>
</html>

View File

@@ -23,6 +23,30 @@
{% endif %}
{% endmacro %}
{% macro render_included_field(field, disabled_fields=[], include_fields=[]) %}
{% set disabled = field.name in disabled_fields %}
{% if field.name in include_fields %}
{% if field.type == 'BooleanField' %}
<div class="form-check">
{{ field(class="form-check-input", type="checkbox", id="flexSwitchCheckDefault") }}
{{ field.label(class="form-check-label", for="flexSwitchCheckDefault", disabled=disabled) }}
</div>
{% else %}
<div class="form-group">
{{ field.label(class="form-label") }}
{{ field(class="form-control", disabled=disabled) }}
{% if field.errors %}
<div class="invalid-feedback">
{% for error in field.errors %}
{{ error }}
{% endfor %}
</div>
{% endif %}
</div>
{% endif %}
{% endif %}
{% endmacro %}
{% macro render_table(headers, rows) %}
<div class="card">
<div class="table-responsive">

View File

@@ -71,7 +71,8 @@
{{ dropdown('User Mgmt', 'contacts', [
{'name': 'Tenant List', 'url': '/user/select_tenant', 'roles': ['Super User']},
{'name': 'Tenant Registration', 'url': '/user/tenant', 'roles': ['Super User']},
{'name': 'Generate API Key', 'url': '/user/generate_api_key', 'roles': ['Super User']},
{'name': 'Generate Chat API Key', 'url': '/user/generate_chat_api_key', 'roles': ['Super User']},
{'name': 'Tenant Overview', 'url': '/user/tenant_overview', 'roles': ['Super User', 'Tenant Admin']},
{'name': 'Tenant Domains', 'url': '/user/view_tenant_domains/' + session['tenant']['id']|string, 'roles': ['Super User', 'Tenant Admin']},
{'name': 'Tenant Domain Registration', 'url': '/user/tenant_domain', 'roles': ['Super User', 'Tenant Admin']},
{'name': 'User List', 'url': '/user/view_users/' + session['tenant']['id']|string, 'roles': ['Super User', 'Tenant Admin']},

View File

@@ -1,8 +1,14 @@
<!-- Optional JavaScript -->
<!-- jQuery first, then Popper.js, then Bootstrap JS -->
<script src="https://code.jquery.com/jquery-3.3.1.slim.min.js" integrity="sha384-q8i/X+965DzO0rT7abK41JStQIAqVgRVzpbzo5smXKp4YfRvH+8abtTE1Pi6jizo" crossorigin="anonymous"></script>
<script src="https://ajax.googleapis.com/ajax/libs/jquery/3.7.1/jquery.min.js"></script>
<script src="https://cdn.datatables.net/1.10.21/js/jquery.dataTables.min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.14.7/umd/popper.min.js" integrity="sha384-UO2eT0CpHqdSJQ6hJty5KVphtPhzWj9WO1clHTMGa3JDZwrnQq4sF86dIHNDz0W1" crossorigin="anonymous"></script>
<script src="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/js/bootstrap.min.js" integrity="sha384-JjSmVgyd0p3pXB1rRibZUAYoIIy6OrQ6VrjIEaFf/nJGzIxFDsf4x0xIM+B07jRM" crossorigin="anonymous"></script>
{% block scripts %}
{%- endblock scripts %}
<script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.14.7/perfect-scrollbar.min.js"></script>
<script src="{{url_for('static', filename='/assets/js/plugins/typedjs.js')}}"></script>
<script src="{{url_for('static', filename='/assets/js/plugins/prism.min.js')}}"></script>
<script src="{{url_for('static', filename='/assets/js/plugins/highlight.min.js')}}"></script>
<script src="{{url_for('static', filename='/assets/js/plugins/parallax.min.js')}}"></script>
<script src="{{url_for('static', filename='/assets/js/plugins/nouislider.min.js')}}"></script>
<script src="{{url_for('static', filename='/assets/js/plugins/anime.min.js')}}"></script>
<script src="{{url_for('static', filename='assets/js/material-kit-pro.min.js')}}?v=3.0.4 type="text/javascript"></script>

View File

@@ -0,0 +1,70 @@
{% extends 'base.html' %}
{% from "macros.html" import render_field %}
{% block title %}Generate Chat API Key{% endblock %}
{% block content_title %}Generate Chat API Key{% endblock %}
{% block content_description %}Generate an API key to enable chat.{% endblock %}
{% block content %}
<!-- Bootstrap Modal HTML for API Key Registration -->
<div class="modal fade" id="confirmModal" tabindex="-1" role="dialog" aria-labelledby="confirmModalLabel" aria-hidden="true">
<div class="modal-dialog" role="document">
<div class="modal-content">
<div class="modal-header">
<h5 class="modal-title" id="confirmModalLabel">Confirm New API Key</h5>
<button type="button" class="close" data-dismiss="modal" aria-label="Close">
<span aria-hidden="true">&times;</span>
</button>
</div>
<div class="modal-body">
<div id="confirmation-message">
There is already an API key defined for this tenant. Are you sure you want to generate a new API key? This will require updating the key in all locations using the chat interface.
</div>
<div id="api-key-message" style="display:none;">
<p>New API Key:</p>
<p id="new-api-key" style="font-weight: bold;"></p>
</div>
</div>
<div class="modal-footer">
<button type="button" class="btn btn-secondary" data-dismiss="modal">Cancel</button>
<button type="button" class="btn btn-primary" id="confirmNewKeyBtn">Confirm</button>
</div>
</div>
</div>
</div>
{% endblock %}
{% block content_footer %}
{% endblock %}
{% block scripts %}
<script>
// Trigger the modal when the user tries to generate a new API key
function promptForNewApiKey() {
$('#confirmModal').modal('show');
}
// Handle the confirmation button click
document.getElementById('confirmNewKeyBtn').addEventListener('click', function () {
generateNewApiKey();
});
function generateNewApiKey() {
// Make an AJAX request to your back-end to generate the new API key
$.ajax({
url: '/generate-api-key',
type: 'POST',
data: JSON.stringify({ tenant_id: tenantId }), // Pass the tenant ID as needed
contentType: 'application/json',
success: function(response) {
// Display the new API key in the modal
$('#confirmation-message').hide();
$('#api-key-message').show();
$('#new-api-key').text(response.new_api_key);
},
error: function(error) {
alert('Error generating new API key: ' + error.responseText);
}
});
}
</script>
{% endblock %}

View File

@@ -0,0 +1,171 @@
{% extends 'base.html' %}
{% from "macros.html" import render_field, render_included_field %}
{% block title %}Tenant Overview{% endblock %}
{% block content_title %}Tenant Overview{% endblock %}
{% block content_description %}Tenant information{% endblock %}
{% block content %}
<form method="post">
{{ form.hidden_tag() }}
<!-- Main Tenant Information -->
{% set main_fields = ['name', 'website', 'default_language', 'allowed_languages'] %}
{% for field in form %}
{{ render_included_field(field, disabled_fields=main_fields, include_fields=main_fields) }}
{% endfor %}
<!-- Nav Tabs -->
<div class="row">
<div class="col-lg-12">
<div class="nav-wrapper position-relative end-0">
<ul class="nav nav-pills nav-fill p-1" role="tablist">
<li class="nav-item" role="presentation">
<a class="nav-link mb-0 px-0 py-1 active" data-toggle="tab" href="#model-info-tab" role="tab" aria-controls="model-info" aria-selected="true">
Model Information
</a>
</li>
<li class="nav-item">
<a class="nav-link mb-0 px-0 py-1" data-toggle="tab" href="#license-info-tab" role="tab" aria-controls="license-info" aria-selected="false">
License Information
</a>
</li>
<li class="nav-item">
<a class="nav-link mb-0 px-0 py-1" data-toggle="tab" href="#html-chunking-tab" role="tab" aria-controls="html-chunking" aria-selected="false">
HTML Chunking
</a>
</li>
<li class="nav-item">
<a class="nav-link mb-0 px-0 py-1" data-toggle="tab" href="#domains-tab" role="tab" aria-controls="domains" aria-selected="false">
Domains
</a>
</li>
<li class="nav-item">
<a class="nav-link mb-0 px-0 py-1" data-toggle="tab" href="#users-tab" role="tab" aria-controls="users" aria-selected="false">
Users
</a>
</li>
</ul>
</div>
<div class="tab-content tab-space">
<!-- Model Information Tab -->
<div class="tab-pane fade show active" id="model-info-tab" role="tabpanel">
{% set model_fields = ['embedding_model', 'llm_model'] %}
{% for field in form %}
{{ render_included_field(field, disabled_fields=model_fields, include_fields=model_fields) }}
{% endfor %}
</div>
<!-- License Information Tab -->
<div class="tab-pane fade" id="license-info-tab" role="tabpanel">
{% set license_fields = ['license_start_date', 'license_end_date', 'allowed_monthly_interactions', ] %}
{% for field in form %}
{{ render_included_field(field, disabled_fields=license_fields, include_fields=license_fields) }}
{% endfor %}
<!-- Register API Key Button -->
<button type="button" class="btn btn-primary" onclick="checkAndRegisterApiKey()">Register API Key</button>
</div>
<!-- HTML Chunking Settings Tab -->
<div class="tab-pane fade" id="html-chunking-tab" role="tabpanel">
{% set html_fields = ['html_tags', 'html_end_tags', 'html_included_elements', 'html_excluded_elements', ] %}
{% for field in form %}
{{ render_included_field(field, disabled_fields=html_fields, include_fields=html_fields) }}
{% endfor %}
</div>
<!-- Domains Tab -->
<div class="tab-pane fade" id="domains-tab" role="tabpanel">
<ul>
UNDER CONSTRUCTION
</ul>
</div>
<!-- Users Tab -->
<div class="tab-pane fade" id="users-tab" role="tabpanel">
<ul>
UNDER CONSTRUCTION
</ul>
</div>
</div>
</div>
</div>
</form>
<!-- Modal HTML -->
<div class="modal fade" id="confirmModal" tabindex="-1" role="dialog" aria-labelledby="confirmModalLabel" aria-hidden="true">
<div class="modal-dialog" role="document">
<div class="modal-content">
<div class="modal-header">
<h5 class="modal-title" id="confirmModalLabel">Confirm New API Key</h5>
<button type="button" class="close" data-dismiss="modal" aria-label="Close">
<span aria-hidden="true">&times;</span>
</button>
</div>
<div class="modal-body" id="modal-body-content">
Are you sure you want to register a new API key? This will replace the existing key.
</div>
<div class="modal-footer">
<button type="button" class="btn btn-secondary" data-dismiss="modal">Cancel</button>
<button type="button" class="btn btn-primary" id="confirmNewKeyBtn">Confirm</button>
</div>
</div>
</div>
</div>
{% endblock %}
{% block content_footer %}
{% endblock %}
{% block scripts %}
<script>
function checkAndRegisterApiKey() {
// First, check if an API key already exists
$.ajax({
url: '/user/check_chat_api_key',
type: 'POST',
contentType: 'application/json',
success: function(response) {
if (response.api_key_exists) {
$('#confirmModal').modal('show');
} else {
generateNewApiKey();
}
},
error: function(error) {
alert('Error checking API key: ' + error.responseText);
}
});
}
document.getElementById('confirmNewKeyBtn').addEventListener('click', function () {
generateNewApiKey();
});
function generateNewApiKey() {
$.ajax({
url: '/user/generate_chat_api_key',
type: 'POST',
contentType: 'application/json',
success: function(response) {
$('#modal-body-content').html(`
<p>New API key generated: <span id="new-api-key">${response.api_key}</span></p>
<button class="btn btn-primary" onclick="copyToClipboard('#new-api-key')">Copy to Clipboard</button>
<p id="copy-message" style="display:none;color:green;">API key copied to clipboard</p>
`);
$('#confirmNewKeyBtn').hide();
$('.btn-secondary').text('OK');
},
error: function(error) {
alert('Error generating new API key: ' + error.responseText);
}
});
}
function copyToClipboard(element) {
const text = $(element).text();
navigator.clipboard.writeText(text).then(function() {
$('#copy-message').show().delay(2000).fadeOut();
}).catch(function(error) {
alert('Failed to copy text: ' + error);
});
}
</script>
{% endblock %}

View File

@@ -1,16 +1,17 @@
# from . import user_bp
import uuid
from datetime import datetime as dt, timezone as tz
from flask import request, redirect, url_for, flash, render_template, Blueprint, session, current_app
from flask import request, redirect, url_for, flash, render_template, Blueprint, session, current_app, jsonify
from flask_security import hash_password, roles_required, roles_accepted, current_user
from sqlalchemy.exc import SQLAlchemyError
import ast
from common.models.user import User, Tenant, Role, TenantDomain
from common.extensions import db
from common.extensions import db, kms_client
from .user_forms import TenantForm, CreateUserForm, EditUserForm, TenantDomainForm
from common.utils.database import Database
from common.utils.view_assistants import prepare_table_for_macro
from common.utils.key_encryption import generate_api_key
user_bp = Blueprint('user_bp', __name__, url_prefix='/user')
@@ -324,6 +325,45 @@ def edit_tenant_domain(tenant_domain_id):
return render_template('user/edit_tenant_domain.html', form=form, tenant_domain_id=tenant_domain_id)
@user_bp.route('/check_chat_api_key', methods=['POST'])
@roles_accepted('Super User', 'Tenant Admin')
def check_chat_api_key():
tenant_id = session['tenant']['id']
tenant = Tenant.query.get_or_404(tenant_id)
if tenant.encrypted_chat_api_key:
return jsonify({'api_key_exists': True})
return jsonify({'api_key_exists': False})
@user_bp.route('/generate_chat_api_key', methods=['POST'])
@roles_accepted('Super User', 'Tenant Admin')
def generate_chat_api_key():
tenant = Tenant.query.get_or_404(session['tenant']['id'])
new_api_key = generate_api_key(prefix="EveAI-CHAT")
tenant.encrypted_chat_api_key = kms_client.encrypt_api_key(new_api_key)
update_logging_information(tenant, dt.now(tz.utc))
try:
db.session.add(tenant)
db.session.commit()
except SQLAlchemyError as e:
db.session.rollback()
current_app.logger.error(f'Unable to store api key for tenant {tenant.id}. Error: {str(e)}')
return jsonify({'new_api_key': 'API key generated successfully.', 'api_key': new_api_key}), 200
@user_bp.route('/tenant_overview', methods=['GET'])
@roles_accepted('Super User', 'Tenant Admin')
def tenant_overview():
tenant_id = session['tenant']['id']
tenant = Tenant.query.get_or_404(tenant_id)
form = TenantForm(obj=tenant)
return render_template('user/tenant_overview.html', form=form)
def set_logging_information(obj, timestamp):
obj.created_at = timestamp
obj.updated_at = timestamp
@@ -332,5 +372,5 @@ def set_logging_information(obj, timestamp):
def update_logging_information(obj, timestamp):
obj.created_by = current_user.id
obj.updated_at = timestamp
obj.updated_by = current_user.id