Adding api-key based authentication for Darwin APIs

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

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}")

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}")

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

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

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

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")

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("")

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):

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(

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):

backend/danswer/server/gpts/api.py

@@ -6,18 +6,18 @@
 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
 from danswer.search.pipeline import SearchPipeline
 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:

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"

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")

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")

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()
 
 

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(

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

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")

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")

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("")