add api key auth

api/api/routers/admin.py

@@ -3,7 +3,7 @@
 import logging
 import os
 
-from .auth import key_auth
+from .auth import admin_key_auth
 from ..src.repos.mongo.client import get_db
 
 logger = logging.getLogger(__name__)
@@ -12,7 +12,7 @@
 
 
 @router.get("/logs/{tail}", include_in_schema=False)
-def logs(isAdmin: bool = Depends(key_auth), tail: int = 10):
+def logs(isAdmin: bool = Depends(admin_key_auth), tail: int = 10):
     try:
         # return last lines of log file
         with open("/tmp/eotdl-api.log", "r") as f:
@@ -24,7 +24,7 @@ def logs(isAdmin: bool = Depends(key_auth), tail: int = 10):
 
 
 @router.get("/env", include_in_schema=False)
-def env_vars(isAdmin: bool = Depends(key_auth)):
+def env_vars(isAdmin: bool = Depends(admin_key_auth)):
     try:
         return os.environ
     except Exception as e:
@@ -33,7 +33,7 @@ def env_vars(isAdmin: bool = Depends(key_auth)):
 
 
 @router.get("/init-db", include_in_schema=False)
-def initialize_db(isAdmin: bool = Depends(key_auth)):
+def initialize_db(isAdmin: bool = Depends(admin_key_auth)):
     try:
         db = get_db()
         collections = db.list_collection_names()

api/api/routers/auth/__init__.py

@@ -1 +1 @@
-from .main import get_current_user, key_auth
+from .main import get_current_user, key_auth, admin_key_auth

api/api/routers/auth/main.py

@@ -4,44 +4,42 @@
 import logging
 import os
 
-from ...src.usecases.user import persist_user
+from ...src.usecases.user import persist_user, retrieve_user_by_key
 from ...src.usecases.auth import parse_token
 
 logger = logging.getLogger(__name__)
 
 token_auth_scheme = HTTPBearer(auto_error=False)
 api_key_auth_scheme = APIKeyHeader(name="X-API-Key", auto_error=False)
 
-def key_auth(api_key: str = Depends(api_key_auth_scheme)):
+
+def admin_key_auth(api_key: str = Depends(api_key_auth_scheme)):
     if not api_key or api_key != os.environ.get("ADMIN_API_KEY"):
         raise HTTPException(status_code=401, detail="Invalid API key")
     return
 
 
-# def token_auth(token: str = Depends(token_auth_scheme)):
-def get_current_user(token: str = Depends(token_auth_scheme)):
+def key_auth(api_key: str = Depends(api_key_auth_scheme)):
+    if not api_key:
+        return None
+    try:
+        return retrieve_user_by_key(api_key).model_dump()
+    except Exception:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+
+def token_auth(token: str = Depends(token_auth_scheme)):
     if not token:
         return None
     try:
-        data = parse_token(token.credentials)
-        return persist_user(data)
+        return parse_token(token.credentials)
     except Exception as e:
         logger.exception("get_current_user")
         raise HTTPException(status_code=401, detail="Invalid token")
 
 
-# async def get_current_user(api_key_user=Depends(key_auth), token_user=Depends(token_auth)):
-#     if not token_user:
-#         raise HTTPException(status_code=401, detail="Not authenticated")
-#     return token_user
-
-
-
-
-
-
-
-
-
-
-
+def get_current_user(api_key_user=Depends(key_auth), token: str = Depends(token_auth)):
+    if not (api_key_user or token):
+        raise HTTPException(status_code=401, detail="Not authenticated")
+    data = api_key_user or token
+    return persist_user(data)

api/api/routers/datasets/delete_dataset.py

@@ -2,7 +2,7 @@
 from fastapi import APIRouter, status, Depends
 import logging
 
-from ..auth import key_auth
+from ..auth import admin_key_auth
 from ...src.usecases.datasets import delete_dataset
 
 router = APIRouter()
@@ -12,7 +12,7 @@
 @router.delete("/{name}", include_in_schema=False)
 def delete(
     name: str,
-    isAdmin: bool = Depends(key_auth),  # only admin can delete datasets
+    isAdmin: bool = Depends(admin_key_auth),  # only admin can delete datasets
 ):
     try:
         message = delete_dataset(name)

