Update requirements (#406)

* Updating to use dataset and datafreeze
     * Use a new datafreeze serializer to get around Python 3 issues. 
* Update requirements.txt
* Add simple test for export_ctf()

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()
+            result_file = six.BytesIO()
-            dataset.freeze(result, format='json', fileobj=result_file)
+            datafreeze.freeze(result, format='ctfd', fileobj=result_file)
             result_file.seek(0)
             backup_zip.writestr('db/{}.json'.format(item), result_file.read())
 

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-Caching==1.3.3
-Flask-Migrate==2.0.4
+Flask-Migrate==2.1.1
-SQLAlchemy==1.1.11
+SQLAlchemy==1.1.14
-SQLAlchemy-Utils==0.32.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
+gunicorn==19.7.1
-dataset==0.8.0
+dataset==1.0.2
 mistune==0.7.4
 netaddr==0.7.19
 redis==2.10.6
+datafreeze==0.1.0

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'])
 

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)