From a1c7c5a223d22017508998599388c5adf9a90713 Mon Sep 17 00:00:00 2001 From: clarkzjw Date: Mon, 27 Feb 2023 15:48:09 -0800 Subject: bot: add check_user decorator to check user login status --- util.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) (limited to 'util.py') diff --git a/util.py b/util.py index 61efa0b..a0cda5e 100644 --- a/util.py +++ b/util.py @@ -1,4 +1,7 @@ from cryptography.fernet import Fernet +from telegram import Update +from dbstore.peewee_store import db, User, get_user_by_id +import functools def encrypt(input: str, key: str) -> str: @@ -11,6 +14,30 @@ def decrypt(input: str, key: str) -> str: return f.decrypt(input).decode('utf-8') +def check_user(fn): + """ + Decorator: loads User model and passes it to the function or stops the request. + Ref: https://shallowdepth.online/posts/2021/12/using-python-decorators-to-process-and-authorize-requests/ + """ + + @functools.wraps(fn) + async def wrapper(*args, **kwargs): + # Expects that Update object is always the first arg + update: Update = args[0] + user = get_user_by_id(str(update.effective_user.id)) + + if len(user) == 0: + await update.effective_message.reply_text("You are not logged in. Use `/login` to link your account first") + return + else: + # TODO + # check access_key is still valid + pass + return await fn(*args, **kwargs, user=user) + + return wrapper + + # if __name__ == "__main__": # key = Fernet.generate_key().decode('utf-8') # print(key) -- cgit v1.2.3