From 06d86e1046e9d507f6e958b4a424784844e3a0e7 Mon Sep 17 00:00:00 2001 From: Jeremy Zhang Date: Sun, 26 Mar 2017 16:37:27 -0700 Subject: [PATCH] Add a bit of rate limiting --- titanembeds/app.py | 3 ++- titanembeds/blueprints/api/api.py | 36 ++++++++++++++++++++++++----- titanembeds/blueprints/user/user.py | 1 + titanembeds/utils.py | 20 +++++++++++++--- 4 files changed, 50 insertions(+), 10 deletions(-) diff --git a/titanembeds/app.py b/titanembeds/app.py index 5be6132..e4c3573 100644 --- a/titanembeds/app.py +++ b/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") diff --git a/titanembeds/blueprints/api/api.py b/titanembeds/blueprints/api/api.py index 951f3c4..1042337 100644 --- a/titanembeds/blueprints/api/api.py +++ b/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 = discord_api.create_message(channel_id, content) - return jsonify(message=message, status=status) + message = {} + status_code = 401 + else: + message = discord_api.create_message(channel_id, content) + 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'] diff --git a/titanembeds/blueprints/user/user.py b/titanembeds/blueprints/user/user.py index 80bdc4b..49720ac 100644 --- a/titanembeds/blueprints/user/user.py +++ b/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") diff --git a/titanembeds/utils.py b/titanembeds/utils.py index d31e49d..9452221 100644 --- a/titanembeds/utils.py +++ b/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) - sess = session['cachestring'] + session['sessionunique'] = rand_str(25) + 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