From f2b81db424e6b47a15b353a17b56710f87d03aea Mon Sep 17 00:00:00 2001
From: swati354 <swatitiwari354@gmail.com>
Date: Wed, 26 Mar 2025 11:08:46 +0530
Subject: [PATCH] Add api key authentication for DarwinAPIs

---
 backend/danswer/auth/api_key.py               | 32 +++++++++++++++++++
 .../danswer/server/danswer_api/ingestion.py   |  4 +--
 backend/danswer/server/documents/cc_pair.py   |  3 +-
 backend/danswer/server/documents/connector.py |  3 +-
 .../danswer/server/documents/credential.py    |  3 +-
 backend/danswer/server/documents/document.py  |  3 +-
 .../server/features/document_set/api.py       |  3 +-
 backend/danswer/server/features/folder/api.py |  3 +-
 .../danswer/server/features/persona/api.py    |  5 +--
 backend/danswer/server/features/prompt/api.py |  3 +-
 backend/danswer/server/features/tool/api.py   |  5 +--
 backend/danswer/server/gpts/api.py            |  4 +--
 .../danswer/server/manage/administrative.py   |  3 +-
 backend/danswer/server/manage/get_state.py    |  4 ++-
 backend/danswer/server/manage/llm/api.py      |  5 +--
 .../danswer/server/manage/secondary_index.py  |  3 +-
 backend/danswer/server/manage/slack_bot.py    |  3 +-
 backend/danswer/server/manage/users.py        |  3 +-
 .../server/query_and_chat/chat_backend.py     |  4 ++-
 .../server/query_and_chat/query_backend.py    |  5 +--
 backend/danswer/server/settings/api.py        |  5 +--
 .../danswer/server/token_rate_limits/api.py   |  3 +-
 22 files changed, 81 insertions(+), 28 deletions(-)
 create mode 100644 backend/danswer/auth/api_key.py

diff --git a/backend/danswer/auth/api_key.py b/backend/danswer/auth/api_key.py
new file mode 100644
index 00000000000..64eec777b75
--- /dev/null
+++ b/backend/danswer/auth/api_key.py
@@ -0,0 +1,32 @@
+from danswer.db.engine import get_session
+from danswer.db.models import ApiKey
+from danswer.utils.logger import setup_logger
+
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import Request
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+
+
+logger = setup_logger()
+
+_API_KEY_HEADER = "X-API-Key"
+
+def validate_api_key(request: Request, db_session: Session = Depends(get_session)):
+    if _API_KEY_HEADER not in request.headers:
+        return None
+
+    api_key_value = request.headers.get(_API_KEY_HEADER)
+    if not api_key_value:
+        raise HTTPException(status_code=401, detail="Missing API key")
+
+    api_key = db_session.scalar(
+        select(ApiKey).where(ApiKey.hashed_api_key == api_key_value)
+    )
+    if not api_key:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    return None
diff --git a/backend/danswer/server/danswer_api/ingestion.py b/backend/danswer/server/danswer_api/ingestion.py
index 1b6e6d9852f..bd5d3f85e0f 100644
--- a/backend/danswer/server/danswer_api/ingestion.py
+++ b/backend/danswer/server/danswer_api/ingestion.py
@@ -21,11 +21,11 @@
 from danswer.server.danswer_api.models import IngestionResult
 from danswer.utils.logger import setup_logger
 from ee.danswer.auth.users import api_key_dep
-
+from danswer.auth.api_key import validate_api_key
 logger = setup_logger()
 
 # not using /api to avoid confusion with nginx api path routing
-router = APIRouter(prefix="/danswer-api")
+router = APIRouter(prefix="/danswer-api", dependencies=[Depends(validate_api_key)])
 
 
 @router.get("/connector-docs/{cc_pair_id}")
diff --git a/backend/danswer/server/documents/cc_pair.py b/backend/danswer/server/documents/cc_pair.py
index c6026401d6c..c0c17892239 100644
--- a/backend/danswer/server/documents/cc_pair.py
+++ b/backend/danswer/server/documents/cc_pair.py
@@ -4,6 +4,7 @@
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.background.celery.celery_utils import get_deletion_status
@@ -19,7 +20,7 @@
 from danswer.server.documents.models import ConnectorCredentialPairMetadata
 from danswer.server.models import StatusResponse
 
-router = APIRouter(prefix="/manage")
+router = APIRouter(prefix="/manage", dependencies=[Depends(validate_api_key)])
 
 
 @router.get("/admin/cc-pair/{cc_pair_id}")
diff --git a/backend/danswer/server/documents/connector.py b/backend/danswer/server/documents/connector.py
index ad25523817d..6ca9d23ab90 100644
--- a/backend/danswer/server/documents/connector.py
+++ b/backend/danswer/server/documents/connector.py
@@ -11,6 +11,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.background.celery.celery_utils import get_deletion_status
@@ -90,7 +91,7 @@
 _GOOGLE_DRIVE_CREDENTIAL_ID_COOKIE_NAME = "google_drive_credential_id"
 
 
