Merge pull request MAECProject#57 from MAECProject/output_options

ikiril01 · ikiril01 · commit 8bb305e01672 · 2014-11-18T13:42:06.000-05:00
Output Options
diff --git a/maec/__init__.py b/maec/__init__.py
@@ -46,11 +46,21 @@ def to_xml_file(self, filename, namespace_dict=None, custom_header=None):
         elif not namespace_dict and hasattr(self, '__input_namespaces__'):
             namespace_dict = self.__input_namespaces__
         out_file  = open(filename, 'w')
-        if custom_header:
+        out_file.write("<?xml version='1.0' encoding='UTF-8'?>\n")
+        # Write out the custom header, if included
+        if isinstance(custom_header, list):
             for line in custom_header:
                 out_file.write[line]
-        else:
-            out_file.write("<?xml version='1.0' encoding='UTF-8'?>\n")
+        elif isinstance(custom_header, dict):
+            out_file.write("<!--\n")
+            for key, value in custom_header.iteritems():
+                sanitized_key = str(key).replace("-->", "\\-\\->")
+                sanitized_value = str(value).replace("-->", "\\-\\->")
+                out_file.write(sanitized_key + ": " + sanitized_value + "\n")
+            out_file.write("-->\n")
+        elif isinstance(custom_header, basestring):
+            out_file.write(custom_header)
+            
         self.to_obj().export(out_file.write, 0, namespacedef_ = self._get_namespace_def(namespace_dict))
         out_file.close()
 
diff --git a/maec/misc/options.py b/maec/misc/options.py
@@ -19,4 +19,11 @@ class ScriptOptions(object):
     def __init__(self):
         self.deduplicate_bundles = False
         self.dereference_bundles = False
-        self.normalize_bundles = False
+        self.normalize_bundles = False
+        
+    def to_dict(self):
+        return {
+            "deduplicate_bundles": self.deduplicate_bundles,
+            "dereference_bundles": self.dereference_bundles,
+            "normalize_bundles": self.normalize_bundles
+        }
