Refactor

starlette_apitally/__init__.py

@@ -1,5 +1 @@
-from starlette_apitally.middleware import ApitallyMiddleware
-
-
-__all__ = ["ApitallyMiddleware"]
 __version__ = "0.0.0"

starlette_apitally/client.py

@@ -9,9 +9,7 @@
 
 import backoff
 import httpx
-from starlette.types import ASGIApp
 
-from starlette_apitally.app_info import get_app_info
 from starlette_apitally.keys import Keys
 from starlette_apitally.requests import Requests
 
@@ -59,6 +57,12 @@ def __init__(self, client_id: str, env: str, enable_keys: bool = False, send_eve
         self._stop_sync_loop = False
         self.start_sync_loop()
 
+    @classmethod
+    def get_instance(cls) -> ApitallyClient:
+        if cls._instance is None:
+            raise RuntimeError("Apitally client not initialized")
+        return cls._instance
+
     def get_http_client(self) -> httpx.AsyncClient:
         base_url = f"{HUB_BASE_URL}/{HUB_VERSION}/{self.client_id}/{self.env}"
         return httpx.AsyncClient(base_url=base_url)
@@ -83,8 +87,7 @@ async def _run_sync_loop(self) -> None:
     def stop_sync_loop(self) -> None:
         self._stop_sync_loop = True
 
-    def send_app_info(self, app: ASGIApp, app_version: Optional[str], openapi_url: Optional[str]) -> None:
-        app_info = get_app_info(app, app_version, openapi_url)
+    def send_app_info(self, app_info: Dict[str, Any]) -> None:
         payload = {
             "instance_uuid": self.instance_uuid,
             "message_uuid": str(uuid4()),

starlette_apitally/fastapi.py

@@ -3,11 +3,17 @@
 from fastapi.exceptions import HTTPException
 from fastapi.openapi.models import APIKey, APIKeyIn
 from fastapi.requests import Request
+from fastapi.security import SecurityScopes
 from fastapi.security.base import SecurityBase
 from fastapi.security.utils import get_authorization_scheme_param
 from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
 
-from starlette_apitally.keys import Key, Keys
+from starlette_apitally.client import ApitallyClient
+from starlette_apitally.keys import Key
+from starlette_apitally.starlette import ApitallyMiddleware
+
+
+__all__ = ["ApitallyMiddleware", "Key", "api_key_auth"]
 
 
 class AuthorizationAPIKeyHeader(SecurityBase):
@@ -20,7 +26,7 @@ def __init__(self, *, auto_error: bool = True):
         self.scheme_name = "Authorization header with ApiKey scheme"
         self.auto_error = auto_error
 
-    async def __call__(self, request: Request) -> Optional[Key]:
+    async def __call__(self, request: Request, security_scopes: SecurityScopes) -> Optional[Key]:
         authorization = request.headers.get("Authorization")
         scheme, param = get_authorization_scheme_param(authorization)
         if not authorization or scheme.lower() != "apikey":
@@ -31,23 +37,19 @@ async def __call__(self, request: Request) -> Optional[Key]:
                     headers={"WWW-Authenticate": "ApiKey"},
                 )
             else:
-                return None
-        keys = self._get_keys()
-        key = keys.get(param)
+                return None  # pragma: no cover
+        key = ApitallyClient.get_instance().keys.get(param)
         if key is None and self.auto_error:
             raise HTTPException(
                 status_code=HTTP_403_FORBIDDEN,
                 detail="Invalid API key",
             )
+        if key is not None and self.auto_error and not key.check_scopes(security_scopes.scopes):
+            raise HTTPException(
+                status_code=HTTP_403_FORBIDDEN,
+                detail="Permission denied",
+            )
         return key
 
-    def _get_keys(self) -> Keys:
-        from starlette_apitally.client import ApitallyClient
-
-        client = ApitallyClient._instance
-        if client is None:
-            raise RuntimeError("ApitallyClient not initialized")
-        return client.keys
-
 
 api_key_auth = AuthorizationAPIKeyHeader()

starlette_apitally/keys.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from hashlib import scrypt
 from typing import Any, Dict, List, Optional, Set
@@ -9,16 +9,23 @@
 @dataclass(frozen=True)
 class Key:
     key_id: int
+    name: str = ""
+    scopes: List[str] = field(default_factory=list)
     expires_at: Optional[datetime] = None
 
     @property
     def is_expired(self) -> bool:
         return self.expires_at is not None and self.expires_at < datetime.now()
 
+    def check_scopes(self, scopes: List[str]) -> bool:
+        return all(scope in self.scopes for scope in scopes)
+
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> Key:
         return cls(
             key_id=data["key_id"],
+            name=data.get("name", ""),
+            scopes=data.get("scopes", []),
             expires_at=(
                 datetime.now() + timedelta(seconds=data["expires_in_seconds"])
                 if data["expires_in_seconds"] is not None