updated rate limiting

Author: marketcalls

.env.sample

@@ -1,10 +1,20 @@
 # Required for Flask sessions and user authentication
-SECRET_KEY=your_secret_key_here
+SECRET_KEY=sketchmaker_secret_key_2024_secure_random_string
+
+# Database Configuration
+DB_PATH=sqlite:///sketchmaker.db
+
+# Rate Limiting Configuration
+RATE_LIMIT_PER_MINUTE=20
+RATE_LIMIT_PER_HOUR=200
+RATE_LIMIT_PER_DAY=1000
+RATE_LIMIT_STORAGE=memory
 
 # Optional default values for new installations
 # These can be changed through the settings interface
 DEFAULT_PROVIDER=OpenAI  # Options: OpenAI, Anthropic, Google Gemini
-DEFAULT_MODEL=gpt-4o    # Should match a model name from the selected provider
+DEFAULT_MODEL=gpt-4o-mini    # Should match a model name from the selected provider
+
 
 # Optional: Set to 'True' to enable debug mode
 # FLASK_DEBUG=False

app.py

@@ -1,12 +1,5 @@
-from flask import Flask, render_template
-from flask_login import LoginManager
-from models import db, User, APIProvider, AIModel
-from blueprints.auth import auth_bp
-from blueprints.core import core_bp
-from blueprints.generate import generate_bp
-from blueprints.gallery import gallery_bp
-from blueprints.download import download_bp
-from blueprints.admin import admin
+from flask import Flask, render_template, jsonify, request
+from extensions import db, login_manager, limiter, get_rate_limit_string
 import os
 from dotenv import load_dotenv
 
@@ -18,7 +11,7 @@ def create_app():
 
     # Configuration
     app.config['SECRET_KEY'] = os.getenv('SECRET_KEY')
-    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///sketchmaker.db'
+    app.config['SQLALCHEMY_DATABASE_URI'] = os.getenv('DB_PATH', 'sqlite:///sketchmaker.db')
     app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
 
     # Ensure required environment variables are set
@@ -29,94 +22,127 @@ def create_app():
 
     # Initialize extensions
     db.init_app(app)
-    login_manager = LoginManager()
-    login_manager.login_view = 'auth.login'
+    limiter.init_app(app)
     login_manager.init_app(app)
-    
-    # Create database tables and initialize providers
-    with app.app_context():
-        db.create_all()
-        init_api_providers(app)
-    
-    # Register blueprints
-    app.register_blueprint(auth_bp)
-    app.register_blueprint(core_bp)
-    app.register_blueprint(generate_bp)
-    app.register_blueprint(gallery_bp)
-    app.register_blueprint(download_bp)
-    app.register_blueprint(admin)
-    
-    # Error handlers
-    @app.errorhandler(404)
-    def page_not_found(e):
-        return render_template('404.html'), 404
-    
-    # User loader
-    @login_manager.user_loader
-    def load_user(user_id):
-        return db.session.get(User, int(user_id))
-    
-    return app
 
-def init_api_providers(app):
-    """Initialize default API providers and models"""
     with app.app_context():
-        # Only initialize if no providers exist
-        if APIProvider.query.first() is None:
-            # Create providers
-            openai = APIProvider(name="OpenAI")
-            anthropic = APIProvider(name="Anthropic")
-            gemini = APIProvider(name="Google Gemini")
-            groq = APIProvider(name="Groq")
-            db.session.add_all([openai, anthropic, gemini, groq])
-            db.session.commit()
+        # Import models here to avoid circular imports
+        from models import User, APIProvider, AIModel
+        
+        # Import blueprints
+        from blueprints.auth import auth_bp
+        from blueprints.core import core_bp
+        from blueprints.generate import generate_bp
+        from blueprints.gallery import gallery_bp
+        from blueprints.download import download_bp
+        from blueprints.admin import admin
+        
+        # Register blueprints
+        app.register_blueprint(auth_bp)
+        app.register_blueprint(core_bp)
+        app.register_blueprint(generate_bp)
+        app.register_blueprint(gallery_bp)
+        app.register_blueprint(download_bp)
+        app.register_blueprint(admin)
+        
+        # Error handlers
+        @app.errorhandler(404)
+        def page_not_found(e):
+            return render_template('errors/404.html'), 404
 
