Add CSV importing feature (#1922)

* Closes #1888 
* Adds code to import CSVs for challenges, users, and teams
This commit is contained in:
Kevin Chung
2021-06-26 18:04:14 -04:00
committed by GitHub
parent 31e8261bad
commit 61507bb12a
6 changed files with 215 additions and 17 deletions

View File

@@ -43,6 +43,7 @@ from CTFd.models import (
) )
from CTFd.utils import config as ctf_config from CTFd.utils import config as ctf_config
from CTFd.utils import get_config, set_config from CTFd.utils import get_config, set_config
from CTFd.utils.csv import load_challenges_csv, load_teams_csv, load_users_csv
from CTFd.utils.decorators import admins_only from CTFd.utils.decorators import admins_only
from CTFd.utils.exports import export_ctf as export_ctf_util from CTFd.utils.exports import export_ctf as export_ctf_util
from CTFd.utils.exports import import_ctf as import_ctf_util from CTFd.utils.exports import import_ctf as import_ctf_util
@@ -116,6 +117,33 @@ def export_ctf():
) )
@admin.route("/admin/import/csv", methods=["POST"])
@admins_only
def import_csv():
csv_type = request.form["csv_type"]
# Try really hard to load data in properly no matter what nonsense Excel gave you
raw = request.files["csv_file"].stream.read()
try:
csvdata = raw.decode("utf-8-sig")
except UnicodeDecodeError:
try:
csvdata = raw.decode("cp1252")
except UnicodeDecodeError:
csvdata = raw.decode("latin-1")
csvfile = StringIO(csvdata)
loaders = {
"challenges": load_challenges_csv,
"users": load_users_csv,
"teams": load_teams_csv,
}
loader = loaders[csv_type]
reader = csv.DictReader(csvfile)
loader(reader)
return redirect(url_for("admin.config"))
@admin.route("/admin/export/csv") @admin.route("/admin/export/csv")
@admins_only @admins_only
def export_csv(): def export_csv():

View File