-router = APIRouter(prefix="/manage")
+router = APIRouter(prefix="/manage", dependencies=[Depends(validate_api_key)])
 
 
 """Admin only API endpoints"""
diff --git a/backend/danswer/server/documents/credential.py b/backend/danswer/server/documents/credential.py
index a5e9098046a..82641355381 100644
--- a/backend/danswer/server/documents/credential.py
+++ b/backend/danswer/server/documents/credential.py
@@ -3,6 +3,7 @@
 from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.schemas import UserRole
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
@@ -19,7 +20,7 @@
 from danswer.server.models import StatusResponse
 
 
-router = APIRouter(prefix="/manage")
+router = APIRouter(prefix="/manage", dependencies=[Depends(validate_api_key)])
 
 
 """Admin-only endpoints"""
diff --git a/backend/danswer/server/documents/document.py b/backend/danswer/server/documents/document.py
index 3b0adea246c..059b70758ce 100644
--- a/backend/danswer/server/documents/document.py
+++ b/backend/danswer/server/documents/document.py
@@ -4,6 +4,7 @@
 from fastapi import Query
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_user
 from danswer.db.embedding_model import get_current_db_embedding_model
 from danswer.db.engine import get_session
@@ -16,7 +17,7 @@
 from danswer.server.documents.models import DocumentInfo
 
 
-router = APIRouter(prefix="/document")
+router = APIRouter(prefix="/document", dependencies=[Depends(validate_api_key)])
 
 
 # Have to use a query parameter as FastAPI is interpreting the URL type document_ids
diff --git a/backend/danswer/server/features/document_set/api.py b/backend/danswer/server/features/document_set/api.py
index f939329bf9a..3cdaf7b9c21 100644
--- a/backend/danswer/server/features/document_set/api.py
+++ b/backend/danswer/server/features/document_set/api.py
@@ -3,6 +3,7 @@
 from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.db.document_set import check_document_sets_are_public
@@ -23,7 +24,7 @@
 from danswer.server.features.document_set.models import DocumentSetUpdateRequest
 
 
-router = APIRouter(prefix="/manage")
+router = APIRouter(prefix="/manage", dependencies=[Depends(validate_api_key)])
 
 
 @router.post("/admin/document-set")
diff --git a/backend/danswer/server/features/folder/api.py b/backend/danswer/server/features/folder/api.py
index 000207370d6..754e3693dab 100644
--- a/backend/danswer/server/features/folder/api.py
+++ b/backend/danswer/server/features/folder/api.py
@@ -4,6 +4,7 @@
 from fastapi import Path
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_user
 from danswer.db.chat import get_chat_session_by_id
 from danswer.db.engine import get_session
@@ -24,7 +25,7 @@
 from danswer.server.models import DisplayPriorityRequest
 from danswer.server.query_and_chat.models import ChatSessionDetails
 
-router = APIRouter(prefix="/folder")
+router = APIRouter(prefix="/folder", dependencies=[Depends(validate_api_key)])
 
 
 @router.get("")
diff --git a/backend/danswer/server/features/persona/api.py b/backend/danswer/server/features/persona/api.py
index 6739da46606..bbe48cf50cc 100644
--- a/backend/danswer/server/features/persona/api.py
+++ b/backend/danswer/server/features/persona/api.py
@@ -5,6 +5,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.db.engine import get_session
@@ -27,8 +28,8 @@
 logger = setup_logger()
 
 
-admin_router = APIRouter(prefix="/admin/persona")
-basic_router = APIRouter(prefix="/persona")
+admin_router = APIRouter(prefix="/admin/persona", dependencies=[Depends(validate_api_key)])
+basic_router = APIRouter(prefix="/persona", dependencies=[Depends(validate_api_key)])
 
 
 class IsVisibleRequest(BaseModel):
diff --git a/backend/danswer/server/features/prompt/api.py b/backend/danswer/server/features/prompt/api.py
index aebcbb8434d..6a4ddbec18d 100644
--- a/backend/danswer/server/features/prompt/api.py
+++ b/backend/danswer/server/features/prompt/api.py
@@ -4,6 +4,7 @@
 from sqlalchemy.orm import Session
 from starlette import status
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_user
 from danswer.db.engine import get_session
 from danswer.db.models import User
@@ -22,7 +23,7 @@
 
 logger = setup_logger()
 
