From 43e3bc4ef54ba2b74aa0ed376e426f86dbd2a752 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 21 Mar 2024 13:22:51 -0300
Subject: [PATCH] Refactor authentication code and update settings

---
 src/backend/langflow/api/v1/login.py           |  4 +---
 src/backend/langflow/services/auth/utils.py    | 18 +++++++-----------
 src/backend/langflow/services/settings/auth.py | 15 +++++++--------
 3 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/src/backend/langflow/api/v1/login.py b/src/backend/langflow/api/v1/login.py
index 881b42797..4f44e79bc 100644
--- a/src/backend/langflow/api/v1/login.py
+++ b/src/backend/langflow/api/v1/login.py
@@ -89,9 +89,7 @@ async def auto_login(
 
 
 @router.post("/refresh")
-async def refresh_token(
-    request: Request, response: Response, settings_service=Depends(get_settings_service)
-):
+async def refresh_token(request: Request, response: Response, settings_service=Depends(get_settings_service)):
     auth_settings = settings_service.auth_settings
 
     token = request.cookies.get("refresh_token_lf")
diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py
index 23ce1e702..6babee5bf 100644
--- a/src/backend/langflow/services/auth/utils.py
+++ b/src/backend/langflow/services/auth/utils.py
@@ -9,13 +9,9 @@ from jose import JWTError, jwt
 from sqlmodel import Session
 from starlette.websockets import WebSocket
 
-from langflow.services.database.models.api_key.model import ApiKey
 from langflow.services.database.models.api_key.crud import check_key
-from langflow.services.database.models.user.crud import (
-    get_user_by_id,
-    get_user_by_username,
-    update_user_last_login_at,
-)
+from langflow.services.database.models.api_key.model import ApiKey
+from langflow.services.database.models.user.crud import get_user_by_id, get_user_by_username, update_user_last_login_at
 from langflow.services.database.models.user.model import User
 from langflow.services.deps import get_session, get_settings_service
 
@@ -107,13 +103,13 @@ async def get_current_user_by_jwt(
     if isinstance(token, Coroutine):
         token = await token
 
-    if settings_service.auth_settings.SECRET_KEY is None:
+    if settings_service.auth_settings.SECRET_KEY.get_secret_value() is None:
         raise credentials_exception
 
     try:
         payload = jwt.decode(
             token,
-            settings_service.auth_settings.SECRET_KEY,
+            settings_service.auth_settings.SECRET_KEY.get_secret_value(),
             algorithms=[settings_service.auth_settings.ALGORITHM],
         )
         user_id: UUID = payload.get("sub")  # type: ignore
@@ -183,7 +179,7 @@ def create_token(data: dict, expires_delta: timedelta):
 
     return jwt.encode(
         to_encode,
-        settings_service.auth_settings.SECRET_KEY,
+        settings_service.auth_settings.SECRET_KEY.get_secret_value(),
         algorithm=settings_service.auth_settings.ALGORITHM,
     )
 
@@ -287,7 +283,7 @@ def create_refresh_token(refresh_token: str, db: Session = Depends(get_session))
     try:
         payload = jwt.decode(
             refresh_token,
-            settings_service.auth_settings.SECRET_KEY,
+            settings_service.auth_settings.SECRET_KEY.get_secret_value(),
             algorithms=[settings_service.auth_settings.ALGORITHM],
         )
         user_id: UUID = payload.get("sub")  # type: ignore
@@ -326,7 +322,7 @@ def add_padding(s):
 
 
 def get_fernet(settings_service=Depends(get_settings_service)):
-    SECRET_KEY = settings_service.auth_settings.SECRET_KEY
+    SECRET_KEY = settings_service.auth_settings.SECRET_KEY.get_secret_value()
     # It's important that your secret key is 32 url-safe base64-encoded byte
     padded_secret_key = add_padding(SECRET_KEY)
     fernet = Fernet(padded_secret_key)
diff --git a/src/backend/langflow/services/settings/auth.py b/src/backend/langflow/services/settings/auth.py
index 8463d0781..11e08b1c4 100644
--- a/src/backend/langflow/services/settings/auth.py
+++ b/src/backend/langflow/services/settings/auth.py
@@ -1,23 +1,23 @@
 import secrets
 from pathlib import Path
-from typing import Optional
+
+from loguru import logger
+from passlib.context import CryptContext
+from pydantic import Field, SecretStr, validator
+from pydantic_settings import BaseSettings
 
 from langflow.services.settings.constants import (
     DEFAULT_SUPERUSER,
     DEFAULT_SUPERUSER_PASSWORD,
 )
 from langflow.services.settings.utils import read_secret_from_file, write_secret_to_file
-from loguru import logger
-from passlib.context import CryptContext
-from pydantic import Field, validator
-from pydantic_settings import BaseSettings
 
 
 class AuthSettings(BaseSettings):
     # Login settings
     CONFIG_DIR: str
-    SECRET_KEY: str = Field(
-        default="",
+    SECRET_KEY: SecretStr = Field(
+        default=None,
         description="Secret key for JWT. If not provided, a random one will be generated.",
         frozen=False,
     )
@@ -26,7 +26,6 @@ class AuthSettings(BaseSettings):
     REFRESH_TOKEN_EXPIRE_MINUTES: int = 60 * 12 * 7
 
     # API Key to execute /process endpoint
-    API_KEY_SECRET_KEY: Optional[str] = "b82818e0ad4ff76615c5721ee21004b07d84cd9b87ba4d9cb42374da134b841a"
     API_KEY_ALGORITHM: str = "HS256"
     API_V1_STR: str = "/api/v1"