Merge pull request #13 from eadwinCode/guard_test

eadwinCode · web-flow · commit 7e8ab42eea79 · 2022-08-15T00:11:40.000+05:30
Guard Module Testing
diff --git a/README.md b/README.md
@@ -107,7 +107,6 @@ You will see the automatic interactive API documentation (provided by <a href="h
 Project is still in development
 - Remaining testing modules:
     - configuration
-    - guard
 - Project CLI scaffolding
 - Documentation
 - Database Plugin with [Encode/ORM](https://github.com/encode/orm)
diff --git a/ellar/core/guard/base.py b/ellar/core/guard/base.py
@@ -3,28 +3,32 @@
 
 from pydantic import BaseModel
 from starlette.exceptions import HTTPException
-from starlette.status import HTTP_403_FORBIDDEN
+from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
 from ellar.core.connection import HTTPConnection
 from ellar.core.context import ExecutionContext
 from ellar.exceptions import APIException
 
 
 class GuardCanActivate(ABC, metaclass=ABCMeta):
-    _exception_class: t.Type[HTTPException] = HTTPException
-    _status_code: int = HTTP_403_FORBIDDEN
-    _detail: str = "Not authenticated"
+    exception_class: t.Type[HTTPException] = HTTPException
+    status_code: int = HTTP_403_FORBIDDEN
+    detail: str = "Not authenticated"
 
     @abstractmethod
     async def can_activate(self, context: ExecutionContext) -> bool:
         pass
 
     def raise_exception(self) -> None:
-        raise self._exception_class(status_code=self._status_code, detail=self._detail)
+        raise self.exception_class(status_code=self.status_code, detail=self.detail)
 
 
 class BaseAuthGuard(GuardCanActivate, ABC, metaclass=ABCMeta):
+    status_code = HTTP_401_UNAUTHORIZED
     openapi_scope: t.List = []
+    openapi_in: t.Optional[str] = None
+    openapi_description: t.Optional[str] = None
+    openapi_name: t.Optional[str] = None
 
     @abstractmethod
     async def handle_request(self, *, connection: HTTPConnection) -> t.Optional[t.Any]:
@@ -40,6 +44,7 @@ async def can_activate(self, context: ExecutionContext) -> bool:
         result = await self.handle_request(connection=connection)
         if result:
             # auth parameter on request
+            connection.scope["user"] = result
             return True
         return False
 
@@ -55,9 +60,8 @@ class HTTPAuthorizationCredentials(BaseModel):
 
 
 class BaseAPIKey(BaseAuthGuard, ABC, metaclass=ABCMeta):
-    openapi_in: t.Optional[str] = None
+    exception_class = APIException
     parameter_name: str = "key"
-    openapi_description: t.Optional[str] = None
 
     def __init__(self) -> None:
         self.name = self.parameter_name
@@ -66,9 +70,7 @@ def __init__(self) -> None:
     async def handle_request(self, connection: HTTPConnection) -> t.Optional[t.Any]:
         key = self._get_key(connection)
         if not key:
-            raise APIException(
-                status_code=HTTP_403_FORBIDDEN, detail="Not authenticated"
-            )
+            self.raise_exception()
         return await self.authenticate(connection, key)
 
     @abstractmethod
@@ -88,12 +90,11 @@ def get_guard_scheme(cls) -> t.Dict:
             "type": "apiKey",
             "description": cls.openapi_description,
             "in": cls.openapi_in,
-            "name": cls.__name__,
+            "name": cls.openapi_name or cls.__name__,
         }
 
 
 class BaseHttpAuth(BaseAuthGuard, ABC, metaclass=ABCMeta):
-    openapi_description: t.Optional[str] = None
     openapi_scheme: t.Optional[str] = None
     realm: t.Optional[str] = None
 
@@ -130,5 +131,5 @@ def get_guard_scheme(cls) -> t.Dict:
             "type": "http",
             "description": cls.openapi_description,
             "scheme": cls.openapi_scheme,
-            "name": cls.__name__,
+            "name": cls.openapi_name or cls.__name__,
         }
diff --git a/ellar/core/guard/http.py b/ellar/core/guard/http.py
@@ -3,15 +3,14 @@
 from abc import ABC
 from base64 import b64decode
 
-from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
-
 from ellar.core.connection import HTTPConnection
 from ellar.exceptions import APIException, AuthenticationFailed
 
 from .base import BaseHttpAuth, HTTPAuthorizationCredentials, HTTPBasicCredentials
 
 
 class HttpBearerAuth(BaseHttpAuth, ABC):
+    exception_class = APIException
     openapi_scheme: str = "bearer"
     openapi_bearer_format: t.Optional[str] = None
     header: str = "Authorization"
