Add a bit of rate limiting

titanembeds/app.py

@@ -1,7 +1,7 @@
 from config import config
 from database import db
 from flask import Flask, render_template, request, session, url_for, redirect, jsonify
-from titanembeds.utils import cache
+from titanembeds.utils import cache, rate_limiter
 import blueprints.api
 import blueprints.user
 import blueprints.embed
@@ -16,6 +16,7 @@ app.secret_key = config['app-secret']
 
 db.init_app(app)
 cache.init_app(app, config={'CACHE_TYPE': 'simple'})
+rate_limiter.init_app(app)
 
 app.register_blueprint(blueprints.api.api, url_prefix="/api", template_folder="/templates")
 app.register_blueprint(blueprints.user.user, url_prefix="/user", template_folder="/templates")

titanembeds/blueprints/api/api.py

@@ -1,6 +1,6 @@
 from titanembeds.database import db, Guilds, UnauthenticatedUsers, UnauthenticatedBans, AuthenticatedUsers
 from titanembeds.decorators import valid_session_required, discord_users_only
-from titanembeds.utils import get_client_ipaddr, discord_api
+from titanembeds.utils import get_client_ipaddr, discord_api, rate_limiter, channel_ratelimit_key, guild_ratelimit_key
 from flask import Blueprint, abort, jsonify, session, request
 from sqlalchemy import and_
 import random
@@ -78,11 +78,23 @@ def update_user_status(guild_id, username, user_key=None):
             dbUser.ip_address = ip_address
             db.session.commit()
     else:
-        pass #authenticated user todo
+        status = {
+            'username': username,
+            'guild_id': guild_id,
+            'user_id': session['user_id'],
+            'banned': checkUserBanned(guild_id),
+            'revoked': checkUserRevoke(guild_id)
+        }
+        if status['banned'] or status['revoked']:
+            return status
+        dbUser = db.session.query(AuthenticatedUsers).filter(AuthenticatedUsers.guild_id == guild_id, AuthenticatedUsers.client_id == status['user_id']).first()
+        dbUser.bumpTimestamp()
     return status
 
 @api.route("/fetch", methods=["GET"])
 @valid_session_required(api=True)
+@rate_limiter.limit("2500/hour")
+@rate_limiter.limit("12/minute", key_func = channel_ratelimit_key)
 def fetch():
     channel_id = request.args.get('channel_id')
     after_snowflake = request.args.get('after', None, type=int)
@@ -93,12 +105,18 @@ def fetch():
     status = update_user_status(channel_id, session['username'], key)
     if status['banned'] or status['revoked']:
         messages = {}
+        status_code = 401
     else:
         messages = discord_api.get_channel_messages(channel_id, after_snowflake)
-    return jsonify(messages=messages, status=status)
+        status_code = messages['code']
+    response = jsonify(messages=messages.get('content', messages), status=status)
+    resonse.status_code = status_code
+    return response
 
 @api.route("/post", methods=["POST"])
 @valid_session_required(api=True)
+@rate_limiter.limit("1200/hour")
+@rate_limiter.limit("6/minute", key_func = channel_ratelimit_key)
 def post():
     channel_id = request.form.get('channel_id')
     content = request.form.get('content')
@@ -108,11 +126,17 @@ def post():
         key = None
     status = update_user_status(channel_id, session['username'], key)
     if status['banned'] or status['revoked']:
-        return jsonify(status=status)
+        message = {}
+        status_code = 401
+    else:
         message = discord_api.create_message(channel_id, content)
-    return jsonify(message=message, status=status)
+        status_code = messages['code']
+    response = jsonify(message=message.get('content', message), status=status)
+    response.status_code = status_code
+    return response
 
 @api.route("/create_unauthenticated_user", methods=["POST"])
+@rate_limiter.limit("4/hour", key_func=guild_ratelimit_key)
 def create_unauthenticated_user():
     session['unauthenticated'] = True
     username = request.form['username']

titanembeds/blueprints/user/user.py

@@ -6,6 +6,7 @@ from titanembeds.database import db, Guilds, UnauthenticatedUsers, Unauthenticat
 from titanembeds.oauth import authorize_url, token_url, make_authenticated_session, get_current_authenticated_user, get_user_managed_servers, check_user_can_administrate_guild, check_user_permission, generate_avatar_url, generate_guild_icon_url, generate_bot_invite_url
 
 user = Blueprint("user", __name__)
+
 @user.route("/login_authenticated", methods=["GET"])
 def login_authenticated():
     session["redirect"] = request.args.get("redirect")

titanembeds/utils.py

@@ -1,6 +1,7 @@
 from titanembeds.discordrest import DiscordREST
 from flask import request, session
 from flask.ext.cache import Cache
+from flask_limiter import Limiter
 from config import config
 import random
 import string
@@ -15,11 +16,11 @@ def get_client_ipaddr():
         return request.remote_addr
 
 def generate_session_key():
-    sess = session.get("cachestring", None)
+    sess = session.get("sessionunique", None)
     if not sess:
         rand_str = lambda n: ''.join([random.choice(string.lowercase) for i in xrange(n)])
-        session['cachestring'] = rand_str(25)
+        session['sessionunique'] = rand_str(25)
-        sess = session['cachestring']
+        sess = session['sessionunique']
     return sess #Totally unique
 
 def make_cache_key(*args, **kwargs):
@@ -33,3 +34,16 @@ def make_guilds_cache_key():
     sess = generate_session_key()
     ip = get_client_ipaddr()
     return (sess + ip + "user_guilds").encode('utf-8')
+
+def channel_ratelimit_key(): # Generate a bucket with given channel & unique session key
+    sess = generate_session_key()
+    channel_id = request.args.get('channel_id', "0")
+    return (sess + channel_id).encode('utf-8')
+
+def guild_ratelimit_key():
+    sess = generate_session_key()
+    guild_id = request.args.get('guild_id', "0")
+    return (sess + guild_id).encode('utf-8')
+
+
+rate_limiter = Limiter(key_func=get_client_ipaddr) # Default limit by ip address