Merge pull request #84 from mwielgoszewski/fix-csv-python-3

Fix csv node download in python 3

doorman/compat.py

@@ -18,12 +18,6 @@
     basestring = (str, bytes)
 
 
-try:
-    from StringIO import StringIO
-except ImportError:
-    from io import StringIO
-
-
 def with_metaclass(meta, *bases):
     """Create a base class with a metaclass."""
     class metaclass(meta):

doorman/manage/views.py

@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
+from io import BytesIO
 from operator import itemgetter
-import csv
 import json
 import datetime as dt
+import unicodecsv as csv
 
 from flask import (
-    Blueprint, current_app, flash, jsonify, make_response, redirect,
-    render_template, request, url_for
+    Blueprint, current_app, flash, jsonify, redirect, render_template,
+    request, send_file, url_for
 )
 from flask_login import login_required
 from flask_paginate import Pagination
@@ -26,7 +27,6 @@
     UpdateRuleForm,
     UpdateNodeForm,
 )
-from doorman.compat import StringIO
 from doorman.database import db
 from doorman.models import (
     DistributedQuery, DistributedQueryTask, DistributedQueryResult,
@@ -112,10 +112,10 @@ def nodes_csv():
     column_names = map(itemgetter(0), current_app.config['DOORMAN_CAPTURE_NODE_INFO'])
     labels = map(itemgetter(1), current_app.config['DOORMAN_CAPTURE_NODE_INFO'])
     headers.extend(labels)
-    headers = map(str.title, headers)
+    headers = list(map(str.title, headers))
 
-    sio = StringIO()
-    writer = csv.writer(sio)
+    bio = BytesIO()
+    writer = csv.writer(bio)
     writer.writerow(headers)
 
     for node in Node.query:
@@ -130,9 +130,15 @@ def nodes_csv():
         row.extend([node.node_info.get(column, '') for column in column_names])
         writer.writerow(row)
 
-    response = make_response(sio.getvalue())
-    response.headers["Content-Disposition"] = "attachment; filename=nodes.csv"
-    response.headers["Content-Type"] = "text/csv"
+    bio.seek(0)
+
+    response = send_file(
+        bio,
+        mimetype='text/csv',
+        as_attachment=True,
+        attachment_filename='nodes.csv'
+    )
+
     return response
 
 

requirements/dev.txt

@@ -53,6 +53,7 @@ scales==1.0.9
 six==1.10.0
 SQLAlchemy==1.0.12
 waitress==0.9.0
+unicodecsv==0.14.1
 webassets==0.11.1
 WebOb==1.6.0
 WebTest==2.0.21

requirements/prod.txt

@@ -41,6 +41,7 @@ requests==2.10.0
 requests-oauthlib==0.6.1
 scales==1.0.9
 SQLAlchemy==1.0.12
+unicodecsv==0.14.1
 webassets==0.11.1
 Werkzeug==0.11.8
 WTForms==2.1

setup.py

@@ -6,7 +6,7 @@
     name='doorman',
     description='an osquery fleet manager',
     url='https://github.com/mwielgoszewski/doorman',
-    version='0.5',
+    version='0.5.1',
     packages=find_packages(
         exclude=[
             'tests*',

tests/test_functional.py

@@ -1538,3 +1538,30 @@ def test_node_info_not_updated_on_erroneous_data(self, node, testapp):
 
         learn_from_result(result, node.to_dict())
         assert 'foobar' not in node.node_info
+
+
+class TestCSVExport:
+    def test_node_csv_download(self, node, testapp):
+        import unicodecsv as csv
+
+        node.enrolled_on = dt.datetime.utcnow()
+        node.last_checkin = dt.datetime.utcnow()
+        node.last_ip = '1.1.1.1'
+        node.node_info = {'hardware_vendor': "Honest Achmed's Computer Supply"}
+        node.save()
+
+        resp = testapp.get(url_for('manage.nodes_csv'))
+
+        assert resp.headers['Content-Type'] == 'text/csv; charset=utf-8'
+        assert resp.headers['Content-Disposition'] == 'attachment; filename=nodes.csv'
+
+        reader = csv.DictReader(io.BytesIO(resp.body))
+        row = next(reader)
+
+        assert row['Display Name'] == node.display_name
+        assert row['Host Identifier'] == node.host_identifier
+        assert row['Enrolled On'] == str(node.enrolled_on)
+        assert row['Last Check-In'] == str(node.last_checkin)
+        assert row['Last Ip Address'] == node.last_ip
+        assert row['Is Active'] == 'True'
+        assert row['Make'] == node.node_info['hardware_vendor']