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:
dni ⚡
2024-05-13 15:26:25 +01:00
committed by GitHub
parent 9f8942a921
commit 6730c6ed67
7 changed files with 101 additions and 177 deletions
+76 -163
View File
@@ -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