40 lines
976 B
Python
40 lines
976 B
Python
|
from functools import wraps
|
||
|
|
||
|
import basicauth
|
||
|
from flask import Response, request
|
||
|
|
||
|
from llm_server import opts
|
||
|
|
||
|
|
||
|
def check_auth(token):
|
||
|
password = None
|
||
|
if token.startswith('Basic '):
|
||
|
try:
|
||
|
_, password = basicauth.decode(token)
|
||
|
except:
|
||
|
return False
|
||
|
elif token.startswith('Bearer '):
|
||
|
password = token.split('Bearer ', maxsplit=1)
|
||
|
del password[0]
|
||
|
password = ''.join(password)
|
||
|
print(password, token)
|
||
|
return password == opts.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
|