In this chapter, we'll use csv.DictWriter instead of csv.writer, just like we did in Chapter 5, Databases in Python, and Chapter 6, Extracting Artifacts from Binary Files. As a reminder, the difference is that the DictWriter writes dictionary objects to a CSV file and the csv.writer function is more suited for writing lists.
The great thing about csv.DictWriter is that it requires an argument, fieldnames, when creating the writer object. The fieldnames argument should be a list that represents the desired order of columns in the output. In addition, all possible keys must be included in the fieldnames list. If a key exists that isn't contained in the list, an exception will be raised. On the other hand, if a key isn't present in the dictionary but is in the fieldnames list, then that column will simply be skipped for that entry:
001 from __future__ import print_function
002 import sys
003 import os
004 if sys.version_info[0] == 2:
005 import unicodecsv as csv
006 elif sys.version_info[0] == 3:
007 import csv
008 import logging
... 040 def csv_writer(output_data, headers, output_dir, output_name):
041 """
042 The csv_writer function uses the csv DictWriter module to
043 write the list of dictionaries. The DictWriter can take
044 a fieldnames argument, as a list, which represents the
045 desired order of columns.
046 :param output_data: The list of dictionaries containing
047 embedded metadata.
048 :param headers: A list of keys in the dictionary that
049 represent the desired order of columns in the output.
050 :param output_dir: The folder to write the output CSV to.
051 :param output_name: The name of the output CSV.
052 :return:
053 """
054 msg = 'Writing ' + output_name + ' CSV output.'
055 print('[+]', msg)
056 logging.info(msg)
057
058 out_file = os.path.join(output_dir, output_name)
059
060 if sys.version_info[0] == 2:
061 csvfile = open(out_file, "wb")
062 elif sys.version_info[0] == 3:
063 csvfile = open(out_file, "w", newline='',
064 encoding='utf-8')
On line 69, we create our csv.DictWriter function, passing in the output file and the headers as a list of fieldnames from our plugin function. To write the headers for our CSV file, we can simply call the writeheader function, which uses the fieldnames list as its list of headers. Finally, we need to iterate through each dictionary in our metadata container list and write them using the writerow() function in line 76, as follows:
066 with csvfile:
067 # We use DictWriter instead of Writer to write
068 # dictionaries to CSV.
069 writer = csv.DictWriter(csvfile, fieldnames=headers)
070
071 # Writerheader writes the header based on the supplied
072 # headers object
073 writer.writeheader()
074 for dictionary in output_data:
075 if dictionary:
076 writer.writerow(dictionary)