From 0f3bf9192d97dea8c19c1a4edf7bdf019b906794 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Sat, 26 Aug 2023 17:38:43 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=92=20chore(utils.py):=20refactor=20im?=
 =?UTF-8?q?port=20statements=20and=20remove=20unused=20imports=20for=20bet?=
 =?UTF-8?q?ter=20code=20organization=20and=20readability=20=F0=9F=94=92=20?=
 =?UTF-8?q?chore(utils.py):=20refactor=20auth=5Fscheme=5Fdependency=20to?=
 =?UTF-8?q?=20use=20OAuth2PasswordBearer=20for=20better=20security=20?=
 =?UTF-8?q?=F0=9F=94=92=20chore(utils.py):=20refactor=20validate=5Fapi=5Fk?=
 =?UTF-8?q?ey=20to=20use=20check=5Fkey=20function=20from=20crud=20module?=
 =?UTF-8?q?=20for=20better=20code=20reuse=20=F0=9F=94=92=20chore(utils.py)?=
 =?UTF-8?q?:=20add=20api=5Fkey=5Fsecurity=20function=20to=20handle=20API?=
 =?UTF-8?q?=20key=20authentication=20logic=20=F0=9F=94=92=20chore(utils.py?=
 =?UTF-8?q?):=20refactor=20get=5Fcurrent=5Fuser=20to=20use=20oauth2=5Flogi?=
 =?UTF-8?q?n=20dependency=20for=20better=20security=20=F0=9F=94=92=20chore?=
 =?UTF-8?q?(utils.py):=20remove=20unused=20validate=5Fapi=5Fkey=20function?=
 =?UTF-8?q?=20=F0=9F=94=92=20chore(auth.py):=20remove=20unused=20oauth2=5F?=
 =?UTF-8?q?scheme=20variable=20for=20better=20code=20organization?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/services/auth/utils.py   | 60 +++++++++++++------
 .../langflow/services/settings/auth.py        |  2 -
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py
index 4904a179b..f89bd6c88 100644
--- a/src/backend/langflow/services/auth/utils.py
+++ b/src/backend/langflow/services/auth/utils.py
@@ -1,10 +1,10 @@
 from datetime import datetime, timedelta, timezone
-from fastapi import Depends, HTTPException, Request, status
+from fastapi import Depends, HTTPException, Security, status
+from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer
 from jose import JWTError, jwt
 from typing import Annotated, Coroutine
 from uuid import UUID
-from langflow.services.auth.service import AuthManager
-from langflow.services.database.models.api_key.api_key import ApiKey
+from langflow.services.database.models.api_key.crud import check_key
 from langflow.services.database.models.user.user import User
 from langflow.services.database.models.user.crud import (
     get_user_by_id,
@@ -12,19 +12,51 @@ from langflow.services.database.models.user.crud import (
     update_user_last_login_at,
 )
 from langflow.services.utils import get_session, get_settings_manager
-from sqlmodel import Session, select
+from sqlmodel import Session
+
+oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login")
+
+API_KEY_NAME = "api-key"
+
+api_key_query = APIKeyQuery(
+    name=API_KEY_NAME, scheme_name="API key query", auto_error=False
+)
+api_key_header = APIKeyHeader(
+    name=API_KEY_NAME, scheme_name="API key header", auto_error=False
+)
 
 
-async def auth_scheme_dependency(request: Request):
-    settings_manager = (
-        get_settings_manager()
-    )  # Assuming get_settings_manager is defined
+# Source: https://github.com/mrtolkien/fastapi_simple_security/blob/master/fastapi_simple_security/security_api_key.py
+async def api_key_security(
+    query_param: str = Security(api_key_query),
+    header_param: str = Security(api_key_header),
+    db: Session = Depends(get_session),
+):
+    settings_manager = get_settings_manager()
+    if settings_manager.auth_settings.AUTO_LOGIN:
+        return settings_manager.auth_settings.API_KEY_SECRET_KEY
 
-    return await AuthManager(settings_manager).run_oauth2_scheme(request)
+    elif not query_param and not header_param:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="An API key must be passed as query or header",
+        )
+
+    elif query_param and check_key(db, query_param):
+        return query_param
+
+    elif header_param and check_key(db, header_param):
+        return header_param
+
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Invalid or missing API key",
+        )
 
 
 async def get_current_user(
-    token: Annotated[str, Depends(auth_scheme_dependency)],
+    token: Annotated[str, Depends(oauth2_login)],
     db: Session = Depends(get_session),
 ) -> User:
     settings_manager = get_settings_manager()
@@ -58,14 +90,6 @@ async def get_current_user(
     return user
 
 
-async def validate_api_key(
-    token: Annotated[str, Depends(auth_scheme_dependency)],
-    db: Session = Depends(get_session),
-) -> bool:
-    hashed_api_key = get_password_hash(token)
-    return db.exec(select(ApiKey).where(ApiKey.hashed_api_key == hashed_api_key)).one_or_none()  # type: ignore
-
-
 def get_current_active_user(current_user: Annotated[User, Depends(get_current_user)]):
     if not current_user.is_active:
         raise HTTPException(status_code=400, detail="Inactive user")
diff --git a/src/backend/langflow/services/settings/auth.py b/src/backend/langflow/services/settings/auth.py
index 2aa4e17bc..40147ba15 100644
--- a/src/backend/langflow/services/settings/auth.py
+++ b/src/backend/langflow/services/settings/auth.py
@@ -3,7 +3,6 @@ import secrets
 
 from pydantic import BaseSettings
 from passlib.context import CryptContext
-from fastapi.security import OAuth2PasswordBearer
 
 
 class AuthSettings(BaseSettings):
@@ -27,7 +26,6 @@ class AuthSettings(BaseSettings):
     FIRST_SUPERUSER_PASSWORD: str = "langflow"
 
     pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
-    oauth2_scheme = OAuth2PasswordBearer(tokenUrl=f"{API_V1_STR}/login")
 
     class Config:
         validate_assignment = True