diff --git a/CTFd/utils.py b/CTFd/utils.py index 2fe5539..7f8747d 100644 --- a/CTFd/utils.py +++ b/CTFd/utils.py @@ -17,6 +17,7 @@ import sys import tempfile import time import dataset +import datafreeze import zipfile import io @@ -31,6 +32,9 @@ from werkzeug.utils import secure_filename from CTFd.models import db, WrongKeys, Pages, Config, Tracking, Teams, Files, ip2long, long2ip +from datafreeze.format import SERIALIZERS +from datafreeze.format.fjson import JSONSerializer, JSONEncoder + if six.PY2: text_type = unicode binary_type = str @@ -45,6 +49,39 @@ plugin_scripts = [] plugin_stylesheets = [] +class CTFdSerializer(JSONSerializer): + """ + Slightly modified datafreeze serializer so that we can properly + export the CTFd database into a zip file. + """ + def close(self): + for path, result in self.buckets.items(): + result = self.wrap(result) + + if self.fileobj is None: + fh = open(path, 'wb') + else: + fh = self.fileobj + + data = json.dumps(result, + cls=JSONEncoder, + indent=self.export.get_int('indent')) + + callback = self.export.get('callback') + if callback: + data = "%s && %s(%s);" % (callback, callback, data) + + if six.PY3: + fh.write(bytes(data, encoding='utf-8')) + else: + fh.write(data) + if self.fileobj is None: + fh.close() + + +SERIALIZERS['ctfd'] = CTFdSerializer # Load the custom serializer + + def init_logs(app): logger_keys = logging.getLogger('keys') logger_logins = logging.getLogger('logins') @@ -627,15 +664,16 @@ def export_ctf(segments=None): } # Backup database - backup = io.BytesIO() + backup = six.BytesIO() + backup_zip = zipfile.ZipFile(backup, 'w') for segment in segments: group = groups[segment] for item in group: result = db[item].all() - result_file = io.BytesIO() - dataset.freeze(result, format='json', fileobj=result_file) + result_file = six.BytesIO() + datafreeze.freeze(result, format='ctfd', fileobj=result_file) result_file.seek(0) backup_zip.writestr('db/{}.json'.format(item), result_file.read()) diff --git a/requirements.txt b/requirements.txt index dd4ffbe..578beda 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,18 +1,19 @@ Flask==0.12.2 -Flask-SQLAlchemy==2.2 +Flask-SQLAlchemy==2.3.1 Flask-Session==0.3.1 -Flask-Caching==1.2.0 -Flask-Migrate==2.0.4 -SQLAlchemy==1.1.11 -SQLAlchemy-Utils==0.32.14 +Flask-Caching==1.3.3 +Flask-Migrate==2.1.1 +SQLAlchemy==1.1.14 +SQLAlchemy-Utils>=0.32.17 passlib==1.7.1 bcrypt==3.1.3 -six==1.10.0 +six==1.11.0 itsdangerous==0.24 requests==2.18.1 PyMySQL==0.7.11 -gunicorn==19.7.0 -dataset==0.8.0 +gunicorn==19.7.1 +dataset==1.0.2 mistune==0.7.4 netaddr==0.7.19 redis==2.10.6 +datafreeze==0.1.0 diff --git a/tests/helpers.py b/tests/helpers.py index 68adf26..7987b22 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -4,6 +4,7 @@ from sqlalchemy_utils import database_exists, create_database, drop_database from sqlalchemy.engine.url import make_url import datetime import six +import gc if six.PY2: text_type = unicode @@ -42,6 +43,7 @@ def destroy_ctfd(app): with app.app_context(): app.db.session.commit() app.db.session.close_all() + gc.collect() # Garbage collect (necessary in the case of dataset freezes to clean database connections) app.db.drop_all() drop_database(app.config['SQLALCHEMY_DATABASE_URI']) diff --git a/tests/test_utils.py b/tests/test_utils.py index c15c20e..c8f0cdb 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,7 +3,7 @@ from tests.helpers import * from CTFd.models import ip2long, long2ip -from CTFd.utils import get_config, set_config, override_template, sendmail, verify_email, ctf_started, ctf_ended +from CTFd.utils import get_config, set_config, override_template, sendmail, verify_email, ctf_started, ctf_ended, export_ctf from CTFd.utils import register_plugin_script, register_plugin_stylesheet from CTFd.utils import base64encode, base64decode from freezegun import freeze_time @@ -324,3 +324,28 @@ def test_register_plugin_stylesheet(): assert '/fake/stylesheet/path.css' in output assert 'http://ctfd.io/fake/stylesheet/path.css' in output destroy_ctfd(app) + + +def test_export_ctf(): + """Test that CTFd can properly export the database""" + app = create_ctfd() + with app.app_context(): + register_user(app) + chal = gen_challenge(app.db, name=text_type('🐺')) + chal_id = chal.id + hint = gen_hint(app.db, chal_id) + + client = login_as_user(app) + with client.session_transaction() as sess: + data = { + "nonce": sess.get('nonce') + } + r = client.post('/hints/1', data=data) + output = r.get_data(as_text=True) + output = json.loads(output) + app.db.session.commit() + backup = export_ctf() + backup.seek(0) + with open('export.zip', 'wb') as f: + f.write(backup.getvalue()) + destroy_ctfd(app)