@@ -1,4 +1,4 @@
from wtforms import BooleanField, SelectField, StringField, TextAreaField from wtforms import BooleanField, FileField, SelectField, StringField, TextAreaField
from wtforms.fields.html5 import IntegerField, URLField from wtforms.fields.html5 import IntegerField, URLField
from wtforms.widgets.html5 import NumberInput from wtforms.widgets.html5 import NumberInput
@@ -83,6 +83,15 @@ class ExportCSVForm(BaseForm):
submit = SubmitField("Download CSV") submit = SubmitField("Download CSV")
class ImportCSVForm(BaseForm):
csv_type = SelectField(
"CSV Type",
choices=[("users", "Users"), ("teams", "Teams"), ("challenges", "Challenges")],
description="Type of CSV data",
)
csv_file = FileField("CSV File", description="CSV file contents")
class LegalSettingsForm(BaseForm): class LegalSettingsForm(BaseForm):
tos_url = URLField( tos_url = URLField(
"Terms of Service URL", "Terms of Service URL",

View File

@@ -10,6 +10,9 @@
<li class="nav-item"> <li class="nav-item">
<a class="nav-link " href="#save-csv" aria-controls="save-csv" role="tab" data-toggle="tab">Download CSV</a> <a class="nav-link " href="#save-csv" aria-controls="save-csv" role="tab" data-toggle="tab">Download CSV</a>
</li> </li>
<li class="nav-item">
<a class="nav-link " href="#import-csv" aria-controls="import-csv" role="tab" data-toggle="tab">Import CSV</a>
</li>
</ul> </ul>
<div class="tab-content"> <div class="tab-content">
<div role="tabpanel" class="tab-pane active" id="export-ctf"> <div role="tabpanel" class="tab-pane active" id="export-ctf">
@@ -49,5 +52,27 @@
</form> </form>
{% endwith %} {% endwith %}
</div> </div>
<div role="tabpanel" class="tab-pane" id="import-csv">
{% with form = Forms.config.ImportCSVForm() %}
<form method="POST" action="{{ url_for('admin.import_csv') }}" enctype="multipart/form-data">
<div class="form-group">
<b>{{ form.csv_type.label }}</b>
{{ form.csv_type(class="form-control custom-select") }}
<small class="form-text text-muted">
{{ form.csv_type.description }}
</small>
</div>
<div class="form-group">
<b>{{ form.csv_file.label }}</b>
{{ form.csv_file(class="form-control-file") }}
<small class="form-text text-muted">
{{ form.csv_file.description }}
</small>
</div>
{{ form.nonce() }}
<input type="submit" class="btn btn-warning" value="Import CSV">
</form>
{% endwith %}
</div>
</div> </div>
</div> </div>

View File

@@ -0,0 +1,59 @@
import json
from CTFd.models import Flags, Hints, Tags, Teams, Users, db
from CTFd.plugins.challenges import get_chal_class
def load_users_csv(dict_reader):
for line in dict_reader:
result = Users(**line)
db.session.add(result)
db.session.commit()
return True
def load_teams_csv(dict_reader):
for line in dict_reader:
result = Teams(**line)
db.session.add(result)
db.session.commit()
return True
def load_challenges_csv(dict_reader):
for line in dict_reader:
flags = line.pop("flags", None)
tags = line.pop("tags", None)
hints = line.pop("hints", None)
challenge_type = line.pop("type", "standard")
# Load in custome type_data
type_data = json.loads(line.pop("type_data", "{}"))
line.update(type_data)
ChallengeClass = get_chal_class(challenge_type)
challenge = ChallengeClass.challenge_model(**line)
db.session.add(challenge)
db.session.commit()
if flags:
flags = [flag.strip() for flag in flags.split(",")]
for flag in flags:
f = Flags(type="static", challenge_id=challenge.id, content=flag,)
db.session.add(f)
db.session.commit()
if tags:
tags = [tag.strip() for tag in tags.split(",")]
for tag in tags:
t = Tags(challenge_id=challenge.id, value=tag,)
db.session.add(t)
db.session.commit()
if hints:
hints = [hint.strip() for hint in hints.split(",")]
for hint in hints:
h = Hints(challenge_id=challenge.id, content=hint,)
db.session.add(h)
db.session.commit()
return True

93
tests/admin/test_csv.py Normal file
View File

@@ -0,0 +1,93 @@
import io
from CTFd.models import Challenges, Teams, Users
from CTFd.utils.crypto import verify_password
from tests.helpers import create_ctfd, destroy_ctfd, gen_challenge, login_as_user
def test_export_csv_works():
"""Test that CSV exports work properly"""
app = create_ctfd()
with app.app_context():
gen_challenge(app.db)
client = login_as_user(app, name="admin", password="password")
csv_data = client.get("/admin/export/csv?table=challenges").get_data(
as_text=True
)
assert len(csv_data) > 0
destroy_ctfd(app)
def test_import_csv_works():
"""Test that CSV imports work properly"""
USERS_CSV = b"""name,email,password
user1,user1@examplectf.com,password
user2,user2@examplectf.com,password"""
TEAMS_CSV = b"""name,email,password
team1,team1@examplectf.com,password
team2,team2@examplectf.com,password"""
CHALLENGES_CSV = b"""name,category,description,value,flags,tags,hints
challenge1,category1,description1,100,"flag1,flag2,flag3","tag1,tag2,tag3","hint1,hint2,hint3"""
app = create_ctfd()
with app.app_context():
client = login_as_user(app, name="admin", password="password")
with client.session_transaction() as sess:
data = {
"csv_type": "users",
"csv_file": (io.BytesIO(USERS_CSV), "users.csv"),
"nonce": sess.get("nonce"),
}
client.post("/admin/import/csv", data=data, content_type="multipart/form-data")
assert Users.query.count() == 3
user = Users.query.filter_by(id=2).first()
assert user.name == "user1"
assert user.email == "user1@examplectf.com"
assert verify_password("password", user.password)
user = Users.query.filter_by(id=3).first()
assert user.name == "user2"
assert user.email == "user2@examplectf.com"
assert verify_password("password", user.password)
with client.session_transaction() as sess:
data = {
"csv_type": "teams",
"csv_file": (io.BytesIO(TEAMS_CSV), "users.csv"),
"nonce": sess.get("nonce"),
}
client.post("/admin/import/csv", data=data, content_type="multipart/form-data")
assert Teams.query.count() == 2
team = Teams.query.filter_by(id=1).first()
assert team.name == "team1"
assert team.email == "team1@examplectf.com"
assert verify_password("password", team.password)
team = Teams.query.filter_by(id=2).first()
assert team.name == "team2"
assert team.email == "team2@examplectf.com"
assert verify_password("password", team.password)
with client.session_transaction() as sess:
data = {
"csv_type": "challenges",
"csv_file": (io.BytesIO(CHALLENGES_CSV), "challenges.csv"),
"nonce": sess.get("nonce"),
}
client.post("/admin/import/csv", data=data, content_type="multipart/form-data")
assert Challenges.query.count() == 1
challenge = Challenges.query.filter_by(id=1).first()
assert challenge.name == "challenge1"
assert challenge.category == "category1"
assert challenge.description == "description1"
assert challenge.value == 100
assert len(challenge.flags) == 3
assert len(challenge.tags) == 3
assert len(challenge.hints) == 3
destroy_ctfd(app)

View File

@@ -1,16 +0,0 @@
from tests.helpers import create_ctfd, destroy_ctfd, gen_challenge, login_as_user
def test_export_csv_works():
"""Test that CSV exports work properly"""
app = create_ctfd()
with app.app_context():
gen_challenge(app.db)
client = login_as_user(app, name="admin", password="password")
csv_data = client.get("/admin/export/csv?table=challenges").get_data(
as_text=True
)
assert len(csv_data) > 0
destroy_ctfd(app)