add rate limits

app.py

@@ -1,9 +1,12 @@
-from flask import Flask
+from flask import Flask, g, request
 from flask_sqlalchemy import SQLAlchemy
 from sqlalchemy import MetaData
 from flask_migrate import Migrate
 import version
 from lib import cachebust
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
+from lib import get_viewer
 
 app = Flask(__name__)
 
@@ -44,3 +47,15 @@ def inject_static():
     def static(filename, **kwargs):
         return url_for('static', filename=filename, **kwargs)
     return {'st': static}
+
+def rate_limit_key():
+    viewer = get_viewer()
+    if viewer:
+        return viewer.id
+    for address in request.access_route:
+        if address != '127.0.0.1':
+            print(address)
+            return address
+    return request.remote_addr
+
+limiter = Limiter(app, key_func=rate_limit_key)

lib/__init__.py

@@ -2,4 +2,4 @@ from .auth import require_auth
 from .interval import decompose_interval
 from .interval import SCALES as interval_scales
 from .cachebust import cachebust
-from .session import set_session_cookie
+from .session import set_session_cookie, get_viewer_session, get_viewer

lib/session.py

@@ -1,5 +1,18 @@
+from flask import request
+
 def set_session_cookie(session, response, secure=True):
     response.set_cookie('forget_sid', session.id,
         max_age=60*60*48,
         httponly=True,
         secure=secure)
+
+def get_viewer_session():
+    from model import Session
+    sid = request.cookies.get('forget_sid', None)
+    if sid:
+        return Session.query.get(sid)
+
+def get_viewer():
+    session = get_viewer_session()
+    if session:
+        return session.account

requirements.txt

@@ -6,6 +6,7 @@ celery==4.1.0
 click==6.7
 contextlib2==0.5.5
 Flask==0.12.2
+Flask-Limiter==0.9.5
 Flask-Migrate==2.0.4
 Flask-Script==2.0.5
 Flask-SQLAlchemy==2.2
@@ -14,6 +15,7 @@ honcho==1.0.1
 itsdangerous==0.24
 Jinja2==2.9.6
 kombu==4.1.0
+limits==1.2.1
 Mako==1.0.6
 MarkupSafe==1.0
 psycopg2==2.7.1

routes.py

@@ -4,8 +4,9 @@ import lib.twitter
 import lib
 from lib import require_auth
 from lib import set_session_cookie
+from lib import get_viewer_session
 from model import Account, Session, Post, TwitterArchive
-from app import app, db, sentry
+from app import app, db, sentry, limiter
 import tasks
 from zipfile import BadZipFile
 from twitter import TwitterError
@@ -14,16 +15,13 @@ import version
 
 @app.before_request
 def load_viewer():
-    g.viewer = None
-    sid = request.cookies.get('forget_sid', None)
-    if sid:
-        g.viewer = Session.query.get(sid)
-        if g.viewer and sentry:
-            sentry.user_context({
-                 'id': g.viewer.account.id,
-                 'username': g.viewer.account.screen_name,
-                 'service': g.viewer.account.service
-                })
+    g.viewer = get_viewer_session()
+    if g.viewer and sentry:
+        sentry.user_context({
+                'id': g.viewer.account.id,
+                'username': g.viewer.account.screen_name,
+                'service': g.viewer.account.service
+            })
 
 @app.context_processor
 def inject_version():
@@ -31,7 +29,7 @@ def inject_version():
 
 @app.after_request
 def touch_viewer(resp):
-    if g.viewer:
+    if 'viewer' in g and g.viewer:
         set_session_cookie(g.viewer, resp, app.config.get('HTTPS'))
         g.viewer.touch()
         db.session.commit()
@@ -49,6 +47,7 @@ def index():
                 twitter_login_error = 'twitter_login_error' in request.args)
 
 @app.route('/login/twitter')
+@limiter.limit('3/minute')
 def twitter_login_step1():
     try:
         return redirect(lib.twitter.get_login_url(
@@ -59,6 +58,7 @@ def twitter_login_step1():
         return redirect(url_for('index', twitter_login_error='', _anchor='log_in'))
 
 @app.route('/login/twitter/callback')
+@limiter.limit('3/minute')
 def twitter_login_step2():
     try:
         oauth_token = request.args['oauth_token']
@@ -78,6 +78,7 @@ def twitter_login_step2():
         return redirect(url_for('index', twitter_login_error='', _anchor='log_in'))
 
 @app.route('/upload_tweet_archive', methods=('POST',))
+@limiter.limit('10/10 minutes')
 @require_auth
 def upload_tweet_archive():
     ta = TwitterArchive(account = g.viewer.account,