lnbits-legend/lnbits/extensions/lndhub/decorators.py

from base64 import b64decode
from fastapi.param_functions import Security

from fastapi.security.api_key import APIKeyHeader

from fastapi import Request, status
from starlette.exceptions import HTTPException

from lnbits.decorators import WalletTypeInfo, get_key_type  # type: ignore


api_key_header_auth = APIKeyHeader(
    name="AUTHORIZATION",
    auto_error=False,
    description="Admin or Invoice key for LNDHub API's",
)


async def check_wallet(
    r: Request, api_key_header_auth: str = Security(api_key_header_auth)
) -> WalletTypeInfo:
    if not api_key_header_auth:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid auth key"
        )

    t = api_key_header_auth.split(" ")[1]
    _, token = b64decode(t).decode("utf-8").split(":")

    return await get_key_type(r, api_key_header=token)


async def require_admin_key(
    r: Request, api_key_header_auth: str = Security(api_key_header_auth)
):
    wallet = await check_wallet(r, api_key_header_auth)
    if wallet.wallet_type != 0:
        # If wallet type is not admin then return the unauthorized status
        # This also covers when the user passes an invalid key type
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail="Admin key required."
        )
    else:
        return wallet