added Api key auth flow

Author: avirajsingh7

backend/app/api/deps.py

@@ -1,10 +1,10 @@
 from collections.abc import Generator
-from typing import Annotated
+from typing import Annotated, Optional
 
 import jwt
-from fastapi import Depends, HTTPException, status, Request
+from fastapi import Depends, HTTPException, status, Request, Header, Security
 from fastapi.responses import JSONResponse
-from fastapi.security import OAuth2PasswordBearer
+from fastapi.security import OAuth2PasswordBearer, APIKeyHeader
 from jwt.exceptions import InvalidTokenError
 from pydantic import ValidationError
 from sqlmodel import Session, select
@@ -13,42 +13,66 @@
 from app.core.config import settings
 from app.core.db import engine
 from app.utils import APIResponse
-from app.models import TokenPayload, User, UserProjectOrg, ProjectUser, Project, Organization
+from app.crud.organization import validate_organization
+from app.crud.api_key import get_api_key_by_value
+from app.models import TokenPayload, User, UserProjectOrg, UserOrganization, ProjectUser, Project, Organization
 
-reusable_oauth2 = OAuth2PasswordBearer(
-    tokenUrl=f"{settings.API_V1_STR}/login/access-token"
-)
+# reusable_oauth2 = OAuth2PasswordBearer(
+#     tokenUrl=f"{settings.API_V1_STR}/login/access-token"
+# )
 
 
 def get_db() -> Generator[Session, None, None]:
     with Session(engine) as session:
         yield session
 
-
+api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
 SessionDep = Annotated[Session, Depends(get_db)]
-TokenDep = Annotated[str, Depends(reusable_oauth2)]
-
-
-def get_current_user(session: SessionDep, token: TokenDep) -> User:
-    try:
-        payload = jwt.decode(
-            token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
-        )
-        token_data = TokenPayload(**payload)
-    except (InvalidTokenError, ValidationError):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="Could not validate credentials",
-        )
-    user = session.get(User, token_data.sub)
-    if not user:
-        raise HTTPException(status_code=404, detail="User not found")
-    if not user.is_active:
-        raise HTTPException(status_code=400, detail="Inactive user")
-    return user
-
-
-CurrentUser = Annotated[User, Depends(get_current_user)]
+# TokenDep = Annotated[str, Depends(reusable_oauth2)]
+
+def get_current_user(
+    session: SessionDep,
+    auth_header: str = Security(api_key_header),
+) -> UserOrganization:
+    """Authenticate user via API Key first, fallback to JWT token."""
+
+    if auth_header.startswith("ApiKey "):
+        api_key = auth_header.split(" ", 1)[1]
+        api_key_record = get_api_key_by_value(session, api_key)
+        if not api_key_record:
+            raise HTTPException(status_code=401, detail="Invalid API Key")
+
+        user = session.get(User, api_key_record.user_id)
+        if not user:
+            raise HTTPException(status_code=404, detail="User linked to API Key not found")
+
+        validate_organization(session, api_key_record.organization_id)
+
+        # Return UserOrganization model with organization ID
+        return UserOrganization(**user.model_dump(), organization_id=api_key_record.organization_id)
+
+    if auth_header.startswith("Bearer "):
+        try:
+            token = auth_header.split(" ", 1)[1]
+            payload = jwt.decode(
+                token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+            )
+            token_data = TokenPayload(**payload)
+        except (InvalidTokenError, ValidationError):
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="Could not validate credentials",
+            )
+        user = session.get(User, token_data.sub)
+        if not user:
+            raise HTTPException(status_code=404, detail="User not found")
+        if not user.is_active:
+            raise HTTPException(status_code=400, detail="Inactive user")
+
+        return UserOrganization(**user.model_dump(), organization_id=None)
+    raise HTTPException(status_code=401, detail="Invalid Authorization header format")
+
+CurrentUser = Annotated[UserOrganization, Depends(get_current_user)]
 
 
 def get_current_active_superuser(current_user: CurrentUser) -> User:
