diff options
| -rw-r--r-- | pysite/base_route.py | 17 | ||||
| -rw-r--r-- | pysite/decorators.py | 21 | ||||
| -rw-r--r-- | pysite/views/api/bot/tag.py | 54 |
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}) |