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

from functools import wraps

import basicauth
from flask import Response, request

from llm_server import opts


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 opts.admin_token:
        # The admin token is not set/enabled.
        # Default: deny all.
        return False
    return parse_token(token) == 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