local-llm-server/llm_server/routes/auth.py

47 lines
1.2 KiB
Python

from functools import wraps
import basicauth
from flask import Response, request
from llm_server.config.global_config import GlobalConfig
def parse_token(input_token):
password = None
if input_token.startswith('Basic '):
try:
_, password = basicauth.decode(input_token)
except:
return False
elif input_token.startswith('Bearer '):
password = input_token.split('Bearer ', maxsplit=1)
del password[0]
password = ''.join(password)
return password
def check_auth(token):
if not GlobalConfig.get().admin_token:
# The admin token is not set/enabled.
# Default: deny all.
return False
return parse_token(token) == GlobalConfig.get().admin_token
def authenticate():
"""Sends a 401 response that enables basic auth"""
return Response(
'AUTHENTICATION REQUIRED', 401,
{'WWW-Authenticate': 'Basic realm="Login Required"'})
def requires_auth(f):
@wraps(f)
def decorated(*args, **kwargs):
auth = request.headers.get('Authorization')
if not auth or not check_auth(auth):
return authenticate()
return f(*args, **kwargs)
return decorated