@@ -28,18 +27,17 @@ def _get_credentials(
         authorization: str = connection.headers.get(self.header)
         scheme, _, credentials = self.authorization_partitioning(authorization)
         if not (authorization and scheme and credentials):
-            raise APIException(
-                status_code=HTTP_403_FORBIDDEN, detail="Not authenticated"
-            )
-        if scheme.lower() != self.openapi_scheme:
-            raise APIException(
-                status_code=HTTP_403_FORBIDDEN,
+            self.raise_exception()
+        if scheme and str(scheme).lower() != self.openapi_scheme:
+            raise self.exception_class(
+                status_code=self.status_code,
                 detail="Invalid authentication credentials",
             )
         return HTTPAuthorizationCredentials(scheme=scheme, credentials=credentials)
 
 
 class HttpBasicAuth(BaseHttpAuth, ABC):
+    exception_class = APIException
     openapi_scheme: str = "basic"
     realm: t.Optional[str] = None
     header = "Authorization"
@@ -50,33 +48,39 @@ def _not_unauthorized_exception(self, message: str) -> None:
         else:
             unauthorized_headers = {"WWW-Authenticate": "Basic"}
         raise AuthenticationFailed(
-            status_code=HTTP_401_UNAUTHORIZED,
+            status_code=self.status_code,
             detail=message,
             headers=unauthorized_headers,
         )
 
     def _get_credentials(self, connection: HTTPConnection) -> HTTPBasicCredentials:
         authorization: str = connection.headers.get(self.header)
-        scheme, _, credentials = self.authorization_partitioning(authorization)
+        parts = authorization.split(" ") if authorization else []
+        scheme, credentials = str(), str()
+
+        if len(parts) == 1:
+            credentials = parts[0]
+            scheme = "basic"
+        elif len(parts) == 2:
+            credentials = parts[1]
+            scheme = parts[0].lower()
 
         if (
             not (authorization and scheme and credentials)
             or scheme.lower() != self.openapi_scheme
         ):
-            raise APIException(
-                status_code=HTTP_403_FORBIDDEN, detail="Not authenticated"
-            )
+            self.raise_exception()
+
         data: t.Optional[t.Union[str, bytes]] = None
         try:
             data = b64decode(credentials).decode("ascii")
         except (ValueError, UnicodeDecodeError, binascii.Error):
             self._not_unauthorized_exception("Invalid authentication credentials")
 
         username, separator, password = (
-            str(data).partition(":") if data else None,
-            None,
-            None,
+            str(data).partition(":") if data else (None, None, None)
         )
+
         if not separator:
             self._not_unauthorized_exception("Invalid authentication credentials")
         return HTTPBasicCredentials(username=username, password=password)
diff --git a/tests/notstarted/test_guard/__init__.py b/tests/notstarted/test_guard/__init__.py
diff --git a/tests/test_guard.py b/tests/test_guard.py
@@ -0,0 +1,210 @@
+import pytest
+from starlette.status import HTTP_401_UNAUTHORIZED
+
+from ellar.common import Req, get, guards
+from ellar.core import AppFactory, TestClient
+from ellar.core.guard import (
+    APIKeyCookie,
+    APIKeyHeader,
+    APIKeyQuery,
+    HttpBasicAuth,
+    HttpBearerAuth,
+    HttpDigestAuth,
+)
+from ellar.exceptions import APIException
+from ellar.openapi import OpenAPIDocumentBuilder
+from ellar.serializer import serialize_object
+
+
+class CustomException(APIException):
+    pass
+
+
+class QuerySecretKey(APIKeyQuery):
+    async def authenticate(self, connection, key):
+        if key == "querysecretkey":
+            return key
+
+
+class HeaderSecretKey(APIKeyHeader):
+    async def authenticate(self, connection, key):
+        if key == "headersecretkey":
+            return key
+
+
+class HeaderSecretKeyCustomException(HeaderSecretKey):
+    exception_class = CustomException
+
+
+class CookieSecretKey(APIKeyCookie):
+    openapi_name = "API Key Auth"
+
+    async def authenticate(self, connection, key):
+        if key == "cookiesecretkey":
+            return key
+
+
+class BasicAuth(HttpBasicAuth):
+    openapi_name = "API Authentication"
+
+    async def authenticate(self, connection, credentials):
+        if credentials.username == "admin" and credentials.password == "secret":
+            return credentials.username
+
+
+class BearerAuth(HttpBearerAuth):
+    openapi_name = "JWT Authentication"
+
+    async def authenticate(self, connection, credentials):
+        if credentials.credentials == "bearertoken":
+            return credentials.credentials
+
+
+class DigestAuth(HttpDigestAuth):
+    async def authenticate(self, connection, credentials):
+        if credentials.credentials == "digesttoken":
+            return credentials.credentials
+
+
+app = AppFactory.create_app()
+
+
+for _path, auth in [
+    ("apikeyquery", QuerySecretKey()),
+    ("apikeyheader", HeaderSecretKey()),
+    ("apikeycookie", CookieSecretKey()),
+    ("basic", BasicAuth()),
+    ("bearer", BearerAuth()),
+    ("digest", DigestAuth()),
+    ("customexception", HeaderSecretKeyCustomException()),
+]:
+
+    @get(f"/{_path}")
+    @guards(auth)
+    def auth_demo_endpoint(request: Req()):
+        return {"authentication": request.user}
+
+    app.router.append(auth_demo_endpoint)
+
+client = TestClient(app)
+
+BODY_UNAUTHORIZED_DEFAULT = {"detail": "Not authenticated"}
+
+
+@pytest.mark.parametrize(
+    "path,kwargs,expected_code,expected_body",
+    [
+        ("/apikeyquery", {}, HTTP_401_UNAUTHORIZED, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/apikeyquery?key=querysecretkey",
+            {},
+            200,
+            dict(authentication="querysecretkey"),
+        ),
+        ("/apikeyheader", {}, HTTP_401_UNAUTHORIZED, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/apikeyheader",
+            dict(headers={"key": "headersecretkey"}),
+            200,
+            dict(authentication="headersecretkey"),
+        ),
+        ("/apikeycookie", {}, HTTP_401_UNAUTHORIZED, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/apikeycookie",
+            dict(cookies={"key": "cookiesecretkey"}),
+            200,
+            dict(authentication="cookiesecretkey"),
+        ),
+        ("/basic", {}, HTTP_401_UNAUTHORIZED, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/basic",
+            dict(headers={"Authorization": "Basic YWRtaW46c2VjcmV0"}),
+            200,
+            dict(authentication="admin"),
+        ),
+        (
+            "/basic",
+            dict(headers={"Authorization": "YWRtaW46c2VjcmV0"}),
+            200,
+            dict(authentication="admin"),
+        ),
+        (
+            "/basic",
+            dict(headers={"Authorization": "Basic invalid"}),
+            HTTP_401_UNAUTHORIZED,
+            {"detail": "Invalid authentication credentials"},
+        ),
+        (
+            "/basic",
+            dict(headers={"Authorization": "some invalid value"}),
+            HTTP_401_UNAUTHORIZED,
+            BODY_UNAUTHORIZED_DEFAULT,
+        ),
+        ("/bearer", {}, 401, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/bearer",
+            dict(headers={"Authorization": "Bearer bearertoken"}),
+            200,
+            dict(authentication="bearertoken"),
+        ),
+        (
+            "/bearer",
+            dict(headers={"Authorization": "Invalid bearertoken"}),
+            HTTP_401_UNAUTHORIZED,
+            {"detail": "Invalid authentication credentials"},
+        ),
+        ("/digest", {}, 401, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/digest",
+            dict(headers={"Authorization": "Digest digesttoken"}),
+            200,
+            dict(authentication="digesttoken"),
+        ),
+        (
+            "/digest",
+            dict(headers={"Authorization": "Invalid digesttoken"}),
+            HTTP_401_UNAUTHORIZED,
+            {"detail": "Invalid authentication credentials"},
+        ),
+        ("/customexception", {}, HTTP_401_UNAUTHORIZED, BODY_UNAUTHORIZED_DEFAULT),
+        (
+            "/customexception",
+            dict(headers={"key": "headersecretkey"}),
+            200,
+            dict(authentication="headersecretkey"),
+        ),
+    ],
+)
+def test_auth(path, kwargs, expected_code, expected_body):
+    response = client.get(path, **kwargs)
+    assert response.status_code == expected_code
+    assert response.json() == expected_body
+
+
+def test_auth_schema():
+    document = serialize_object(OpenAPIDocumentBuilder().build_document(app))
+    assert document["components"]["securitySchemes"] == {
+        "API Key Auth": {"type": "apiKey", "in": "cookie", "name": "API Key Auth"},
+        "HeaderSecretKey": {
+            "type": "apiKey",
+            "in": "header",
+            "name": "HeaderSecretKey",
+        },
+        "QuerySecretKey": {"type": "apiKey", "in": "query", "name": "QuerySecretKey"},
+        "API Authentication": {
+            "type": "http",
+            "scheme": "basic",
+            "name": "API Authentication",
+        },
+        "JWT Authentication": {
+            "type": "http",
+            "scheme": "bearer",
+            "name": "JWT Authentication",
+        },
+        "HeaderSecretKeyCustomException": {
+            "type": "apiKey",
+            "in": "header",
+            "name": "HeaderSecretKeyCustomException",
+        },
+        "DigestAuth": {"type": "http", "scheme": "digest", "name": "DigestAuth"},
+    }
diff --git a/tests/test_openapi/test_module.py b/tests/test_openapi/test_module.py
