1 files changed, 40 insertions, 0 deletions
diff --git a/frest/decorators.py b/frest/decorators.py
new file mode 100644
index 0000000..7ce79d7
--- /dev/null
+++ b/frest/decorators.py
@@ -0,0 +1,40 @@
+from flask import request, abort
+from auth.models import Token
+from functools import wraps
+
+
+def check_token(f):
+    @wraps(f)
+    def inner(*args, **kwargs):
+        userid = request.url.split("/")[-1]
+        headers = request.headers
+        if not headers.get("Authorization"):
+            abort(403)
+
+        auth = request.headers.get("Authorization")
+        token = Token.query.filter_by(string=auth).first()
+        if not token:
+            abort(403)
+
+        if userid.isdigit():
+            if int(userid) != token.user.userId and not token.user.is_admin:
+                abort(403)
+
+        return f(*args, **kwargs)
+
+    return inner
+
+
+def admin_required(f):
+    @wraps(f)
+    def inner(*args, **kwargs):
+        header = request.headers
+
+        auth = request.headers.get("Authorization")
+        token = Token.query.filter_by(string=auth).first()
+        if not token.user.is_admin:
+            abort(403)
+
+        return f(*args, **kwargs)
+
+    return inner