@@ -78,6 +102,8 @@ def verify_user_project_organization(
     Verify that the authenticated user is part of the project
     and that the project belongs to the organization.
     """
+    if current_user.organization_id and current_user.organization_id != organization_id:
+        raise HTTPException(status_code=403, detail="User does not belong to the specified organization")
 
     project_organization = db.exec(
         select(Project, Organization)
@@ -105,9 +131,11 @@ def verify_user_project_organization(
         raise HTTPException(status_code=403, detail="Project does not belong to the organization")
 
 
+    current_user.organization_id = organization_id
+    
     # Superuser bypasses all checks
     if current_user.is_superuser:
-        return UserProjectOrg(**current_user.model_dump(), project_id=project_id, organization_id=organization_id)
+        return UserProjectOrg(**current_user.model_dump(), project_id=project_id)
 
     # Check if the user is part of the project
     user_in_project = db.exec(
@@ -121,4 +149,4 @@ def verify_user_project_organization(
     if not user_in_project:
         raise HTTPException(status_code=403, detail="User is not part of the project")
 
-    return UserProjectOrg(**current_user.model_dump(), project_id=project_id, organization_id=organization_id)
+    return UserProjectOrg(**current_user.model_dump(), project_id=project_id)

backend/app/api/routes/api_keys.py

@@ -2,7 +2,7 @@
 from fastapi import APIRouter, Depends, HTTPException
 from sqlmodel import Session
 from app.api.deps import get_db, get_current_active_superuser
-from app.crud.api_key import create_api_key, get_api_key, get_api_keys_by_organization, delete_api_key
+from app.crud.api_key import create_api_key, get_api_key, get_api_keys_by_organization, delete_api_key, get_api_key_by_user_org
 from app.crud.organization import get_organization_by_id, validate_organization
 from app.crud.project_user import is_user_part_of_organization
 from app.models import APIKeyPublic, User
@@ -30,6 +30,10 @@ def create_key(
         if not is_user_part_of_organization(session, user_id, organization_id):
             raise HTTPException(status_code=403, detail="User is not part of any project in the organization")
 
+        existing_api_key = get_api_key_by_user_org(session, organization_id, user_id)
+        if existing_api_key:
+            raise HTTPException(status_code=400, detail="API Key already exists for this user and organization")
+    
         # Create and return API key
         api_key = create_api_key(session, organization_id=organization_id, user_id=user_id)
         return APIResponse.success_response(api_key)

backend/app/crud/api_key.py

@@ -12,7 +12,7 @@ def create_api_key(session: Session, organization_id: uuid.UUID, user_id: uuid.U
     Generates a new API key for an organization and associates it with a user.
     """
     api_key = APIKey(
-        key=secrets.token_urlsafe(32),
+        key='ApiKey '+secrets.token_urlsafe(32),
         organization_id=organization_id,
         user_id=user_id
     )
@@ -63,3 +63,19 @@ def delete_api_key(session: Session, api_key_id: int) -> None:
 
     session.add(api_key)
     session.commit()
+
+def get_api_key_by_value(session: Session, api_key_value: str) -> APIKey | None:
+    """
+    Retrieve an API Key record by its value.
+    """
+    return session.exec(select(APIKey).where(APIKey.key == api_key_value, APIKey.is_deleted == False)).first()
+
+def get_api_key_by_user_org(session: Session, organization_id: int, user_id: str) -> APIKey | None:
+    """
+    Retrieve an API key for a specific user and organization.
+    """
+    statement = select(APIKey).where(
+        APIKey.organization_id == organization_id,
+        APIKey.user_id == user_id
+    )
+    return session.exec(statement).first()

backend/app/models/__init__.py

@@ -40,5 +40,6 @@
     UserUpdateMe,
     NewPassword,
     UpdatePassword,
-    UserProjectOrg
+    UserProjectOrg,
+    UserOrganization
 )

backend/app/models/user.py

@@ -11,12 +11,14 @@ class UserBase(SQLModel):
     is_superuser: bool = False
     full_name: str | None = Field(default=None, max_length=255)
 
+class UserOrganization(UserBase):
+    id: uuid.UUID
+    organization_id: int | None
 
-class UserProjectOrg(UserBase):
-    id: uuid.UUID  # User ID
+class UserProjectOrg(UserOrganization):
     project_id: int
-    organization_id: int
 
+
 # Properties to receive via API on creation
 class UserCreate(UserBase):
     password: str = Field(min_length=8, max_length=40)