-            # OpenAI models
-            openai_models = [
-                "gpt-4o",
-                "gpt-4o-mini",
-                "o1-preview",
-                "o1-mini"
-            ]
-            for model in openai_models:
-                db.session.add(AIModel(name=model, provider_id=openai.id))
+        @app.errorhandler(429)
+        def ratelimit_handler(e):
+            # Default retry time of 60 seconds if not provided
+            retry_after = getattr(e, 'retry_after', 60)
+            if retry_after is None:
+                retry_after = 60
 
-            # Anthropic models
-            anthropic_models = [
-                "claude-3-5-sonnet-20241022",
-                "claude-3-5-haiku-20241022",
-                "claude-3-opus-20240229",
-                "claude-3-haiku-20240307"
-            ]
-            for model in anthropic_models:
-                db.session.add(AIModel(name=model, provider_id=anthropic.id))
+            # For API endpoints, return JSON response
+            if request.path.startswith('/api/'):
+                return jsonify({
+                    "error": "Rate limit exceeded",
+                    "message": "Too many requests. Please try again later.",
+                    "retry_after": retry_after
+                }), 429
+            
+            # For web pages, render template
+            return render_template('errors/429.html', 
+                retry_after=retry_after
+            ), 429
 
-            # Google Gemini models
-            gemini_models = [
-                "gemini-1.5-flash-002",
-                "gemini-1.5-flash-exp-0827",
-                "gemini-1.5-flash-8b-exp-0827",
-                "gemini-1.5-pro-002",
-                "gemini-1.5-pro-exp-0827"
-            ]
-            for model in gemini_models:
-                db.session.add(AIModel(name=model, provider_id=gemini.id))
+        @app.errorhandler(500)
+        def internal_server_error(e):
+            return render_template('errors/500.html'), 500
+        
+        # User loader
+        @login_manager.user_loader
+        def load_user(user_id):
+            return db.session.get(User, int(user_id))
 
-            # Groq models
-            groq_models = [
-                "llama-3.1-70b-versatile",
-                "llama-3.1-8b-instant",
-                "llama-3.2-11b-text-preview",
-                "llama-3.2-11b-vision-preview",
-                "llama-3.2-1b-preview",
-                "llama-3.2-3b-preview",
-                "llama-3.2-90b-text-preview",
-                "llama-3.2-90b-vision-preview"
-            ]
-            for model in groq_models:
-                db.session.add(AIModel(name=model, provider_id=groq.id))
+        def init_api_providers():
+            """Initialize default API providers and models"""
+            # Only initialize if no providers exist
+            if APIProvider.query.first() is None:
+                # Create providers
+                openai = APIProvider(name="OpenAI")
+                anthropic = APIProvider(name="Anthropic")
+                gemini = APIProvider(name="Google Gemini")
+                groq = APIProvider(name="Groq")
+                db.session.add_all([openai, anthropic, gemini, groq])
+                db.session.commit()
 
-            db.session.commit()
+                # OpenAI models
+                openai_models = [
+                    "gpt-4o",
+                    "gpt-4o-mini",
+                    "o1-preview",
+                    "o1-mini"
+                ]
+                for model in openai_models:
+                    db.session.add(AIModel(name=model, provider_id=openai.id))
+
+                # Anthropic models
+                anthropic_models = [
+                    "claude-3-5-sonnet-20241022",
+                    "claude-3-5-haiku-20241022",
+                    "claude-3-opus-20240229",
+                    "claude-3-haiku-20240307"
+                ]
+                for model in anthropic_models:
+                    db.session.add(AIModel(name=model, provider_id=anthropic.id))
+
+                # Google Gemini models
+                gemini_models = [
+                    "gemini-1.5-flash-002",
+                    "gemini-1.5-flash-exp-0827",
+                    "gemini-1.5-flash-8b-exp-0827",
+                    "gemini-1.5-pro-002",
+                    "gemini-1.5-pro-exp-0827"
+                ]
+                for model in gemini_models:
+                    db.session.add(AIModel(name=model, provider_id=gemini.id))
+
+                # Groq models
+                groq_models = [
+                    "llama-3.1-70b-versatile",
+                    "llama-3.1-8b-instant",
+                    "llama-3.2-11b-text-preview",
+                    "llama-3.2-11b-vision-preview",
+                    "llama-3.2-1b-preview",
+                    "llama-3.2-3b-preview",
+                    "llama-3.2-90b-text-preview",
+                    "llama-3.2-90b-vision-preview"
+                ]
+                for model in groq_models:
+                    db.session.add(AIModel(name=model, provider_id=groq.id))
+
+                db.session.commit()
+        
+        # Create database tables and initialize providers
+        db.create_all()
+        init_api_providers()
+    
+    return app
 
 # Create the application instance for gunicorn
 app = create_app()

