77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
import functools
|
|
|
|
import jingrow
|
|
import pyotp
|
|
from jingrow import auth
|
|
|
|
|
|
def enabled(user_key: str = "user", raise_error: bool = False):
|
|
"""
|
|
Guard to check if two-factor authentication is enabled for a user.
|
|
|
|
:param user_key: The key to retrieve the user from the request form if the session user is "Guest". Defaults to "user".
|
|
:param raise_error: Whether to raise a PermissionError if 2FA is not enabled. Defaults to False.
|
|
:raises jingrow.PermissionError: If 2FA is not enabled and `raise_error` is True.
|
|
"""
|
|
|
|
def get_user() -> str:
|
|
if jingrow.session.user == "Guest":
|
|
return (jingrow.request.json or jingrow.request.form).get(user_key, jingrow.session.user)
|
|
return jingrow.session.user
|
|
|
|
def wrapper(fn):
|
|
@functools.wraps(fn)
|
|
def inner(*args, **kwargs):
|
|
if jingrow.get_value("User 2FA", get_user(), "enabled"):
|
|
return fn(*args, **kwargs)
|
|
if raise_error:
|
|
message = "Two-factor authentication is not enabled."
|
|
jingrow.throw(message, jingrow.PermissionError)
|
|
return None
|
|
|
|
return inner
|
|
|
|
return wrapper
|
|
|
|
|
|
def verify(user_key: str = "user", code_key: str = "totp_code", raise_error: bool = False):
|
|
"""
|
|
Guard to verify two-factor authentication for a user.
|
|
|
|
:param user_key: The key to retrieve the user from the request form if the session user is "Guest". Defaults to "user".
|
|
:param code_key: The key to retrieve the TOTP code from the request form. Defaults to "totp_code".
|
|
:param raise_error: Whether to raise a PermissionError if 2FA verification fails. Defaults to False.
|
|
:raises jingrow.PermissionError: If 2FA verification fails and `raise_error` is True.
|
|
"""
|
|
|
|
def get_user() -> str:
|
|
if jingrow.session.user == "Guest":
|
|
return (jingrow.request.json or jingrow.request.form).get(user_key, jingrow.session.user)
|
|
return jingrow.session.user
|
|
|
|
def get_code() -> str:
|
|
return (jingrow.request.json or jingrow.request.form).get(code_key, "")
|
|
|
|
def verify_code(user: str, code: str):
|
|
secret = auth.get_decrypted_password("User 2FA", user, "totp_secret")
|
|
return secret and pyotp.TOTP(secret).verify(code)
|
|
|
|
def wrapper(fn):
|
|
@functools.wraps(fn)
|
|
def inner(*args, **kwargs):
|
|
user = get_user()
|
|
code = get_code()
|
|
|
|
if user and not jingrow.get_value("User 2FA", user, "enabled"):
|
|
return fn(*args, **kwargs)
|
|
if user and code and verify_code(user, code):
|
|
return fn(*args, **kwargs)
|
|
if raise_error:
|
|
message = "Two-factor authentication verification failed."
|
|
jingrow.throw(message, jingrow.PermissionError)
|
|
return None
|
|
|
|
return inner
|
|
|
|
return wrapper
|