Fix CSV exports in Python 3 by converting StringIO to BytesIO (#1107)

* Fix CSV exports in Python 3 by converting StringIO to BytesIO
This commit is contained in:
Kevin Chung
2019-09-10 03:22:01 -04:00
committed by GitHub
parent 607c517d28
commit b8c1970b8e
2 changed files with 25 additions and 2 deletions

View File

@@ -125,8 +125,8 @@ def export_csv():
if model is None: if model is None:
abort(404) abort(404)
output = six.StringIO() temp = six.StringIO()
writer = csv.writer(output) writer = csv.writer(temp)
header = [column.name for column in model.__mapper__.columns] header = [column.name for column in model.__mapper__.columns]
writer.writerow(header) writer.writerow(header)
@@ -138,7 +138,14 @@ def export_csv():
[getattr(curr, column.name) for column in model.__mapper__.columns] [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) output.seek(0)
temp.close()
return send_file( return send_file(
output, output,
as_attachment=True, as_attachment=True,

View File

@@ -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)