blueprints/admin.py

@@ -2,6 +2,7 @@
 from flask_login import login_required, current_user
 from functools import wraps
 from models import db, User
+from extensions import limiter, get_rate_limit_string
 
 admin = Blueprint('admin', __name__)
 
@@ -15,13 +16,15 @@ def decorated_function(*args, **kwargs):
     return decorated_function
 
 @admin.route('/manage')
+@limiter.limit(get_rate_limit_string())
 @login_required
 @admin_required
 def manage():
     users = User.query.all()
     return render_template('admin/manage.html', users=users)
 
 @admin.route('/manage/user/<int:user_id>', methods=['POST'])
+@limiter.limit(get_rate_limit_string())
 @login_required
 @admin_required
 def update_user(user_id):

blueprints/auth.py

@@ -2,10 +2,12 @@
 from flask_login import login_user, logout_user, login_required, current_user
 from werkzeug.security import generate_password_hash, check_password_hash
 from models import db, User
+from extensions import limiter, get_rate_limit_string
 
 auth_bp = Blueprint('auth', __name__)
 
 @auth_bp.route('/login', methods=['GET', 'POST'])
+@limiter.limit(get_rate_limit_string())
 def login():
     # Redirect to dashboard if user is already logged in
     if current_user.is_authenticated:
@@ -37,6 +39,7 @@ def login():
     return render_template('auth/login.html')
 
 @auth_bp.route('/register', methods=['GET', 'POST'])
+@limiter.limit(get_rate_limit_string())
 def register():
     # Redirect to dashboard if user is already logged in
     if current_user.is_authenticated:
@@ -78,6 +81,7 @@ def register():
     return render_template('auth/register.html')
 
 @auth_bp.route('/logout')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def logout():
     logout_user()

blueprints/core.py

@@ -2,19 +2,23 @@
 from flask_login import login_required, current_user
 from models import db, User, APIProvider, AIModel
 from datetime import datetime
+from extensions import limiter, get_rate_limit_string
 
 core_bp = Blueprint('core', __name__)
 
 @core_bp.route('/')
+@limiter.limit(get_rate_limit_string())
 def index():
     return render_template('index.html', current_year=datetime.now().year)
 
 @core_bp.route('/dashboard')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def dashboard():
     return render_template('dashboard.html')
 
 @core_bp.route('/settings')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def settings():
     # Get all providers and models
@@ -26,6 +30,7 @@ def settings():
                          models=models)
 
 @core_bp.route('/settings/update', methods=['POST'])
+@limiter.limit(get_rate_limit_string())
 @login_required
 def update_settings():
     try:
@@ -72,6 +77,7 @@ def update_settings():
         return jsonify({'error': str(e)}), 500
 
 @core_bp.route('/settings/models/<int:provider_id>')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def get_provider_models(provider_id):
     try:

blueprints/download.py

@@ -1,11 +1,13 @@
 from flask import Blueprint, send_file, abort
 from flask_login import login_required, current_user
 from models import Image
+from extensions import limiter, get_rate_limit_string
 import os
 
 download_bp = Blueprint('download', __name__)
 
 @download_bp.route('/download/<filename>/<format>')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def download_image(filename, format):
     # Get base filename without extension

blueprints/gallery.py

@@ -1,11 +1,13 @@
 from flask import Blueprint, render_template, abort
 from flask_login import login_required, current_user
 from models import Image
+from extensions import limiter, get_rate_limit_string
 import markdown2
 
 gallery_bp = Blueprint('gallery', __name__)
 
 @gallery_bp.route('/gallery')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def view_gallery():
     # Get user's images ordered by creation date (newest first)
@@ -16,6 +18,7 @@ def view_gallery():
     return render_template('gallery.html', images=images)
 
 @gallery_bp.route('/gallery/<int:image_id>')
+@limiter.limit(get_rate_limit_string())
 @login_required
 def view_image(image_id):
     # Get the specific image and verify ownership