api/api/routers/migrate.py

@@ -4,7 +4,7 @@
 from minio.commonconfig import CopySource
 from boto3.s3.transfer import TransferConfig
 
-from .auth import key_auth
+from .auth import admin_key_auth
 from ..src.repos.mongo.client import get_db
 from ..src.repos.minio.client import get_client
 from ..src.repos.boto3.client import get_client as get_boto3_client
@@ -19,7 +19,7 @@
 
 
 @router.get("", include_in_schema=False)
-def migrate_db(isAdmin: bool = Depends(key_auth)):
+def migrate_db(isAdmin: bool = Depends(admin_key_auth)):
     # return "Done"
     db = get_db()
     collections = db.list_collection_names()

api/api/routers/models/delete_model.py

@@ -2,7 +2,7 @@
 from fastapi import APIRouter, status, Depends
 import logging
 
-from ..auth import key_auth
+from ..auth import admin_key_auth
 from ...src.usecases.models import delete_model
 
 router = APIRouter()
@@ -12,7 +12,7 @@
 @router.delete("/{name}", include_in_schema=False)
 def delete(
     name: str,
-    isAdmin: bool = Depends(key_auth),  # only admin can delete models
+    isAdmin: bool = Depends(admin_key_auth),  # only admin can delete models
 ):
     try:
         message = delete_model(name)

api/api/src/errors/__init__.py

@@ -21,5 +21,6 @@
     TierLimitError,
     UserDoesNotExistError,
     UserAlreadyExistsError,
+    InvalidApiKey,
 )
 from .tags import InvalidTagError

api/api/src/errors/user.py

@@ -4,14 +4,22 @@ class UserUnauthorizedError(Exception):
     def __init__(self):
         super().__init__(self.message)
 
+
 class TierLimitError(Exception):
     def __init__(self, message):
         super().__init__(message)
 
+
 class UserDoesNotExistError(Exception):
     def __init__(self):
         super().__init__("User does not exist")
 
+
 class UserAlreadyExistsError(Exception):
     def __init__(self):
-        super().__init__("User already exists")
+        super().__init__("User already exists")
+
+
+class InvalidApiKey(Exception):
+    def __init__(self):
+        super().__init__("Invalid API key")

api/api/src/repos/mongo/MongoUserRepo.py

@@ -8,6 +8,9 @@ def __init__(self):
     def retrieve_user_by_uid(self, uid):
         return self.retrieve("users", uid, "uid")
 
+    def retrieve_user_by_key(self, key):
+        return self.retrieve("keys", key)
+
     def update_user(self, id, data):
         return self.update("users", id, data)
 

api/api/src/usecases/user/__init__.py

@@ -6,3 +6,4 @@
 from .check_user_can_create_dataset import check_user_can_create_dataset
 from .check_user_can_download_dataset import check_user_can_download_dataset
 from .check_user_can_create_model import check_user_can_create_model
+from .retrieve_user_by_key import retrieve_user_by_key

api/api/src/usecases/user/retrieve_user_by_key.py

@@ -0,0 +1,14 @@
+from ...repos import UserDBRepo
+from ...errors import UserDoesNotExistError, InvalidApiKey
+from ...models import User
+
+
+def retrieve_user_by_key(api_key: str) -> User:
+    repo = UserDBRepo()
+    data = repo.retrieve_user_by_key(api_key)
+    if data is None:
+        raise InvalidApiKey()
+    user = repo.retrieve_user_by_uid(data["uid"])
+    if user is None:
+        raise UserDoesNotExistError()
+    return User(**user)

api/tests/e2e/datasets/_test_delete_datasets.py

@@ -6,7 +6,7 @@
 from fastapi.testclient import TestClient
 from api.main import app
 
-from api.routers.auth import get_current_user, key_auth
+from api.routers.auth import get_current_user, admin_key_auth
 from api.src.models import User
 from tests.e2e.setup import users, db, s3, boto3, bucket
 
@@ -19,12 +19,12 @@ def get_current_user_mock():
     return User(**users[0])
 
 
-def key_auth_mock():
+def admin_key_auth_mock():
     return True
 
 
 app.dependency_overrides[get_current_user] = get_current_user_mock
