Decorator for routes that require a login with one of a set of roles

2 files changed, 25 insertions, 4 deletions

diff --git a/pysite/constants.py b/pysite/constants.py
index 3ccebf18..6339267d 100644
--- a/pysite/constants.py
+++ b/pysite/constants.py
@@ -20,6 +20,7 @@ class ValidationTypes(Enum):
 OWNER_ROLE = 267627879762755584
 ADMIN_ROLE = 267628507062992896
 MODERATOR_ROLE = 267629731250176001
+DEVOPS_ROLE = 409416496733880320
 HELPER_ROLE = 267630620367257601
 
 SERVER_ID = 267624335836053506
diff --git a/pysite/decorators.py b/pysite/decorators.py
index 447c17e4..3e07c6d2 100644
--- a/pysite/decorators.py
+++ b/pysite/decorators.py
@@ -2,14 +2,34 @@
 import os
 from functools import wraps
 from json import JSONDecodeError
-from typing import List
 
-from flask import request
+from flask import request, redirect, url_for
 from schema import Schema, SchemaError
+from werkzeug.exceptions import Forbidden
 
+from pysite.base_route import APIView, BaseView
 from pysite.constants import ErrorCodes, ValidationTypes
 
 
+def require_roles(*roles: int):
+    def inner_decorator(f):
+
+        @wraps(f)
+        def inner(self: BaseView, *args, **kwargs):
+            data = self.user_data
+
+            if data:
+                for role in roles:
+                    if role in data["roles"]:
+                        return f(self, *args, **kwargs)
+
+                raise Forbidden()
+            return redirect(url_for("discord.login"))
+        return inner
+
+    return inner_decorator
+
+
 def api_key(f):
     """
     Decorator to check if X-API-Key is valid.
@@ -18,7 +38,7 @@ def api_key(f):
     """
 
     @wraps(f)
-    def inner(self, *args, **kwargs):
+    def inner(self: APIView, *args, **kwargs):
         if not request.headers.get("X-API-Key") == os.environ.get("BOT_API_KEY"):
             return self.error(ErrorCodes.invalid_api_key)
         return f(self, *args, **kwargs)
@@ -39,7 +59,7 @@ def api_params(schema: Schema, validation_type: ValidationTypes = ValidationType
     def inner_decorator(f):
 
         @wraps(f)
-        def inner(self, *args, **kwargs):
+        def inner(self: BaseView, *args, **kwargs):
             if validation_type == ValidationTypes.json:
                 try:
                     if not request.is_json: