Move API validation decorator to its own file #yxdk

3 files changed, 49 insertions, 43 deletions

diff --git a/pysite/base_route.py b/pysite/base_route.py
index 730b3e10..e1b9c6b2 100644
--- a/pysite/base_route.py
+++ b/pysite/base_route.py
@@ -2,9 +2,8 @@
 import os
 import random
 import string
-from functools import wraps
 
-from flask import Blueprint, g, jsonify, render_template, request
+from flask import Blueprint, g, jsonify, render_template
 from flask.views import MethodView
 
 from rethinkdb.ast import Table
@@ -93,20 +92,6 @@ class APIView(RouteView):
         pool = random.choices(string.ascii_letters + string.digits, k=32)
         return "".join(pool)
 
-    def valid_api_key(f):
-        """
-        Decorator to check if X-API-Key is valid.
-        """
-        @wraps(f)
-        def has_valid_api_key(*args, **kwargs):
-            if not request.headers.get("X-API-Key") == os.environ.get("API_KEY"):
-                resp = jsonify({"error_code": 401, "error_message": "Invalid API-Key"})
-                resp.status_code = 401
-                return resp
-            return f(*args, **kwargs)
-
-        return has_valid_api_key
-
     def error(self, error_code: ErrorCodes):
 
         data = {
diff --git a/pysite/decorators.py b/pysite/decorators.py
new file mode 100644
index 00000000..6951e875
--- /dev/null
+++ b/pysite/decorators.py
@@ -0,0 +1,21 @@
+# coding=utf-8
+import os
+
+from flask import request
+
+from pysite.constants import ErrorCodes
+
+
+def valid_api_key(f):
+    """
+    Decorator to check if X-API-Key is valid.
+
+    Should only be applied to functions on APIView routes.
+    """
+
+    def has_valid_api_key(self, *args, **kwargs):
+        if not request.headers.get("X-API-Key") == os.environ.get("API_KEY"):
+            return self.error(ErrorCodes.invalid_api_key)
+        return f(*args, **kwargs)
+
+    return has_valid_api_key
diff --git a/pysite/views/api/bot/tag.py b/pysite/views/api/bot/tag.py
index 84fd8977..ef17e8fa 100644
--- a/pysite/views/api/bot/tag.py
+++ b/pysite/views/api/bot/tag.py
@@ -4,6 +4,7 @@ from flask import jsonify, request
 
 from pysite.base_route import APIView, DBViewMixin
 from pysite.constants import ErrorCodes
+from pysite.decorators import valid_api_key
 
 
 class TagView(APIView, DBViewMixin):
@@ -12,45 +13,44 @@ class TagView(APIView, DBViewMixin):
     table_name = "tag"
     table_primary_key = "tag_name"
 
+    @valid_api_key
     def get(self):
         """
         Data must be provided as params,
         API key must be provided as header
         """
-        api_key = request.headers.get("X-API-Key")
+
         tag_name = request.args.get("tag_name")
 
-        if self.validate_key(api_key):
-            if tag_name:
-                data = self.db.get(self.table_name, tag_name)
-            else:
-                data = self.db.pluck(self.table_name, "tag_name")
+        if tag_name:
+            data = self.db.get(self.table_name, tag_name)
         else:
-            return self.error(ErrorCodes.invalid_api_key)
+            data = self.db.pluck(self.table_name, "tag_name")
 
         return jsonify(data or {})
 
+    @valid_api_key
     def post(self):
-        """ Data must be provided as JSON. """
-        indata = request.get_json()
-        tag_name = indata.get("tag_name")
-        tag_content = indata.get("tag_content")
-        tag_category = indata.get("tag_category")
-        api_key = request.headers.get("X-API-Key")
-
-        if self.validate_key(api_key):
-            if tag_name and tag_content:
-                self.db.insert(
-                    self.table_name,
-                    {
-                        "tag_name": tag_name,
-                        "tag_content": tag_content,
-                        "tag_category": tag_category
-                    }
-                )
-            else:
-                return self.error(ErrorCodes.missing_parameters)
+        """
+        Data must be provided as JSON.
+        """
+
+        data = request.get_json()
+
+        tag_name = data.get("tag_name")
+        tag_content = data.get("tag_content")
+        tag_category = data.get("tag_category")
+
+        if tag_name and tag_content:
+            self.db.insert(
+                self.table_name,
+                {
+                    "tag_name": tag_name,
+                    "tag_content": tag_content,
+                    "tag_category": tag_category
+                }
+            )
         else:
-            return self.error(ErrorCodes.invalid_api_key)
+            return self.error(ErrorCodes.missing_parameters)
 
         return jsonify({"success": True})