From d96011c522b172049ad5985d55383f27d75275bd Mon Sep 17 00:00:00 2001
From: Shekar Siri <sshekarsiri@gmail.com>
Date: Tue, 7 Nov 2023 16:12:00 +0100
Subject: [PATCH] change(assist-stats): removed path prefix, should use the
 root_path

---
 ee/assist-stats/auth.py | 44 ++++++++++++++++++++++-------------------
 1 file changed, 24 insertions(+), 20 deletions(-)

diff --git a/ee/assist-stats/auth.py b/ee/assist-stats/auth.py
index ec440be58..8c5507303 100644
--- a/ee/assist-stats/auth.py
+++ b/ee/assist-stats/auth.py
@@ -1,8 +1,9 @@
+from fastapi import HTTPException, Depends, status, Security
 from fastapi.security import OAuth2PasswordBearer
-from fastapi import HTTPException, Depends, status
 from decouple import config
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+# Instantiate OAuth2PasswordBearer with automatic error responses disabled
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token", auto_error=False)
 
 
 class AuthHandler:
@@ -10,30 +11,33 @@ class AuthHandler:
         """
         Authorization method using an API key.
         """
-        # Attempt to get the ACCESS_TOKEN, if not set, default to an empty list
-        api_key = config("ACCESS_TOKEN", default=None)
-        self.__api_keys = [api_key] if api_key else []
+        # Attempt to get the ACCESS_TOKEN, if not set, default to None
+        self.api_key = config("ACCESS_TOKEN", default=None)
 
-    def __contains__(self, api_key):
-        # Skip the check if no API keys are configured
-        if not self.__api_keys:
-            return True
-
-        return api_key in self.__api_keys
-
-    def add_key(self, key):
-        """Adds new key for authentication."""
-        if key:  # Ensure we don't add empty keys
-            self.__api_keys.append(key)
+    def verify_api_key(self, api_key: str):
+        return api_key == self.api_key
 
 
-auth_method = AuthHandler()
+auth_handler = AuthHandler()
 
 
-def api_key_auth(api_key: str = Depends(oauth2_scheme)):
-    """Method to verify auth."""
-    if api_key not in auth_method:
+async def api_key_auth(api_key: str = Security(oauth2_scheme)):
+    # If ACCESS_TOKEN is not configured, skip the authorization check
+    if not auth_handler.api_key:
+        return True
+
+    # If the Authorization header is not provided, raise an HTTP 403 Forbidden error
+    if not api_key:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail="Not authenticated"
+        )
+
+    # If the provided API key is invalid, raise an HTTP 401 Unauthorized error
+    if not auth_handler.verify_api_key(api_key):
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Forbidden"
         )
+    # If the API key is valid, continue processing the request
+    return True