Fix rate-limiting of flag submission when using team mode (CTFd#977)

* Fix rate-limiting of flag submission when using team mode (Resolves CTFd#975)
* Add tests for rate-limiting of flag submission

Author: hakatashi

CTFd/api/v1/challenges.py

@@ -1,4 +1,4 @@
-from flask import session, request, abort, url_for
+from flask import request, abort, url_for
 from flask_restplus import Namespace, Resource
 from CTFd.models import (
     db,
@@ -376,7 +376,8 @@ def post(self):
         chal_class = get_chal_class(challenge.type)
 
         # Anti-bruteforce / submitting Flags too quickly
-        if current_user.get_wrong_submissions_per_minute(session['id']) > 10:
+        kpm = current_user.get_wrong_submissions_per_minute(user.account_id)
+        if kpm > 10:
             if ctftime():
                 chal_class.fail(
                     user=user,
@@ -388,7 +389,7 @@ def post(self):
                 'submissions',
                 "[{date}] {name} submitted {submission} with kpm {kpm} [TOO FAST]",
                 submission=request_data['submission'].encode('utf-8'),
-                kpm=current_user.get_wrong_submissions_per_minute(session['id'])
+                kpm=kpm
             )
             # Submitting too fast
             return {
@@ -432,8 +433,7 @@ def post(self):
                     'submissions',
                     "[{date}] {name} submitted {submission} with kpm {kpm} [CORRECT]",
                     submission=request_data['submission'].encode('utf-8'),
-                    kpm=current_user.get_wrong_submissions_per_minute(
-                        session['id'])
+                    kpm=kpm
                 )
                 return {
                     'success': True,
@@ -456,8 +456,7 @@ def post(self):
                     'submissions',
                     "[{date}] {name} submitted {submission} with kpm {kpm} [WRONG]",
                     submission=request_data['submission'].encode('utf-8'),
-                    kpm=current_user.get_wrong_submissions_per_minute(
-                        session['id'])
+                    kpm=kpm
                 )
 
                 if max_tries:
@@ -491,9 +490,7 @@ def post(self):
                 'submissions',
                 "[{date}] {name} submitted {submission} with kpm {kpm} [ALREADY SOLVED]",
                 submission=request_data['submission'].encode('utf-8'),
-                kpm=current_user.get_wrong_submissions_per_minute(
-                    user.account_id
-                )
+                kpm=kpm
             )
             return {
                 'success': True,

tests/api/v1/test_challenges.py

@@ -11,7 +11,8 @@
                            gen_flag,
                            gen_user,
                            gen_team,
-                           gen_solve)
+                           gen_solve,
+                           gen_fail)
 from freezegun import freeze_time
 
 
@@ -308,19 +309,56 @@ def test_api_challenge_attempt_post_private():
     """Can an private user post /api/v1/challenges/attempt"""
     app = create_ctfd()
     with app.app_context():
-        gen_challenge(app.db)
-        gen_flag(app.db, 1)
+        challenge_id = gen_challenge(app.db).id
+        gen_flag(app.db, challenge_id)
         register_user(app)
         with login_as_user(app) as client:
-            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": 1, "submission": "wrong_flag"})
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "wrong_flag"})
             assert r.status_code == 200
             assert r.get_json()['data']['status'] == 'incorrect'
-            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": 1, "submission": "flag"})
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "flag"})
             assert r.status_code == 200
             assert r.get_json()['data']['status'] == 'correct'
-            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": 1, "submission": "flag"})
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "flag"})
             assert r.status_code == 200
             assert r.get_json()['data']['status'] == 'already_solved'
+        challenge_id = gen_challenge(app.db).id
+        gen_flag(app.db, challenge_id)
+        with login_as_user(app) as client:
+            for i in range(10):
+                gen_fail(app.db, user_id=2, challenge_id=challenge_id)
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "flag"})
+            assert r.status_code == 429
+            assert r.get_json()['data']['status'] == 'ratelimited'
+    destroy_ctfd(app)
+
+    app = create_ctfd(user_mode="teams")
+    with app.app_context():
+        challenge_id = gen_challenge(app.db).id
+        gen_flag(app.db, challenge_id)
+        register_user(app)
+        team_id = gen_team(app.db).id
+        user = Users.query.filter_by(id=2).first()
+        user.team_id = team_id
+        app.db.session.commit()
+        with login_as_user(app) as client:
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "wrong_flag"})
+            assert r.status_code == 200
+            assert r.get_json()['data']['status'] == 'incorrect'
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "flag"})
+            assert r.status_code == 200
+            assert r.get_json()['data']['status'] == 'correct'
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "flag"})
+            assert r.status_code == 200
+            assert r.get_json()['data']['status'] == 'already_solved'
+        challenge_id = gen_challenge(app.db).id
+        gen_flag(app.db, challenge_id)
+        with login_as_user(app) as client:
+            for i in range(10):
+                gen_fail(app.db, user_id=2, team_id=team_id, challenge_id=challenge_id)
+            r = client.post('/api/v1/challenges/attempt', json={"challenge_id": challenge_id, "submission": "flag"})
+            assert r.status_code == 429
+            assert r.get_json()['data']['status'] == 'ratelimited'
     destroy_ctfd(app)