From 596fc8ed411e95d8c4b0e7fc1596d6ddb6394ec0 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 18 Oct 2023 17:03:53 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(api=5Fkey.py):=20import=20mi?=
 =?UTF-8?q?ssing=20dependencies=20and=20fix=20import=20paths=20for=20auth?=
 =?UTF-8?q?=5Futils=20and=20get=5Fcurrent=5Factive=5Fuser=20=E2=9C=A8=20fe?=
 =?UTF-8?q?at(api=5Fkey.py):=20add=20new=20route=20to=20save=20store=20API?=
 =?UTF-8?q?=20key=20for=20a=20user?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/api/v1/api_key.py | 42 ++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/backend/langflow/api/v1/api_key.py b/src/backend/langflow/api/v1/api_key.py
index 7de77389e..c804746fc 100644
--- a/src/backend/langflow/api/v1/api_key.py
+++ b/src/backend/langflow/api/v1/api_key.py
@@ -1,7 +1,8 @@
 from uuid import UUID
-from fastapi import APIRouter, HTTPException, Depends
+from fastapi import APIRouter, HTTPException, Depends, Body
+from httpx import get
 from langflow.api.v1.schemas import ApiKeysResponse
-from langflow.services.auth.utils import get_current_active_user
+from langflow.services.auth import utils as auth_utils
 from langflow.services.database.models.api_key.api_key import (
     ApiKeyCreate,
     UnmaskedApiKeyRead,
@@ -14,9 +15,19 @@ from langflow.services.database.models.api_key.crud import (
     delete_api_key,
 )
 from langflow.services.database.models.user.user import User
-from langflow.services.deps import get_session
+from langflow.services.deps import (
+    get_store_service,
+    get_session,
+    get_settings_service,
+)
+from typing import TYPE_CHECKING
+
+
 from sqlmodel import Session
 
+if TYPE_CHECKING:
+    from langflow.services.store.service import StoreService
+    from langflow.services.settings.service import SettingsService
 
 router = APIRouter(tags=["APIKey"], prefix="/api_key")
 
@@ -24,7 +35,7 @@ router = APIRouter(tags=["APIKey"], prefix="/api_key")
 @router.get("/", response_model=ApiKeysResponse)
 def get_api_keys_route(
     db: Session = Depends(get_session),
-    current_user: User = Depends(get_current_active_user),
+    current_user: User = Depends(auth_utils.auth_utils.get_current_active_user),
 ):
     try:
         user_id = current_user.id
@@ -38,7 +49,7 @@ def get_api_keys_route(
 @router.post("/", response_model=UnmaskedApiKeyRead)
 def create_api_key_route(
     req: ApiKeyCreate,
-    current_user: User = Depends(get_current_active_user),
+    current_user: User = Depends(auth_utils.get_current_active_user),
     db: Session = Depends(get_session),
 ):
     try:
@@ -51,7 +62,7 @@ def create_api_key_route(
 @router.delete("/{api_key_id}")
 def delete_api_key_route(
     api_key_id: UUID,
-    current_user=Depends(get_current_active_user),
+    current_user=Depends(auth_utils.get_current_active_user),
     db: Session = Depends(get_session),
 ):
     try:
@@ -59,3 +70,22 @@ def delete_api_key_route(
         return {"detail": "API Key deleted"}
     except Exception as e:
         raise HTTPException(status_code=400, detail=str(e)) from e
+
+
+@router.post("/store")
+def save_store_api_key(
+    api_key: Body(str, embed=True),
+    current_user: User = Depends(auth_utils.get_current_active_user),
+    db: Session = Depends(get_session),
+    settings_service=Depends(get_settings_service),
+):
+    try:
+        # Encrypt the API key
+        encrypted = auth_utils.encrypt_api_key(
+            api_key, fernet=auth_utils.get_fernet(settings_service)
+        )
+        current_user.store_api_key = encrypted
+        db.commit()
+        return {"detail": "API Key saved"}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e)) from e