From b8c1970b8e82ef1d82377d763980d934eca6727f Mon Sep 17 00:00:00 2001
From: Kevin Chung <kchung@nyu.edu>
Date: Tue, 10 Sep 2019 03:22:01 -0400
Subject: [PATCH] Fix CSV exports in Python 3 by converting StringIO to BytesIO
 (#1107)

* Fix CSV exports in Python 3 by converting StringIO to BytesIO
---
 CTFd/admin/__init__.py         | 11 +++++++++--
 tests/admin/test_export_csv.py | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)
 create mode 100644 tests/admin/test_export_csv.py

diff --git a/CTFd/admin/__init__.py b/CTFd/admin/__init__.py
index 8872fe49..89e64be2 100644
--- a/CTFd/admin/__init__.py
+++ b/CTFd/admin/__init__.py
@@ -125,8 +125,8 @@ def export_csv():
     if model is None:
         abort(404)
 
-    output = six.StringIO()
-    writer = csv.writer(output)
+    temp = six.StringIO()
+    writer = csv.writer(temp)
 
     header = [column.name for column in model.__mapper__.columns]
     writer.writerow(header)
@@ -138,7 +138,14 @@ def export_csv():
             [getattr(curr, column.name) for column in model.__mapper__.columns]
         )
 
+    temp.seek(0)
+
+    # In Python 3 send_file requires bytes
+    output = six.BytesIO()
+    output.write(temp.getvalue().encode("utf-8"))
     output.seek(0)
+    temp.close()
+
     return send_file(
         output,
         as_attachment=True,
diff --git a/tests/admin/test_export_csv.py b/tests/admin/test_export_csv.py
new file mode 100644
index 00000000..371eab93
--- /dev/null
+++ b/tests/admin/test_export_csv.py
@@ -0,0 +1,16 @@
+from tests.helpers import create_ctfd, destroy_ctfd, login_as_user, gen_challenge
+
+
+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)