diff --git a/CTFd/admin/__init__.py b/CTFd/admin/__init__.py index 2ff77721..dbd590a1 100644 --- a/CTFd/admin/__init__.py +++ b/CTFd/admin/__init__.py @@ -1,7 +1,7 @@ import csv import datetime import os -from io import BytesIO, StringIO +from io import StringIO from flask import Blueprint, abort from flask import current_app as app @@ -39,11 +39,10 @@ from CTFd.models import ( Unlocks, Users, db, - get_class_by_tablename, ) from CTFd.utils import config as ctf_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.csv import dump_csv, load_challenges_csv, load_teams_csv, load_users_csv from CTFd.utils.decorators import admins_only from CTFd.utils.exports import export_ctf as export_ctf_util from CTFd.utils.exports import import_ctf as import_ctf_util @@ -149,31 +148,7 @@ def import_csv(): def export_csv(): table = request.args.get("table") - # TODO: It might make sense to limit dumpable tables. Config could potentially leak sensitive information. - model = get_class_by_tablename(table) - if model is None: - abort(404) - - temp = StringIO() - writer = csv.writer(temp) - - header = [column.name for column in model.__mapper__.columns] - writer.writerow(header) - - responses = model.query.all() - - for curr in responses: - writer.writerow( - [getattr(curr, column.name) for column in model.__mapper__.columns] - ) - - temp.seek(0) - - # In Python 3 send_file requires bytes - output = BytesIO() - output.write(temp.getvalue().encode("utf-8")) - output.seek(0) - temp.close() + output = dump_csv(name=table) return send_file( output, diff --git a/CTFd/forms/config.py b/CTFd/forms/config.py index a9c1929a..9a7940e7 100644 --- a/CTFd/forms/config.py +++ b/CTFd/forms/config.py @@ -4,7 +4,7 @@ from wtforms.widgets.html5 import NumberInput from CTFd.forms import BaseForm from CTFd.forms.fields import SubmitField -from CTFd.models import db +from CTFd.utils.csv import get_dumpable_tables class ResetInstanceForm(BaseForm): @@ -74,12 +74,7 @@ class AccountSettingsForm(BaseForm): class ExportCSVForm(BaseForm): - table = SelectField( - "Database Table", - choices=list( - zip(sorted(db.metadata.tables.keys()), sorted(db.metadata.tables.keys())) - ), - ) + table = SelectField("Database Table", choices=get_dumpable_tables()) submit = SubmitField("Download CSV") diff --git a/CTFd/utils/csv/__init__.py b/CTFd/utils/csv/__init__.py index ac6cc1ff..d4f15a66 100644 --- a/CTFd/utils/csv/__init__.py +++ b/CTFd/utils/csv/__init__.py @@ -1,7 +1,232 @@ +import csv import json +from io import BytesIO, StringIO -from CTFd.models import Flags, Hints, Tags, Teams, Users, db +from CTFd.models import ( + Flags, + Hints, + Tags, + TeamFields, + Teams, + UserFields, + Users, + db, + get_class_by_tablename, +) from CTFd.plugins.challenges import get_chal_class +from CTFd.utils.config import is_teams_mode, is_users_mode +from CTFd.utils.scores import get_standings + + +def get_dumpable_tables(): + csv_keys = list(CSV_KEYS.keys()) + db_keys = list(db.metadata.tables.keys()) + tables = csv_keys + db_keys + table_keys = list(zip(tables, tables)) + return table_keys + + +def dump_csv(name): + dump_func = CSV_KEYS.get(name) + if dump_func: + return dump_func() + elif get_class_by_tablename(name): + return dump_database_table(tablename=name) + else: + raise KeyError + + +def dump_scoreboard_csv(): + # TODO: Add fields to scoreboard data + temp = StringIO() + writer = csv.writer(temp) + + standings = get_standings() + + # Get all user fields in a specific order + user_fields = UserFields.query.all() + user_field_ids = [f.id for f in user_fields] + user_field_names = [f.name for f in user_fields] + + if is_teams_mode(): + team_fields = TeamFields.query.all() + team_field_ids = [f.id for f in team_fields] + team_field_names = [f.name for f in team_fields] + + header = ( + [ + "place", + "team", + "team id", + "score", + "member name", + "member id", + "member email", + "member score", + ] + + user_field_names + + team_field_names + ) + writer.writerow(header) + + for i, standing in enumerate(standings): + team = Teams.query.filter_by(id=standing.account_id).first() + + # Build field entries using the order of the field values + team_field_entries = {f.field_id: f.value for f in team.field_entries} + team_field_values = [ + team_field_entries.get(f_id, "") for f_id in team_field_ids + ] + team_row = [ + i + 1, + team.name, + team.id, + standing.score, + "", + "", + ] + team_field_values + + writer.writerow(team_row) + + for member in team.members: + user_field_entries = {f.field_id: f.value for f in member.field_entries} + user_field_values = [ + user_field_entries.get(f_id, "") for f_id in user_field_ids + ] + user_row = [ + "", + "", + "", + "", + member.name, + member.id, + member.email, + member.score, + ] + user_field_values + writer.writerow(user_row) + elif is_users_mode(): + header = ["place", "user", "score"] + user_field_names + writer.writerow(header) + + for i, standing in enumerate(standings): + user = Users.query.filter_by(id=standing.account_id).first() + + # Build field entries using the order of the field values + user_field_entries = {f.field_id: f.value for f in user.field_entries} + user_field_values = [ + user_field_entries.get(f_id, "") for f_id in user_field_ids + ] + user_row = [i + 1, user.name, standing.score] + user_field_values + writer.writerow(user_row) + + # In Python 3 send_file requires bytes + output = BytesIO() + output.write(temp.getvalue().encode("utf-8")) + output.seek(0) + temp.close() + + return output + + +def dump_users_with_fields_csv(): + temp = StringIO() + writer = csv.writer(temp) + + user_fields = UserFields.query.all() + user_field_ids = [f.id for f in user_fields] + user_field_names = [f.name for f in user_fields] + + header = [column.name for column in Users.__mapper__.columns] + user_field_names + writer.writerow(header) + + responses = Users.query.all() + + for curr in responses: + user_field_entries = {f.field_id: f.value for f in curr.field_entries} + user_field_values = [ + user_field_entries.get(f_id, "") for f_id in user_field_ids + ] + user_row = [ + getattr(curr, column.name) for column in Users.__mapper__.columns + ] + user_field_values + writer.writerow(user_row) + + temp.seek(0) + + # In Python 3 send_file requires bytes + output = BytesIO() + output.write(temp.getvalue().encode("utf-8")) + output.seek(0) + temp.close() + + return output + + +def dump_teams_with_fields_csv(): + temp = StringIO() + writer = csv.writer(temp) + + team_fields = TeamFields.query.all() + team_field_ids = [f.id for f in team_fields] + team_field_names = [f.name for f in team_fields] + + header = [column.name for column in Teams.__mapper__.columns] + team_field_names + writer.writerow(header) + + responses = Teams.query.all() + + for curr in responses: + team_field_entries = {f.field_id: f.value for f in curr.field_entries} + team_field_values = [ + team_field_entries.get(f_id, "") for f_id in team_field_ids + ] + + team_row = [ + getattr(curr, column.name) for column in Teams.__mapper__.columns + ] + team_field_values + + writer.writerow(team_row) + + temp.seek(0) + + # In Python 3 send_file requires bytes + output = BytesIO() + output.write(temp.getvalue().encode("utf-8")) + output.seek(0) + temp.close() + + return output + + +def dump_database_table(tablename): + # TODO: It might make sense to limit dumpable tables. Config could potentially leak sensitive information. + model = get_class_by_tablename(tablename) + + if model is None: + raise KeyError("Unknown database table") + + temp = StringIO() + writer = csv.writer(temp) + + header = [column.name for column in model.__mapper__.columns] + writer.writerow(header) + + responses = model.query.all() + + for curr in responses: + writer.writerow( + [getattr(curr, column.name) for column in model.__mapper__.columns] + ) + + temp.seek(0) + + # In Python 3 send_file requires bytes + output = BytesIO() + output.write(temp.getvalue().encode("utf-8")) + output.seek(0) + temp.close() + + return output def load_users_csv(dict_reader): @@ -57,3 +282,10 @@ def load_challenges_csv(dict_reader): db.session.add(h) db.session.commit() return True + + +CSV_KEYS = { + "scoreboard": dump_scoreboard_csv, + "users+fields": dump_users_with_fields_csv, + "teams+fields": dump_teams_with_fields_csv, +}