-basic_router = APIRouter(prefix="/prompt")
+basic_router = APIRouter(prefix="/prompt", dependencies=[Depends(validate_api_key)])
 
 
 def create_update_prompt(
diff --git a/backend/danswer/server/features/tool/api.py b/backend/danswer/server/features/tool/api.py
index b1f57a1a924..f403633ce11 100644
--- a/backend/danswer/server/features/tool/api.py
+++ b/backend/danswer/server/features/tool/api.py
@@ -6,6 +6,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.db.engine import get_session
@@ -20,8 +21,8 @@
 from danswer.tools.custom.openapi_parsing import openapi_to_method_specs
 from danswer.tools.custom.openapi_parsing import validate_openapi_schema
 
-router = APIRouter(prefix="/tool")
-admin_router = APIRouter(prefix="/admin/tool")
+router = APIRouter(prefix="/tool", dependencies=[Depends(validate_api_key)])
+admin_router = APIRouter(prefix="/admin/tool", dependencies=[Depends(validate_api_key)])
 
 
 class CustomToolCreate(BaseModel):
diff --git a/backend/danswer/server/gpts/api.py b/backend/danswer/server/gpts/api.py
index 84b0078ee77..3be9c47891e 100644
--- a/backend/danswer/server/gpts/api.py
+++ b/backend/danswer/server/gpts/api.py
@@ -6,6 +6,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.db.engine import get_session
 from danswer.llm.factory import get_default_llms
 from danswer.search.models import SearchRequest
@@ -13,11 +14,10 @@
 from danswer.server.danswer_api.ingestion import api_key_dep
 from danswer.utils.logger import setup_logger
 
-
 logger = setup_logger()
 
 
-router = APIRouter(prefix="/gpts")
+router = APIRouter(prefix="/gpts", dependencies=[Depends(validate_api_key)])
 
 
 def time_ago(dt: datetime) -> str:
diff --git a/backend/danswer/server/manage/administrative.py b/backend/danswer/server/manage/administrative.py
index d6a52917f3b..84d0b390f57 100644
--- a/backend/danswer/server/manage/administrative.py
+++ b/backend/danswer/server/manage/administrative.py
@@ -8,6 +8,7 @@
 from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.configs.app_configs import GENERATIVE_MODEL_ACCESS_CHECK_FREQ
 from danswer.configs.constants import DocumentSource
@@ -32,7 +33,7 @@
 from danswer.server.manage.models import HiddenUpdateRequest
 from danswer.utils.logger import setup_logger
 
-router = APIRouter(prefix="/manage")
+router = APIRouter(prefix="/manage", dependencies=[Depends(validate_api_key)])
 logger = setup_logger()
 
 GEN_AI_KEY_CHECK_TIME = "genai_api_key_last_check_time"
diff --git a/backend/danswer/server/manage/get_state.py b/backend/danswer/server/manage/get_state.py
index 3ca47841b64..5d61e4f4376 100644
--- a/backend/danswer/server/manage/get_state.py
+++ b/backend/danswer/server/manage/get_state.py
@@ -1,13 +1,15 @@
 from fastapi import APIRouter
+from fastapi import Depends
 
 from danswer import __version__
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import user_needs_to_be_verified
 from danswer.configs.app_configs import AUTH_TYPE
 from danswer.server.manage.models import AuthTypeResponse
 from danswer.server.manage.models import VersionResponse
 from danswer.server.models import StatusResponse
 
-router = APIRouter()
+router = APIRouter(dependencies=[Depends(validate_api_key)])
 
 
 @router.get("/health")
diff --git a/backend/danswer/server/manage/llm/api.py b/backend/danswer/server/manage/llm/api.py
index 4df00b529af..71047ac36e2 100644
--- a/backend/danswer/server/manage/llm/api.py
+++ b/backend/danswer/server/manage/llm/api.py
@@ -5,6 +5,7 @@
 from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.db.engine import get_session
@@ -28,8 +29,8 @@
 logger = setup_logger()
 
 
-admin_router = APIRouter(prefix="/admin/llm")
-basic_router = APIRouter(prefix="/llm")
+admin_router = APIRouter(prefix="/admin/llm", dependencies=[Depends(validate_api_key)])
+basic_router = APIRouter(prefix="/llm", dependencies=[Depends(validate_api_key)])
 
 
 @admin_router.get("/built-in/options")
diff --git a/backend/danswer/server/manage/secondary_index.py b/backend/danswer/server/manage/secondary_index.py
index 6f5adf752f6..f52b944dbab 100644
--- a/backend/danswer/server/manage/secondary_index.py
+++ b/backend/danswer/server/manage/secondary_index.py
@@ -4,6 +4,7 @@
 from fastapi import status
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP
@@ -23,7 +24,7 @@
 from danswer.server.models import IdReturn
 from danswer.utils.logger import setup_logger
 
-router = APIRouter(prefix="/secondary-index")
+router = APIRouter(prefix="/secondary-index", dependencies=[Depends(validate_api_key)])
 logger = setup_logger()
 
 
diff --git a/backend/danswer/server/manage/slack_bot.py b/backend/danswer/server/manage/slack_bot.py
index 71ee4df1f27..da4a80988a5 100644
--- a/backend/danswer/server/manage/slack_bot.py
+++ b/backend/danswer/server/manage/slack_bot.py
@@ -3,6 +3,7 @@
 from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.danswerbot.slack.config import validate_channel_names
 from danswer.danswerbot.slack.tokens import fetch_tokens
@@ -24,7 +25,7 @@
 from danswer.server.manage.models import SlackBotTokens
 
 
-router = APIRouter(prefix="/manage")
+router = APIRouter(prefix="/manage", dependencies=[Depends(validate_api_key)])
 
 
 def _form_channel_config(
diff --git a/backend/danswer/server/manage/users.py b/backend/danswer/server/manage/users.py
index c635469919e..8e505755861 100644
--- a/backend/danswer/server/manage/users.py
+++ b/backend/danswer/server/manage/users.py
@@ -9,6 +9,7 @@
 from sqlalchemy import update
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.invited_users import get_invited_users
 from danswer.auth.invited_users import write_invited_users
 from danswer.auth.noauth_user import fetch_no_auth_user
@@ -38,7 +39,7 @@
 
 logger = setup_logger()
 
-router = APIRouter()
+router = APIRouter(dependencies=[Depends(validate_api_key)])
 
 
 USERS_PAGE_SIZE = 10
diff --git a/backend/danswer/server/query_and_chat/chat_backend.py b/backend/danswer/server/query_and_chat/chat_backend.py
index 646660c9fac..f03dcb93397 100644
--- a/backend/danswer/server/query_and_chat/chat_backend.py
+++ b/backend/danswer/server/query_and_chat/chat_backend.py
@@ -11,6 +11,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_user
 from danswer.chat.chat_utils import create_chat_chain
 from danswer.chat.process_message import stream_chat_message
@@ -69,7 +70,8 @@
 
 logger = setup_logger()
 
-router = APIRouter(prefix="/chat")
+router = APIRouter(prefix="/chat", dependencies=[Depends(validate_api_key)])
+# api_router = APIRouter(prefix="/chat", dependencies=[Depends(validate_api_key)])
 
 
 @router.get("/get-user-chat-sessions")
diff --git a/backend/danswer/server/query_and_chat/query_backend.py b/backend/danswer/server/query_and_chat/query_backend.py
index 43192211b79..ff632e0613a 100644
--- a/backend/danswer/server/query_and_chat/query_backend.py
+++ b/backend/danswer/server/query_and_chat/query_backend.py
@@ -4,6 +4,7 @@
 from fastapi.responses import StreamingResponse
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.configs.constants import DocumentSource
@@ -34,8 +35,8 @@
 
 logger = setup_logger()
 
-admin_router = APIRouter(prefix="/admin")
-basic_router = APIRouter(prefix="/query")
+admin_router = APIRouter(prefix="/admin", dependencies=[Depends(validate_api_key)])
+basic_router = APIRouter(prefix="/query", dependencies=[Depends(validate_api_key)])
 
 
 @admin_router.post("/search")
diff --git a/backend/danswer/server/settings/api.py b/backend/danswer/server/settings/api.py
index 422e268c13e..756392cbcc8 100644
--- a/backend/danswer/server/settings/api.py
+++ b/backend/danswer/server/settings/api.py
@@ -2,6 +2,7 @@
 from fastapi import Depends
 from fastapi import HTTPException
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_user
 from danswer.db.models import User
@@ -10,8 +11,8 @@
 from danswer.server.settings.store import store_settings
 
 
-admin_router = APIRouter(prefix="/admin/settings")
-basic_router = APIRouter(prefix="/settings")
+admin_router = APIRouter(prefix="/admin/settings", dependencies=[Depends(validate_api_key)])
+basic_router = APIRouter(prefix="/settings", dependencies=[Depends(validate_api_key)])
 
 
 @admin_router.put("")
diff --git a/backend/danswer/server/token_rate_limits/api.py b/backend/danswer/server/token_rate_limits/api.py
index 245e3391410..e5488cf1acf 100644
--- a/backend/danswer/server/token_rate_limits/api.py
+++ b/backend/danswer/server/token_rate_limits/api.py
@@ -2,6 +2,7 @@
 from fastapi import Depends
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import validate_api_key
 from danswer.auth.users import current_admin_user
 from danswer.db.engine import get_session
 from danswer.db.models import User
@@ -13,7 +14,7 @@
 from ee.danswer.db.token_limit import insert_global_token_rate_limit
 from ee.danswer.db.token_limit import update_token_rate_limit
 
-router = APIRouter(prefix="/admin/token-rate-limits")
+router = APIRouter(prefix="/admin/token-rate-limits", dependencies=[Depends(validate_api_key)])
 
 
 """