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 @@