-app.dependency_overrides[key_auth] = key_auth_mock
+app.dependency_overrides[admin_key_auth] = admin_key_auth_mock
 
 
 @pytest.fixture

api/tests/e2e/datasets/_test_download_datasets.py

@@ -6,7 +6,7 @@
 from fastapi.testclient import TestClient
 from api.main import app
 
-from api.routers.auth import get_current_user, key_auth
+from api.routers.auth import get_current_user, admin_key_auth
 from api.src.models import User
 from ..setup import users, db, s3, boto3
 
@@ -19,12 +19,12 @@ def get_current_user_mock():
     return User(**users[0])
 
 
-def key_auth_mock():
+def admin_key_auth_mock():
     return True
 
 
 app.dependency_overrides[get_current_user] = get_current_user_mock
-app.dependency_overrides[key_auth] = key_auth_mock
+app.dependency_overrides[admin_key_auth] = admin_key_auth_mock
 
 
 @pytest.fixture

api/tests/e2e/datasets/_test_ingest_datasets.py

@@ -6,7 +6,7 @@
 from fastapi.testclient import TestClient
 from api.main import app
 
-from api.routers.auth import get_current_user, key_auth
+from api.routers.auth import get_current_user, admin_key_auth
 from api.src.models import User
 from ..setup import users, db, s3, boto3
 
@@ -20,12 +20,12 @@ def get_current_user_mock():
     return User(**users[0])
 
 
-def key_auth_mock():
+def admin_key_auth_mock():
     return True
 
 
 app.dependency_overrides[get_current_user] = get_current_user_mock
-app.dependency_overrides[key_auth] = key_auth_mock
+app.dependency_overrides[admin_key_auth] = admin_key_auth_mock
 
 
 @pytest.fixture

api/tests/e2e/datasets/_test_ingest_large_datasets.py

@@ -6,7 +6,7 @@
 from fastapi.testclient import TestClient
 from api.main import app
 
-from api.routers.auth import get_current_user, key_auth
+from api.routers.auth import get_current_user, admin_key_auth
 from api.src.models import User
 from ..setup import users, db, s3, boto3
 
@@ -21,12 +21,12 @@ def get_current_user_mock():
     return User(**users[0])
 
 
-def key_auth_mock():
+def admin_key_auth_mock():
     return True
 
 
 app.dependency_overrides[get_current_user] = get_current_user_mock
-app.dependency_overrides[key_auth] = key_auth_mock
+app.dependency_overrides[admin_key_auth] = admin_key_auth_mock
 
 
 @pytest.fixture

api/tests/e2e/datasets/_test_update_datasets.py

@@ -6,7 +6,7 @@
 from fastapi.testclient import TestClient
 from api.main import app
 
-from api.routers.auth import get_current_user, key_auth
+from api.routers.auth import get_current_user, admin_key_auth
 from api.src.models import User
 from ..setup import users, db, s3, boto3
 
@@ -19,12 +19,12 @@ def get_current_user_mock():
     return User(**users[0])
 
 
-def key_auth_mock():
+def admin_key_auth_mock():
     return True
 
 
 app.dependency_overrides[get_current_user] = get_current_user_mock
-app.dependency_overrides[key_auth] = key_auth_mock
+app.dependency_overrides[admin_key_auth] = admin_key_auth_mock
 
 
 @pytest.fixture

api/tests/e2e/datasets/test_retrieve_datasets.py

@@ -6,7 +6,7 @@
 from fastapi.testclient import TestClient
 from api.main import app
 
-from api.routers.auth import get_current_user, key_auth
+from api.routers.auth import get_current_user, admin_key_auth
 from api.src.models import User
 from ..setup import users, db, s3, boto3
 
@@ -19,12 +19,12 @@ def get_current_user_mock():
     return User(**users[0])
 
 
-def key_auth_mock():
+def admin_key_auth_mock():
     return True
 
 
 app.dependency_overrides[get_current_user] = get_current_user_mock
-app.dependency_overrides[key_auth] = key_auth_mock
+app.dependency_overrides[admin_key_auth] = admin_key_auth_mock
 
 
 @pytest.fixture