Skip to content

Commit d2f8b40

Browse files
ColdHeatColdHeat
authored
Adds ondelete='CASCADE' to some models. (CTFd#979)
* Fixes `populate.py` to assign captains to teams. * Adds `ondelete='CASCADE'` to most ForeignKeys in models * Closes CTFd#794 * Test reset in team mode to test removing teams with captains * Test deleting users/teams with awards to test cascading deletion * `gen_team()` test helper now creates users for the team and assigns the first one as captain * Added `Challenges.flags` relationship and moved the `Flags.challenge` relationship to a backref on `Challenges`
1 parent 6fcf143 commit d2f8b40

9 files changed

+334
-64
lines changed

‎CTFd/models/__init__.py

+11-12
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ class Challenges(db.Model):
7676
files = db.relationship("ChallengeFiles", backref="challenge")
7777
tags = db.relationship("Tags", backref="challenge")
7878
hints = db.relationship("Hints", backref="challenge")
79+
flags = db.relationship("Flags", backref="challenge")
7980

8081
__mapper_args__ = {
8182
'polymorphic_identity': 'standard',
@@ -93,7 +94,7 @@ class Hints(db.Model):
9394
__tablename__ = 'hints'
9495
id = db.Column(db.Integer, primary_key=True)
9596
type = db.Column(db.String(80), default='standard')
96-
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id'))
97+
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id', ondelete='CASCADE'))
9798
content = db.Column(db.Text)
9899
cost = db.Column(db.Integer, default=0)
99100
requirements = db.Column(db.JSON)
@@ -125,8 +126,8 @@ def __repr__(self):
125126
class Awards(db.Model):
126127
__tablename__ = 'awards'
127128
id = db.Column(db.Integer, primary_key=True)
128-
user_id = db.Column(db.Integer, db.ForeignKey('users.id'))
129-
team_id = db.Column(db.Integer, db.ForeignKey('teams.id'))
129+
user_id = db.Column(db.Integer, db.ForeignKey('users.id', ondelete='CASCADE'))
130+
team_id = db.Column(db.Integer, db.ForeignKey('teams.id', ondelete='CASCADE'))
130131
type = db.Column(db.String(80), default='standard')
131132
name = db.Column(db.String(80))
132133
description = db.Column(db.Text)
@@ -162,7 +163,7 @@ def __repr__(self):
162163
class Tags(db.Model):
163164
__tablename__ = 'tags'
164165
id = db.Column(db.Integer, primary_key=True)
165-
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id'))
166+
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id', ondelete='CASCADE'))
166167
value = db.Column(db.String(80))
167168

168169
def __init__(self, *args, **kwargs):
@@ -191,7 +192,7 @@ class ChallengeFiles(Files):
191192
__mapper_args__ = {
192193
'polymorphic_identity': 'challenge'
193194
}
194-
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id'))
195+
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id', ondelete='CASCADE'))
195196

196197
def __init__(self, *args, **kwargs):
197198
super(ChallengeFiles, self).__init__(**kwargs)
@@ -210,13 +211,11 @@ def __init__(self, *args, **kwargs):
210211
class Flags(db.Model):
211212
__tablename__ = 'flags'
212213
id = db.Column(db.Integer, primary_key=True)
213-
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id'))
214+
challenge_id = db.Column(db.Integer, db.ForeignKey('challenges.id', ondelete='CASCADE'))
214215
type = db.Column(db.String(80))
215216
content = db.Column(db.Text)
216217
data = db.Column(db.Text)
217218

218-
challenge = db.relationship('Challenges', foreign_keys="Flags.challenge_id", lazy='select')
219-
220219
__mapper_args__ = {
221220
'polymorphic_on': type
222221
}
@@ -454,7 +453,7 @@ class Teams(db.Model):
454453
banned = db.Column(db.Boolean, default=False)
455454

456455
# Relationship for Users
457-
captain_id = db.Column(db.Integer, db.ForeignKey('users.id'))
456+
captain_id = db.Column(db.Integer, db.ForeignKey('users.id', ondelete='SET NULL'))
458457
captain = db.relationship("Users", foreign_keys=[captain_id])
459458

460459
created = db.Column(db.DateTime, default=datetime.datetime.utcnow)
@@ -682,8 +681,8 @@ class Fails(Submissions):
682681
class Unlocks(db.Model):
683682
__tablename__ = 'unlocks'
684683
id = db.Column(db.Integer, primary_key=True)
685-
user_id = db.Column(db.Integer, db.ForeignKey('users.id'))
686-
team_id = db.Column(db.Integer, db.ForeignKey('teams.id'))
684+
user_id = db.Column(db.Integer, db.ForeignKey('users.id', ondelete='CASCADE'))
685+
team_id = db.Column(db.Integer, db.ForeignKey('teams.id', ondelete='CASCADE'))
687686
target = db.Column(db.Integer)
688687
date = db.Column(db.DateTime, default=datetime.datetime.utcnow)
689688
type = db.Column(db.String(32))
@@ -715,7 +714,7 @@ class Tracking(db.Model):
715714
id = db.Column(db.Integer, primary_key=True)
716715
type = db.Column(db.String(32))
717716
ip = db.Column(db.String(46))
718-
user_id = db.Column(db.Integer, db.ForeignKey('users.id'))
717+
user_id = db.Column(db.Integer, db.ForeignKey('users.id', ondelete='CASCADE'))
719718
date = db.Column(db.DateTime, default=datetime.datetime.utcnow)
720719

721720
user = db.relationship('Users', foreign_keys="Tracking.user_id", lazy='select')

‎migrations/versions/b295b033364d_add_ondelete_cascade_to_foreign_keys.py

+140
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
"""Add ondelete cascade to foreign keys
2+
3+
Revision ID: b295b033364d
4+
Revises: b5551cd26764
5+
Create Date: 2019-05-03 19:26:57.746887
6+
7+
"""
8+
from alembic import op
9+
import sqlalchemy as sa
10+
from sqlalchemy.dialects import mysql
11+
12+
# revision identifiers, used by Alembic.
13+
revision = 'b295b033364d'
14+
down_revision = 'b5551cd26764'
15+
branch_labels = None
16+
depends_on = None
17+
18+
19+
def upgrade():
20+
bind = op.get_bind()
21+
url = str(bind.engine.url)
22+
if url.startswith('mysql'):
23+
op.drop_constraint('awards_ibfk_1', 'awards', type_='foreignkey')
24+
op.drop_constraint('awards_ibfk_2', 'awards', type_='foreignkey')
25+
op.create_foreign_key('awards_ibfk_1', 'awards', 'teams', ['team_id'], ['id'], ondelete='CASCADE')
26+
op.create_foreign_key('awards_ibfk_2', 'awards', 'users', ['user_id'], ['id'], ondelete='CASCADE')
27+
28+
op.drop_constraint('files_ibfk_1', 'files', type_='foreignkey')
29+
op.create_foreign_key('files_ibfk_1', 'files', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
30+
31+
op.drop_constraint('flags_ibfk_1', 'flags', type_='foreignkey')
32+
op.create_foreign_key('flags_ibfk_1', 'flags', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
33+
34+
op.drop_constraint('hints_ibfk_1', 'hints', type_='foreignkey')
35+
op.create_foreign_key('hints_ibfk_1', 'hints', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
36+
37+
op.drop_constraint('tags_ibfk_1', 'tags', type_='foreignkey')
38+
op.create_foreign_key('tags_ibfk_1', 'tags', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
39+
40+
op.drop_constraint('team_captain_id', 'teams', type_='foreignkey')
41+
op.create_foreign_key('team_captain_id', 'teams', 'users', ['captain_id'], ['id'], ondelete='SET NULL')
42+
43+
op.drop_constraint('tracking_ibfk_1', 'tracking', type_='foreignkey')
44+
op.create_foreign_key('tracking_ibfk_1', 'tracking', 'users', ['user_id'], ['id'], ondelete='CASCADE')
45+
46+
op.drop_constraint('unlocks_ibfk_1', 'unlocks', type_='foreignkey')
47+
op.drop_constraint('unlocks_ibfk_2', 'unlocks', type_='foreignkey')
48+
op.create_foreign_key('unlocks_ibfk_1', 'unlocks', 'teams', ['team_id'], ['id'], ondelete='CASCADE')
49+
op.create_foreign_key('unlocks_ibfk_2', 'unlocks', 'users', ['user_id'], ['id'], ondelete='CASCADE')
50+
elif url.startswith('postgres'):
51+
op.drop_constraint('awards_team_id_fkey', 'awards', type_='foreignkey')
52+
op.drop_constraint('awards_user_id_fkey', 'awards', type_='foreignkey')
53+
op.create_foreign_key('awards_team_id_fkey', 'awards', 'teams', ['team_id'], ['id'], ondelete='CASCADE')
54+
op.create_foreign_key('awards_user_id_fkey', 'awards', 'users', ['user_id'], ['id'], ondelete='CASCADE')
55+
56+
op.drop_constraint('files_challenge_id_fkey', 'files', type_='foreignkey')
57+
op.create_foreign_key('files_challenge_id_fkey', 'files', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
58+
59+
op.drop_constraint('flags_challenge_id_fkey', 'flags', type_='foreignkey')
60+
op.create_foreign_key('flags_challenge_id_fkey', 'flags', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
61+
62+
op.drop_constraint('hints_challenge_id_fkey', 'hints', type_='foreignkey')
63+
op.create_foreign_key('hints_challenge_id_fkey', 'hints', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
64+
65+
op.drop_constraint('tags_challenge_id_fkey', 'tags', type_='foreignkey')
66+
op.create_foreign_key('tags_challenge_id_fkey', 'tags', 'challenges', ['challenge_id'], ['id'], ondelete='CASCADE')
67+
68+
op.drop_constraint('team_captain_id', 'teams', type_='foreignkey')
69+
op.create_foreign_key('team_captain_id', 'teams', 'users', ['captain_id'], ['id'], ondelete='SET NULL')
70+
71+
op.drop_constraint('tracking_user_id_fkey', 'tracking', type_='foreignkey')
72+
op.create_foreign_key('tracking_user_id_fkey', 'tracking', 'users', ['user_id'], ['id'], ondelete='CASCADE')
73+
74+
op.drop_constraint('unlocks_team_id_fkey', 'unlocks', type_='foreignkey')
75+
op.drop_constraint('unlocks_user_id_fkey', 'unlocks', type_='foreignkey')
76+
op.create_foreign_key('unlocks_team_id_fkey', 'unlocks', 'teams', ['team_id'], ['id'], ondelete='CASCADE')
77+
op.create_foreign_key('unlocks_user_id_fkey', 'unlocks', 'users', ['user_id'], ['id'], ondelete='CASCADE')
78+
79+
80+
def downgrade():
81+
bind = op.get_bind()
82+
url = str(bind.engine.url)
83+
if url.startswith('mysql'):
84+
op.drop_constraint('unlocks_ibfk_1', 'unlocks', type_='foreignkey')
85+
op.drop_constraint('unlocks_ibfk_2', 'unlocks', type_='foreignkey')
86+
op.create_foreign_key('unlocks_ibfk_1', 'unlocks', 'teams', ['team_id'], ['id'])
87+
op.create_foreign_key('unlocks_ibfk_2', 'unlocks', 'users', ['user_id'], ['id'])
88+
89+
op.drop_constraint('tracking_ibfk_1', 'tracking', type_='foreignkey')
90+
op.create_foreign_key('tracking_ibfk_1', 'tracking', 'users', ['user_id'], ['id'])
91+
92+
op.drop_constraint('team_captain_id', 'teams', type_='foreignkey')
93+
op.create_foreign_key('team_captain_id', 'teams', 'users', ['captain_id'], ['id'])
94+
95+
op.drop_constraint('tags_ibfk_1', 'tags', type_='foreignkey')
96+
op.create_foreign_key('tags_ibfk_1', 'tags', 'challenges', ['challenge_id'], ['id'])
97+
98+
op.drop_constraint('hints_ibfk_1', 'hints', type_='foreignkey')
99+
op.create_foreign_key('hints_ibfk_1', 'hints', 'challenges', ['challenge_id'], ['id'])
100+
101+
op.drop_constraint('flags_ibfk_1', 'flags', type_='foreignkey')
102+
op.create_foreign_key('flags_ibfk_1', 'flags', 'challenges', ['challenge_id'], ['id'])
103+
104+
op.drop_constraint('files_ibfk_1', 'files', type_='foreignkey')
105+
op.create_foreign_key('files_ibfk_1', 'files', 'challenges', ['challenge_id'], ['id'])
106+
107+
op.drop_constraint('awards_ibfk_1', 'awards', type_='foreignkey')
108+
op.drop_constraint('awards_ibfk_2', 'awards', type_='foreignkey')
109+
op.create_foreign_key('awards_ibfk_1', 'awards', 'teams', ['team_id'], ['id'])
110+
op.create_foreign_key('awards_ibfk_2', 'awards', 'users', ['user_id'], ['id'])
111+
elif url.startswith('postgres'):
112+
op.drop_constraint('unlocks_team_id_fkey', 'unlocks', type_='foreignkey')
113+
op.drop_constraint('unlocks_user_id_fkey', 'unlocks', type_='foreignkey')
114+
op.create_foreign_key('unlocks_team_id_fkey', 'unlocks', 'teams', ['team_id'], ['id'])
115+
op.create_foreign_key('unlocks_user_id_fkey', 'unlocks', 'users', ['user_id'], ['id'])
116+
117+
op.drop_constraint('tracking_user_id_fkey', 'tracking', type_='foreignkey')
118+
op.create_foreign_key('tracking_user_id_fkey', 'tracking', 'users', ['user_id'], ['id'])
119+
120+
op.drop_constraint('team_captain_id', 'teams', type_='foreignkey')
121+
op.create_foreign_key('team_captain_id', 'teams', 'users', ['captain_id'], ['id'])
122+
123+
op.drop_constraint('tags_challenge_id_fkey', 'tags', type_='foreignkey')
124+
op.create_foreign_key('tags_challenge_id_fkey', 'tags', 'challenges', ['challenge_id'], ['id'])
125+
126+
op.drop_constraint('hints_challenge_id_fkey', 'hints', type_='foreignkey')
127+
op.create_foreign_key('hints_challenge_id_fkey', 'hints', 'challenges', ['challenge_id'], ['id'])
128+
129+
op.drop_constraint('flags_challenge_id_fkey', 'flags', type_='foreignkey')
130+
op.create_foreign_key('flags_challenge_id_fkey', 'flags', 'challenges', ['challenge_id'], ['id'])
131+
132+
op.drop_constraint('files_challenge_id_fkey', 'files', type_='foreignkey')
133+
op.create_foreign_key('files_challenge_id_fkey', 'files', 'challenges', ['challenge_id'], ['id'])
134+
135+
op.drop_constraint('awards_team_id_fkey', 'awards', type_='foreignkey')
136+
op.drop_constraint('awards_user_id_fkey', 'awards', type_='foreignkey')
137+
op.create_foreign_key('awards_team_id_fkey', 'awards', 'teams', ['team_id'], ['id'])
138+
op.create_foreign_key('awards_user_id_fkey', 'awards', 'users', ['user_id'], ['id'])
139+
140+

‎populate.py

+10
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,16 @@ def random_date(start, end):
304304

305305
db.session.commit()
306306

307+
if mode == 'teams':
308+
# Assign Team Captains
309+
print("GENERATING TEAM CAPTAINS")
310+
teams = Teams.query.all()
311+
for team in teams:
312+
captain = Users.query.filter_by(team_id=team.id).order_by(Users.id).limit(1).first()
313+
if captain:
314+
team.captain_id = captain.id
315+
db.session.commit()
316+
307317
# Generating Solves
308318
print("GENERATING SOLVES")
309319
if mode == 'users':

‎tests/admin/test_config.py

+60-12
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,18 @@
1-
from CTFd.models import Users, Challenges, Fails, Solves, Tracking
2-
from tests.helpers import (create_ctfd,
3-
destroy_ctfd,
4-
register_user,
5-
login_as_user,
6-
gen_challenge,
7-
gen_award,
8-
gen_flag,
9-
gen_user,
10-
gen_solve,
11-
gen_fail,
12-
gen_tracking)
1+
from CTFd.models import Users, Teams, Challenges, Fails, Solves, Tracking
2+
from tests.helpers import (
3+
create_ctfd,
4+
destroy_ctfd,
5+
register_user,
6+
login_as_user,
7+
gen_challenge,
8+
gen_award,
9+
gen_flag,
10+
gen_user,
11+
gen_team,
12+
gen_solve,
13+
gen_fail,
14+
gen_tracking
15+
)
1316
import random
1417

1518

@@ -49,3 +52,48 @@ def test_reset():
4952
assert Fails.query.count() == 0
5053
assert Tracking.query.count() == 0
5154
destroy_ctfd(app)
55+
56+
57+
def test_reset_team_mode():
58+
app = create_ctfd(user_mode="teams")
59+
with app.app_context():
60+
base_user = 'user'
61+
base_team = 'team'
62+
63+
for x in range(10):
64+
chal = gen_challenge(app.db, name='chal_name{}'.format(x))
65+
gen_flag(app.db, challenge_id=chal.id, content='flag')
66+
67+
for x in range(10):
68+
user = base_user + str(x)
69+
user_email = user + "@ctfd.io"
70+
user_obj = gen_user(app.db, name=user, email=user_email)
71+
team_obj = gen_team(app.db, name=base_team + str(x), email=base_team + str(x) + '@ctfd.io')
72+
team_obj.members.append(user_obj)
73+
team_obj.captain_id = user_obj.id
74+
app.db.session.commit()
75+
gen_award(app.db, user_id=user_obj.id)
76+
gen_solve(app.db, user_id=user_obj.id, challenge_id=random.randint(1, 10))
77+
gen_fail(app.db, user_id=user_obj.id, challenge_id=random.randint(1, 10))
78+
gen_tracking(app.db, user_id=user_obj.id)
79+
80+
assert Teams.query.count() == 10
81+
assert Users.query.count() == 51 # 10 random users, 40 users (10 teams * 4), 1 admin user
82+
assert Challenges.query.count() == 10
83+
84+
register_user(app)
85+
client = login_as_user(app, name="admin", password="password")
86+
87+
with client.session_transaction() as sess:
88+
data = {
89+
"nonce": sess.get('nonce')
90+
}
91+
client.post('/admin/reset', data=data)
92+
93+
assert Teams.query.count() == 0
94+
assert Users.query.count() == 0
95+
assert Challenges.query.count() == 10
96+
assert Solves.query.count() == 0
97+
assert Fails.query.count() == 0
98+
assert Tracking.query.count() == 0
99+
destroy_ctfd(app)

‎tests/api/v1/teams/test_team_members.py

+10-12
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,10 @@ def test_api_team_get_members():
1414
"""Can a user get /api/v1/teams/<team_id>/members only if admin"""
1515
app = create_ctfd(user_mode="teams")
1616
with app.app_context():
17-
user = gen_user(app.db)
18-
team = gen_team(app.db)
19-
team.members.append(user)
20-
user.team_id = team.id
17+
gen_team(app.db)
2118
app.db.session.commit()
19+
20+
gen_user(app.db, name="user_name")
2221
with login_as_user(app, name="user_name") as client:
2322
r = client.get('/api/v1/teams/1/members', json="")
2423
assert r.status_code == 403
@@ -28,22 +27,20 @@ def test_api_team_get_members():
2827
assert r.status_code == 200
2928

3029
resp = r.get_json()
31-
assert resp['data'] == [2]
30+
# The following data is sorted b/c in Postgres data isn't necessarily returned ordered.
31+
assert sorted(resp['data']) == sorted([2, 3, 4, 5])
3232
destroy_ctfd(app)
3333

3434

3535
def test_api_team_remove_members():
3636
"""Can a user remove /api/v1/teams/<team_id>/members only if admin"""
3737
app = create_ctfd(user_mode="teams")
3838
with app.app_context():
39-
user1 = gen_user(app.db, name="user1", email="[email protected]") # ID 2
40-
user2 = gen_user(app.db, name="user2", email="[email protected]") # ID 3
4139
team = gen_team(app.db)
42-
team.members.append(user1)
43-
team.members.append(user2)
44-
user1.team_id = team.id
45-
user2.team_id = team.id
40+
assert len(team.members) == 4
4641
app.db.session.commit()
42+
43+
gen_user(app.db, name='user1')
4744
with login_as_user(app, name="user1") as client:
4845
r = client.delete('/api/v1/teams/1/members', json={
4946
'id': 2
@@ -57,7 +54,8 @@ def test_api_team_remove_members():
5754
assert r.status_code == 200
5855

5956
resp = r.get_json()
60-
assert resp['data'] == [3]
57+
# The following data is sorted b/c in Postgres data isn't necessarily returned ordered.
58+
assert sorted(resp['data']) == sorted([3, 4, 5])
6159

6260
r = client.delete('/api/v1/teams/1/members', json={
6361
'id': 2

0 commit comments

Comments
 (0)