diff --git a/CTFd/utils/exports/__init__.py b/CTFd/utils/exports/__init__.py index 5092a57..0fc162d 100644 --- a/CTFd/utils/exports/__init__.py +++ b/CTFd/utils/exports/__init__.py @@ -5,12 +5,9 @@ import re import tempfile import zipfile -import datafreeze import dataset import six from alembic.util import CommandError -from datafreeze.format import SERIALIZERS -from datafreeze.format.fjson import JSONEncoder, JSONSerializer from flask import current_app as app from flask_migrate import upgrade from sqlalchemy.exc import OperationalError, ProgrammingError @@ -20,6 +17,7 @@ from CTFd import __version__ as CTFD_VERSION from CTFd.cache import cache from CTFd.models import db, get_class_by_tablename from CTFd.utils import get_app_config, set_config +from CTFd.utils.exports.freeze import freeze_export from CTFd.utils.migrations import ( create_database, drop_database, @@ -29,52 +27,6 @@ from CTFd.utils.migrations import ( from CTFd.utils.uploads import get_uploader -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 - - # Certain databases (MariaDB) store JSON as LONGTEXT. - # Before emitting a file we should standardize to valid JSON (i.e. a dict) - # See Issue #973 - for i, r in enumerate(result["results"]): - data = r.get("requirements") - if data: - try: - if isinstance(data, six.string_types): - result["results"][i]["requirements"] = json.loads(data) - except ValueError: - pass - - 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 export_ctf(): # TODO: For some unknown reason dataset is only able to see alembic_version during tests. # Even using a real sqlite database. This makes this test impossible to pass in sqlite. @@ -89,7 +41,7 @@ def export_ctf(): for table in tables: result = db[table].all() result_file = six.BytesIO() - datafreeze.freeze(result, format="ctfd", fileobj=result_file) + freeze_export(result, fileobj=result_file) result_file.seek(0) backup_zip.writestr("db/{}.json".format(table), result_file.read()) diff --git a/CTFd/utils/exports/encoders.py b/CTFd/utils/exports/encoders.py new file mode 100644 index 0000000..5643b9d --- /dev/null +++ b/CTFd/utils/exports/encoders.py @@ -0,0 +1,11 @@ +import json +from datetime import datetime, date +from decimal import Decimal + + +class JSONEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, (datetime, date)): + return obj.isoformat() + if isinstance(obj, Decimal): + return str(obj) diff --git a/CTFd/utils/exports/freeze.py b/CTFd/utils/exports/freeze.py new file mode 100644 index 0000000..23059ef --- /dev/null +++ b/CTFd/utils/exports/freeze.py @@ -0,0 +1,11 @@ +from CTFd.utils.exports.serializers import JSONSerializer +from sqlalchemy.exc import ProgrammingError, OperationalError + + +def freeze_export(result, fileobj): + try: + query = result + serializer = JSONSerializer(query, fileobj) + serializer.serialize() + except (OperationalError, ProgrammingError) as e: + raise OperationalError("Invalid query: %s" % e) diff --git a/CTFd/utils/exports/serializers.py b/CTFd/utils/exports/serializers.py new file mode 100644 index 0000000..5158097 --- /dev/null +++ b/CTFd/utils/exports/serializers.py @@ -0,0 +1,44 @@ +import json +import six +from collections import OrderedDict + +from CTFd.utils.exports.encoders import JSONEncoder + + +class JSONSerializer(object): + def __init__(self, query, fileobj): + self.query = query + self.fileobj = fileobj + self.buckets = [] + + def serialize(self): + for row in self.query: + self.write(None, row) + self.close() + + def write(self, path, result): + self.buckets.append([result]) + + def wrap(self, result): + result = OrderedDict([("count", len(result)), ("results", result)]) + result["meta"] = {} + return result + + def close(self): + for result in self.buckets: + result = self.wrap(result) + + # Certain databases (MariaDB) store JSON as LONGTEXT. + # Before emitting a file we should standardize to valid JSON (i.e. a dict) + # See Issue #973 + for i, r in enumerate(result["results"]): + data = r.get("requirements") + if data: + try: + if isinstance(data, six.string_types): + result["results"][i]["requirements"] = json.loads(data) + except ValueError: + pass + + data = json.dumps(result, cls=JSONEncoder, indent=2) + self.fileobj.write(data.encode("utf-8")) diff --git a/requirements.txt b/requirements.txt index a4d8229..9881625 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,13 +13,10 @@ itsdangerous==1.1.0 requests>=2.20.0 PyMySQL==0.9.3 gunicorn==19.9.0 -banal==0.4.2 -normality==2.0.0 dataset==1.1.2 mistune==0.8.4 netaddr==0.7.19 redis==3.3.11 -datafreeze==0.1.0 gevent==1.4.0 python-dotenv==0.10.3 flask-restx==0.1.1