bot: add check_user decorator to check user login statusfeature/user_permission

1 files changed, 27 insertions, 0 deletions

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)