From 883988dbabe67570078e9271fa793f1f143455c8 Mon Sep 17 00:00:00 2001 From: Josako Date: Thu, 16 May 2024 23:22:26 +0200 Subject: [PATCH] Implement chat API key generation, and create a tenant_overview --- common/extensions.py | 2 + common/models/user.py | 2 +- common/utils/key_encryption.py | 100 +++++----- config/config.py | 5 +- eveai_app/__init__.py | 3 +- eveai_app/templates/base.html | 1 + eveai_app/templates/macros.html | 24 +++ eveai_app/templates/navbar.html | 3 +- eveai_app/templates/scripts.html | 12 +- .../templates/user/generate_chat_api_key.html | 70 +++++++ eveai_app/templates/user/tenant_overview.html | 171 ++++++++++++++++++ eveai_app/views/user_views.py | 46 ++++- 12 files changed, 387 insertions(+), 52 deletions(-) create mode 100644 eveai_app/templates/user/generate_chat_api_key.html create mode 100644 eveai_app/templates/user/tenant_overview.html diff --git a/common/extensions.py b/common/extensions.py index f706f24..f617f3f 100644 --- a/common/extensions.py +++ b/common/extensions.py @@ -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() diff --git a/common/models/user.py b/common/models/user.py index 868343f..9eb5ffa 100644 --- a/common/models/user.py +++ b/common/models/user.py @@ -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') diff --git a/common/utils/key_encryption.py b/common/utils/key_encryption.py index 060a2b4..5513cba 100644 --- a/common/utils/key_encryption.py +++ b/common/utils/key_encryption.py @@ -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() diff --git a/config/config.py b/config/config.py index a627edc..7911bba 100644 --- a/config/config.py +++ b/config/config.py @@ -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' diff --git a/eveai_app/__init__.py b/eveai_app/__init__.py index 9ffdfa6..66da3be 100644 --- a/eveai_app/__init__.py +++ b/eveai_app/__init__.py @@ -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 diff --git a/eveai_app/templates/base.html b/eveai_app/templates/base.html index da02acc..8ca39d7 100644 --- a/eveai_app/templates/base.html +++ b/eveai_app/templates/base.html @@ -51,5 +51,6 @@
{% include 'footer.html' %} {% include 'scripts.html' %} + {% block scripts %}{% endblock %} \ No newline at end of file diff --git a/eveai_app/templates/macros.html b/eveai_app/templates/macros.html index 6398d43..ff39cf0 100644 --- a/eveai_app/templates/macros.html +++ b/eveai_app/templates/macros.html @@ -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' %} +
+ {{ field(class="form-check-input", type="checkbox", id="flexSwitchCheckDefault") }} + {{ field.label(class="form-check-label", for="flexSwitchCheckDefault", disabled=disabled) }} +
+ {% else %} +
+ {{ field.label(class="form-label") }} + {{ field(class="form-control", disabled=disabled) }} + {% if field.errors %} +
+ {% for error in field.errors %} + {{ error }} + {% endfor %} +
+ {% endif %} +
+ {% endif %} + {% endif %} +{% endmacro %} + {% macro render_table(headers, rows) %}
diff --git a/eveai_app/templates/navbar.html b/eveai_app/templates/navbar.html index 39b5c78..b8e6cf9 100644 --- a/eveai_app/templates/navbar.html +++ b/eveai_app/templates/navbar.html @@ -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']}, diff --git a/eveai_app/templates/scripts.html b/eveai_app/templates/scripts.html index 406c35f..d932840 100644 --- a/eveai_app/templates/scripts.html +++ b/eveai_app/templates/scripts.html @@ -1,8 +1,14 @@ - + - {% block scripts %} - {%- endblock scripts %} \ No newline at end of file + + + + + + + + \ No newline at end of file diff --git a/eveai_app/templates/user/generate_chat_api_key.html b/eveai_app/templates/user/generate_chat_api_key.html new file mode 100644 index 0000000..02e357c --- /dev/null +++ b/eveai_app/templates/user/generate_chat_api_key.html @@ -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 %} + + +{% endblock %} +{% block content_footer %} +{% endblock %} +{% block scripts %} + +{% endblock %} diff --git a/eveai_app/templates/user/tenant_overview.html b/eveai_app/templates/user/tenant_overview.html new file mode 100644 index 0000000..51f4e36 --- /dev/null +++ b/eveai_app/templates/user/tenant_overview.html @@ -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.hidden_tag() }} + + {% 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 %} + + +
+
+
+ +
+
+ +
+ {% set model_fields = ['embedding_model', 'llm_model'] %} + {% for field in form %} + {{ render_included_field(field, disabled_fields=model_fields, include_fields=model_fields) }} + {% endfor %} +
+ +
+ {% 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 %} + + +
+ +
+ {% 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 %} +
+ +
+
    + UNDER CONSTRUCTION +
+
+ +
+
    + UNDER CONSTRUCTION +
+
+
+
+
+
+ + +{% endblock %} + + +{% block content_footer %} + +{% endblock %} + +{% block scripts %} + +{% endblock %} \ No newline at end of file diff --git a/eveai_app/views/user_views.py b/eveai_app/views/user_views.py index 84e3703..5e857c6 100644 --- a/eveai_app/views/user_views.py +++ b/eveai_app/views/user_views.py @@ -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