diff --git a/app.py b/app.py index f7ed7a0..cc46890 100644 --- a/app.py +++ b/app.py @@ -10,12 +10,13 @@ from flask import Flask, jsonify, redirect from flask_cors import CORS -from config import Config +from config import Config, limiter from routes import register_routes app = Flask(__name__) app.config.from_object(Config) CORS(app) +limiter.init_app(app) @app.route("/", methods=["GET"]) @@ -51,5 +52,19 @@ def add_common_headers(response): return response +@app.errorhandler(429) +def ratelimit_error(e): + return ( + jsonify( + { + "error": "Too many requests", + "message": "Rate limit exceeded. Please try again later.", + "rate_limit": e.description, + } + ), + 429, + ) + + if __name__ == "__main__": app.run(host="0.0.0.0", port=5000) diff --git a/config.py b/config.py index d56ce03..cfe1913 100644 --- a/config.py +++ b/config.py @@ -1,7 +1,8 @@ import os from dotenv import load_dotenv - +from flask_limiter import Limiter +from flask_limiter.util import get_remote_address # Load environment variables from .env file load_dotenv() @@ -13,3 +14,9 @@ class Config: MYSQL_PASSWORD = os.getenv("DB_PASSWORD", "myrootpassword") MYSQL_DB = os.getenv("DB_NAME", "0ce") MYSQL_CURSORCLASS = "DictCursor" + + +limiter = Limiter( + key_func=get_remote_address, + default_limits=["1000 per day", "200 per hour", "30 per minute", "3 per second"], +) diff --git a/requirements.txt b/requirements.txt index 2b01573..59da5cf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ argon2-cffi>=23.1.0 flask>=3.0.3 flask-cors>=5.0.0 +Flask-Limiter>=3.8.0 pyjwt>=2.10.0 pymysql>=1.1.1 pytest>=8.3.3