import ssl
import time
import logging
import asyncio
from asyncio.exceptions import CancelledError
import anyio
import starlette
import uvicorn
import fastapi
from starlette.responses import JSONResponse
from starlette.websockets import WebSocket, WebSocketDisconnect
from fastapi import FastAPI, Request, Response
from fastapi.exceptions import HTTPException
from fastapi.responses import RedirectResponse
from fastapi.encoders import jsonable_encoder
from modules.logger import log
import modules.errors as errors
from modules.api.validate import validate_request, validate_log
errors.install()
def validate_subpath(endpoint: str, subpath: str | None):
if (subpath is not None) and (len(subpath) > 0) and (not endpoint.startswith(subpath)):
if not subpath.endswith('/'):
subpath = f'{subpath}/'
if endpoint.startswith('/'):
url = f'{subpath}{endpoint[1:]}'
else:
url = f'{subpath}{endpoint}'
log.trace(f'API: redirect subpath={subpath} url="{endpoint}" redirect="{url}"')
return RedirectResponse(url=url, status_code=308)
return None
class LoopInstrumentorMiddleware:
def __init__(self, app):
self.app = app
self.instrumented = False
async def __call__(self, scope, receive, send):
if not self.instrumented:
loop = asyncio.get_running_loop()
def verbose_task_factory(loop, coro, context=None):
coro_name = getattr(coro, '__qualname__', str(coro))
frame = getattr(coro, 'cr_frame', None)
origin = f"{frame.f_code.co_filename}:{frame.f_lineno}" if frame else "unknown"
log.trace(f"HTTP: coro={coro_name} fn={origin}")
if context is not None:
return asyncio.Task(coro, loop=loop, name=coro_name, context=context)
return asyncio.Task(coro, loop=loop, name=coro_name)
loop.set_task_factory(verbose_task_factory)
self.instrumented = True
await self.app(scope, receive, send)
def setup_logging(debug: bool = False):
level = logging.DEBUG if debug else logging.WARNING
logging.getLogger("httpcore").setLevel(level)
logging.getLogger("httpx").setLevel(level)
logging.getLogger("uvicorn.access").setLevel(level)
logging.getLogger("asyncio").setLevel(level)
if not debug:
logging.getLogger("uvicorn.error").disabled = True
if debug:
asyncio_logger = logging.getLogger("asyncio")
if not asyncio_logger.handlers:
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("[asyncio] %(message)s"))
asyncio_logger.addHandler(handler)
def setup_middleware(app: FastAPI, cmd_opts):
ssl._create_default_https_context = ssl._create_unverified_context
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']
app.middleware_stack = None
app.add_middleware(GZipMiddleware, minimum_size=2048)
if cmd_opts.profile:
app.add_middleware(LoopInstrumentorMiddleware)
if cmd_opts.cors_origins and cmd_opts.cors_regex:
app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_origins.split(','), allow_origin_regex=cmd_opts.cors_regex, allow_methods=['*'], allow_credentials=True, allow_headers=['*'])
elif cmd_opts.cors_origins:
app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_origins.split(','), allow_methods=['*'], allow_credentials=True, allow_headers=['*'])
elif cmd_opts.cors_regex:
app.add_middleware(CORSMiddleware, allow_origin_regex=cmd_opts.cors_regex, allow_methods=['*'], allow_credentials=True, allow_headers=['*'])
@app.middleware("http")
async def api_preprocess(req: Request, call_next):
try:
ts = time.time()
res: Response = await call_next(req)
duration = str(round(time.time() - ts, 4))
res.headers["X-Process-Time"] = duration
endpoint = req.scope.get('path', 'err')
client = req.scope.get('client', ('0:0.0.0', 0))[0]
token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure")
redirect = validate_subpath(endpoint, cmd_opts.subpath)
if redirect:
return redirect
validate_request(client, endpoint)
if cmd_opts.api_log:
if not validate_log(client, endpoint):
return res
log.info('API user={user} code={code} {prot}/{ver} {method} {endpoint} {client} {duration}'.format(
user = app.tokens.get(token) if hasattr(app, 'tokens') else None,
code = res.status_code,
ver = req.scope.get('http_version', '0.0'),
client = client,
prot = req.scope.get('scheme', 'err'),
method = req.scope.get('method', 'err'),
endpoint = endpoint,
duration = duration,
))
return res
except CancelledError:
log.warning('WebSocket closed')
except BaseException as e:
return handle_exception(req, e)
def handle_exception(req: Request, e: Exception):
err = {
"error": type(e).__name__,
"code": vars(e).get('status_code', 500),
"detail": vars(e).get('detail', ''),
"body": vars(e).get('body', ''),
"errors": str(e),
}
if err['code'] == 401 and 'file=' in req.url.path:
return JSONResponse(status_code=err['code'], content=jsonable_encoder(err))
if err['code'] == 404 and 'file=ui/' in req.url.path:
return JSONResponse(status_code=err['code'], content=jsonable_encoder(err))
if err["code"] == 429:
return JSONResponse(status_code=err["code"], content=jsonable_encoder(err))
endpoint = req.scope.get("path", "err")
client = req.scope.get("client", ("0:0.0.0", 0))[0]
if not validate_log(client, endpoint):
log.error(f"API error: {req.method}: {req.url} {err}")
if not isinstance(e, HTTPException) and err['error'] != 'TypeError':
errors.display(e, 'HTTP API', [anyio, fastapi, uvicorn, starlette])
elif err['code'] in [404, 401, 400, 403]:
pass
else:
log.debug(e, exc_info=True)
return JSONResponse(status_code=err['code'], content=jsonable_encoder(err))
@app.exception_handler(HTTPException)
async def http_exception_handler(req: Request, e: HTTPException):
return handle_exception(req, e)
@app.exception_handler(Exception)
async def general_exception_handler(req: Request, e: Exception):
if isinstance(e, TypeError):
return JSONResponse(status_code=500, content=jsonable_encoder(str(e)))
else:
return handle_exception(req, e)
app.build_middleware_stack()
log.debug(f'API middleware: {[m.cls.__name__ for m in app.user_middleware]}')
@app.websocket("/internal/monitor")
async def ws_monitor(ws: WebSocket):
await ws.accept()
try:
while True:
await asyncio.sleep(1.0)
await ws.send_json({"status": "ok"})
except WebSocketDisconnect:
pass
except Exception as e:
log.error(f'WebSocket monitor: {e}')
finally:
try:
await ws.close()
except RuntimeError:
pass