-
-
Save arnydo/4fea283b6e0db7228cce1a5313b90a05 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from fastapi import Security, Depends, FastAPI, HTTPException | |
from fastapi.security.api_key import APIKeyQuery, APIKeyCookie, APIKeyHeader, APIKey | |
from fastapi.openapi.docs import get_swagger_ui_html | |
from fastapi.openapi.utils import get_openapi | |
from starlette.status import HTTP_403_FORBIDDEN | |
from starlette.responses import RedirectResponse, JSONResponse | |
API_KEY = "1234567asdfgh" | |
API_KEY_NAME = "access_token" | |
COOKIE_DOMAIN = "localtest.me" | |
api_key_query = APIKeyQuery(name=API_KEY_NAME, auto_error=False) | |
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False) | |
api_key_cookie = APIKeyCookie(name=API_KEY_NAME, auto_error=False) | |
async def get_api_key( | |
api_key_query: str = Security(api_key_query), | |
api_key_header: str = Security(api_key_header), | |
api_key_cookie: str = Security(api_key_cookie), | |
): | |
if api_key_query == API_KEY: | |
return api_key_query | |
elif api_key_header == API_KEY: | |
return api_key_header | |
elif api_key_cookie == API_KEY: | |
return api_key_cookie | |
else: | |
raise HTTPException( | |
status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials" | |
) | |
app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) | |
@app.get("/") | |
async def homepage(): | |
return "Welcome to the security test!" | |
@app.get("/logout") | |
async def route_logout_and_remove_cookie(): | |
response = RedirectResponse(url="/") | |
response.delete_cookie(API_KEY_NAME, domain=COOKIE_DOMAIN) | |
return response | |
@app.get("/openapi.json", tags=["documentation"]) | |
async def get_open_api_endpoint(api_key: APIKey = Depends(get_api_key)): | |
response = JSONResponse( | |
get_openapi(title="FastAPI security test", version=1, routes=app.routes) | |
) | |
return response | |
@app.get("/documentation", tags=["documentation"]) | |
async def get_documentation(api_key: APIKey = Depends(get_api_key)): | |
response = get_swagger_ui_html(openapi_url="/openapi.json", title="docs") | |
response.set_cookie( | |
API_KEY_NAME, | |
value=api_key, | |
domain=COOKIE_DOMAIN, | |
httponly=True, | |
max_age=1800, | |
expires=1800, | |
) | |
return response | |
@app.get("/secure_endpoint", tags=["test"]) | |
async def get_open_api_endpoint(api_key: APIKey = Depends(get_api_key)): | |
response = "How cool is this?" | |
return response |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Exactly the example I was searching for. Thank you!