decorators: Add `@login_required` decorator

1 files changed, 19 insertions, 0 deletions

diff --git a/wqflask/wqflask/decorators.py b/wqflask/wqflask/decorators.py
index 5930e7ec..cc69902e 100644
--- a/wqflask/wqflask/decorators.py
+++ b/wqflask/wqflask/decorators.py
@@ -1,6 +1,8 @@
 """This module contains gn2 decorators"""
 import hashlib
 import hmac
+import redis
+
 from flask import current_app, g
 from typing import Dict
 from functools import wraps
@@ -13,6 +15,23 @@ def create_hmac(data: str, secret: str) -> str:
     return hmac.new(bytearray(secret, "latin-1"),
                     bytearray(data, "utf-8"),
                     hashlib.sha1).hexdigest[:20]
+
+
+def login_required(f):
+    """Use this for endpoints where login is required"""
+    @wraps(f)
+    def wrap(*args, **kwargs):
+        user_id = (g.user_session.record.get(b"user_id",
+                                        b"").decode("utf-8") or
+                   g.user_session.record.get("user_id", ""))
+        redis_conn = redis.from_url(current_app.config["REDIS_URL"],
+                                    decode_responses=True)
+        if not redis_conn.hget("users", user_id):
+            return "You need to be logged in!", 401
+        return f(*args, **kwargs)
+    return wrap
+
+
 def edit_access_required(f):
     """Use this for endpoints where admins are required"""
     @wraps(f)