refactor: fix duplicate keychecker (#2339)
* refactor: fix duplicate keychecker - refactor KeyChecker to be more approachable - only 1 sql query needed even if you use `get_key_type` - rename `WalletType` to `KeyType` wallet type was misleading fix test sorting * fixup! * revert 404
This commit is contained in:
+76
-163
@@ -17,23 +17,32 @@ from lnbits.core.crud import (
|
||||
get_user,
|
||||
get_wallet_for_key,
|
||||
)
|
||||
from lnbits.core.models import User, Wallet, WalletType, WalletTypeInfo
|
||||
from lnbits.core.models import KeyType, User, WalletTypeInfo
|
||||
from lnbits.db import Filter, Filters, TFilterModel
|
||||
from lnbits.settings import AuthMethods, settings
|
||||
|
||||
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/v1/auth", auto_error=False)
|
||||
|
||||
api_key_header = APIKeyHeader(
|
||||
name="X-API-KEY",
|
||||
auto_error=False,
|
||||
description="Admin or Invoice key for wallet API's",
|
||||
)
|
||||
api_key_query = APIKeyQuery(
|
||||
name="api-key",
|
||||
auto_error=False,
|
||||
description="Admin or Invoice key for wallet API's",
|
||||
)
|
||||
|
||||
|
||||
class KeyChecker(SecurityBase):
|
||||
def __init__(
|
||||
self,
|
||||
scheme_name: Optional[str] = None,
|
||||
auto_error: bool = True,
|
||||
api_key: Optional[str] = None,
|
||||
expected_key_type: Optional[KeyType] = None,
|
||||
):
|
||||
self.scheme_name = scheme_name or self.__class__.__name__
|
||||
self.auto_error: bool = auto_error
|
||||
self._key_type: WalletType = WalletType.invoice
|
||||
self.auto_error: bool = True
|
||||
self.expected_key_type = expected_key_type
|
||||
self._api_key = api_key
|
||||
if api_key:
|
||||
openapi_model = APIKey(
|
||||
@@ -49,185 +58,82 @@ class KeyChecker(SecurityBase):
|
||||
name="X-API-KEY",
|
||||
description="Wallet API Key - HEADER",
|
||||
)
|
||||
self.wallet: Optional[Wallet] = None
|
||||
self.model: APIKey = openapi_model
|
||||
|
||||
async def __call__(self, request: Request):
|
||||
try:
|
||||
key_value = (
|
||||
self._api_key
|
||||
if self._api_key
|
||||
else request.headers.get("X-API-KEY") or request.query_params["api-key"]
|
||||
)
|
||||
# FIXME: Find another way to validate the key. A fetch from DB should be
|
||||
# avoided here. Also, we should not return the wallet here - thats
|
||||
# silly. Possibly store it in a Redis DB
|
||||
wallet = await get_wallet_for_key(key_value, self._key_type)
|
||||
if not wallet:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="Invalid key or wallet.",
|
||||
)
|
||||
self.wallet = wallet
|
||||
except KeyError as exc:
|
||||
async def __call__(self, request: Request) -> WalletTypeInfo:
|
||||
|
||||
key_value = (
|
||||
self._api_key
|
||||
if self._api_key
|
||||
else request.headers.get("X-API-KEY") or request.query_params.get("api-key")
|
||||
)
|
||||
|
||||
if not key_value:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_REQUEST, detail="`X-API-KEY` header missing."
|
||||
) from exc
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="No Api Key provided.",
|
||||
)
|
||||
|
||||
wallet = await get_wallet_for_key(key_value)
|
||||
|
||||
class WalletInvoiceKeyChecker(KeyChecker):
|
||||
"""
|
||||
WalletInvoiceKeyChecker will ensure that the provided invoice
|
||||
wallet key is correct and populate g().wallet with the wallet
|
||||
for the key in `X-API-key`.
|
||||
if not wallet:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.NOT_FOUND,
|
||||
detail="Wallet not found.",
|
||||
)
|
||||
|
||||
The checker will raise an HTTPException when the key is wrong in some ways.
|
||||
"""
|
||||
if self.expected_key_type is KeyType.admin and wallet.adminkey != key_value:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="Invalid adminkey.",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scheme_name: Optional[str] = None,
|
||||
auto_error: bool = True,
|
||||
api_key: Optional[str] = None,
|
||||
):
|
||||
super().__init__(scheme_name, auto_error, api_key)
|
||||
self._key_type = WalletType.invoice
|
||||
if (
|
||||
wallet.user != settings.super_user
|
||||
and wallet.user not in settings.lnbits_admin_users
|
||||
and settings.lnbits_admin_extensions
|
||||
and request["path"].split("/")[1] in settings.lnbits_admin_extensions
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.FORBIDDEN,
|
||||
detail="User not authorized for this extension.",
|
||||
)
|
||||
|
||||
|
||||
class WalletAdminKeyChecker(KeyChecker):
|
||||
"""
|
||||
WalletAdminKeyChecker will ensure that the provided admin
|
||||
wallet key is correct and populate g().wallet with the wallet
|
||||
for the key in `X-API-key`.
|
||||
|
||||
The checker will raise an HTTPException when the key is wrong in some ways.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scheme_name: Optional[str] = None,
|
||||
auto_error: bool = True,
|
||||
api_key: Optional[str] = None,
|
||||
):
|
||||
super().__init__(scheme_name, auto_error, api_key)
|
||||
self._key_type = WalletType.admin
|
||||
|
||||
|
||||
api_key_header = APIKeyHeader(
|
||||
name="X-API-KEY",
|
||||
auto_error=False,
|
||||
description="Admin or Invoice key for wallet API's",
|
||||
)
|
||||
api_key_query = APIKeyQuery(
|
||||
name="api-key",
|
||||
auto_error=False,
|
||||
description="Admin or Invoice key for wallet API's",
|
||||
)
|
||||
key_type = KeyType.admin if wallet.adminkey == key_value else KeyType.invoice
|
||||
return WalletTypeInfo(key_type, wallet)
|
||||
|
||||
|
||||
async def get_key_type(
|
||||
r: Request,
|
||||
request: Request,
|
||||
api_key_header: str = Security(api_key_header),
|
||||
api_key_query: str = Security(api_key_query),
|
||||
) -> WalletTypeInfo:
|
||||
token = api_key_header or api_key_query
|
||||
|
||||
if not token:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="Invoice (or Admin) key required.",
|
||||
)
|
||||
|
||||
for wallet_type, wallet_checker in zip(
|
||||
[WalletType.admin, WalletType.invoice],
|
||||
[WalletAdminKeyChecker, WalletInvoiceKeyChecker],
|
||||
):
|
||||
try:
|
||||
checker = wallet_checker(api_key=token)
|
||||
await checker.__call__(r)
|
||||
if checker.wallet is None:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.NOT_FOUND, detail="Wallet does not exist."
|
||||
)
|
||||
wallet = WalletTypeInfo(wallet_type, checker.wallet)
|
||||
if (
|
||||
wallet.wallet.user != settings.super_user
|
||||
and wallet.wallet.user not in settings.lnbits_admin_users
|
||||
) and (
|
||||
settings.lnbits_admin_extensions
|
||||
and r["path"].split("/")[1] in settings.lnbits_admin_extensions
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.FORBIDDEN,
|
||||
detail="User not authorized for this extension.",
|
||||
)
|
||||
return wallet
|
||||
except HTTPException as exc:
|
||||
if exc.status_code == HTTPStatus.BAD_REQUEST:
|
||||
raise
|
||||
elif exc.status_code == HTTPStatus.UNAUTHORIZED:
|
||||
# we pass this in case it is not an invoice key, nor an admin key,
|
||||
# and then return NOT_FOUND at the end of this block
|
||||
pass
|
||||
else:
|
||||
raise
|
||||
except Exception:
|
||||
raise
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.NOT_FOUND, detail="Wallet does not exist."
|
||||
)
|
||||
check: KeyChecker = KeyChecker(api_key=api_key_header or api_key_query)
|
||||
return await check(request)
|
||||
|
||||
|
||||
async def require_admin_key(
|
||||
r: Request,
|
||||
request: Request,
|
||||
api_key_header: str = Security(api_key_header),
|
||||
api_key_query: str = Security(api_key_query),
|
||||
):
|
||||
token = api_key_header or api_key_query
|
||||
|
||||
if not token:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="Admin key required.",
|
||||
)
|
||||
|
||||
wallet = await get_key_type(r, token)
|
||||
|
||||
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=HTTPStatus.UNAUTHORIZED, detail="Admin key required."
|
||||
)
|
||||
else:
|
||||
return wallet
|
||||
) -> WalletTypeInfo:
|
||||
check: KeyChecker = KeyChecker(
|
||||
api_key=api_key_header or api_key_query,
|
||||
expected_key_type=KeyType.admin,
|
||||
)
|
||||
return await check(request)
|
||||
|
||||
|
||||
async def require_invoice_key(
|
||||
r: Request,
|
||||
request: Request,
|
||||
api_key_header: str = Security(api_key_header),
|
||||
api_key_query: str = Security(api_key_query),
|
||||
):
|
||||
token = api_key_header or api_key_query
|
||||
|
||||
if not token:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="Invoice (or Admin) key required.",
|
||||
)
|
||||
|
||||
wallet = await get_key_type(r, token)
|
||||
|
||||
if (
|
||||
wallet.wallet_type != WalletType.admin
|
||||
and wallet.wallet_type != WalletType.invoice
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.UNAUTHORIZED,
|
||||
detail="Invoice (or Admin) key required.",
|
||||
)
|
||||
else:
|
||||
return wallet
|
||||
) -> WalletTypeInfo:
|
||||
check: KeyChecker = KeyChecker(
|
||||
api_key=api_key_header or api_key_query,
|
||||
expected_key_type=KeyType.invoice,
|
||||
)
|
||||
return await check(request)
|
||||
|
||||
|
||||
async def check_access_token(
|
||||
@@ -255,8 +161,15 @@ async def check_user_exists(
|
||||
user = await get_user(account.id)
|
||||
assert user, "User not found for account."
|
||||
|
||||
if not user.admin and r["path"].split("/")[1] in settings.lnbits_admin_extensions:
|
||||
raise HTTPException(HTTPStatus.FORBIDDEN, "User not authorized for extension.")
|
||||
if (
|
||||
user.id != settings.super_user
|
||||
and user.id not in settings.lnbits_admin_users
|
||||
and settings.lnbits_admin_extensions
|
||||
and r["path"].split("/")[1] in settings.lnbits_admin_extensions
|
||||
):
|
||||
raise HTTPException(
|
||||
HTTPStatus.UNAUTHORIZED, "User not authorized for extension."
|
||||
)
|
||||
|
||||
return user
|
||||
|
||||
|
||||
Reference in New Issue
Block a user