File size: 2,939 Bytes
870ab6b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
from typing import Optional
from fastapi.openapi.models import APIKey, APIKeyIn
from fastapi.security.base import SecurityBase
from starlette.exceptions import HTTPException
from starlette.requests import Request
from starlette.status import HTTP_403_FORBIDDEN
class APIKeyBase(SecurityBase):
pass
class APIKeyQuery(APIKeyBase):
def __init__(
self,
*,
name: str,
scheme_name: Optional[str] = None,
description: Optional[str] = None,
auto_error: bool = True,
):
self.model: APIKey = APIKey(
**{"in": APIKeyIn.query}, # type: ignore[arg-type]
name=name,
description=description,
)
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
async def __call__(self, request: Request) -> Optional[str]:
api_key = request.query_params.get(self.model.name)
if not api_key:
if self.auto_error:
raise HTTPException(
status_code=HTTP_403_FORBIDDEN, detail="Not authenticated"
)
else:
return None
return api_key
class APIKeyHeader(APIKeyBase):
def __init__(
self,
*,
name: str,
scheme_name: Optional[str] = None,
description: Optional[str] = None,
auto_error: bool = True,
):
self.model: APIKey = APIKey(
**{"in": APIKeyIn.header}, # type: ignore[arg-type]
name=name,
description=description,
)
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
async def __call__(self, request: Request) -> Optional[str]:
api_key = request.headers.get(self.model.name)
if not api_key:
if self.auto_error:
raise HTTPException(
status_code=HTTP_403_FORBIDDEN, detail="Not authenticated"
)
else:
return None
return api_key
class APIKeyCookie(APIKeyBase):
def __init__(
self,
*,
name: str,
scheme_name: Optional[str] = None,
description: Optional[str] = None,
auto_error: bool = True,
):
self.model: APIKey = APIKey(
**{"in": APIKeyIn.cookie}, # type: ignore[arg-type]
name=name,
description=description,
)
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
async def __call__(self, request: Request) -> Optional[str]:
api_key = request.cookies.get(self.model.name)
if not api_key:
if self.auto_error:
raise HTTPException(
status_code=HTTP_403_FORBIDDEN, detail="Not authenticated"
)
else:
return None
return api_key
|