46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI
|
|
from starlette.middleware.authentication import AuthenticationMiddleware
|
|
from starlette.requests import Request
|
|
from fastapi.responses import HTMLResponse, PlainTextResponse
|
|
from starlette.responses import JSONResponse
|
|
|
|
from pillar_tool.middleware.basicauth_backend import BasicAuthBackend
|
|
from pillar_tool.util import load_config, config
|
|
from pillar_tool.db import run_db_migrations
|
|
|
|
# load config so everything else can work
|
|
load_config()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def app_lifespan(app: FastAPI):
|
|
run_db_migrations()
|
|
|
|
yield
|
|
|
|
|
|
def on_auth_error(request: Request, exc: Exception):
|
|
response = PlainTextResponse(str(exc), status_code=401)
|
|
response.headers["WWW-Authenticate"] = "Basic"
|
|
|
|
return response
|
|
|
|
def on_general_error(request: Request, exc: Exception):
|
|
print("wtf?")
|
|
response = PlainTextResponse(str(exc), status_code=500)
|
|
|
|
app = FastAPI(lifespan=app_lifespan)
|
|
app.add_middleware(AuthenticationMiddleware, backend=BasicAuthBackend(), on_error=on_auth_error)
|
|
app.exception_handler(Exception)(on_general_error)
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
return {"message": "Hello World"}
|
|
|
|
@app.get("/pillar/{host}")
|
|
async def pillar(host: str):
|
|
return JSONResponse(content=collect_pillar_data(host))
|