Merge pull request #304 from STIXProject/issue303

Make MAECInstance.to_obj() retain all information

stix/extensions/malware/maec_4_1_malware.py

@@ -152,21 +152,18 @@ def from_obj(cls, obj):
             return None
 
         return_obj = cls()
-        
+
         if _MAEC_INSTALLED:
             obj.MAEC = maecPackage.from_obj(obj.MAEC)
         else:
             obj.MAEC = obj.MAEC
-        
+
         super(MAECInstance, cls).from_obj(obj)
 
         return return_obj
-
-    def to_obj(self, return_obj=None, ns_info=None):
-        if not return_obj:
-            return_obj = self._binding_class()
-
-        super(MAECInstance, self).to_obj(ns_info=ns_info)
+
+    def to_obj(self, ns_info=None):
+        return_obj = super(MAECInstance, self).to_obj(ns_info=ns_info)
 
         if mixbox.xml.is_element(self.maec) or mixbox.xml.is_etree(self.maec):
             tree = mixbox.xml.get_etree(self.maec)

stix/test/extensions/malware/maec_4_1_malware_test.py

@@ -1,5 +1,5 @@
 import unittest
-from mixbox.vendor.six import StringIO, BytesIO
+from mixbox.vendor.six import StringIO, BytesIO, text_type
 
 from lxml import etree
 import mixbox.xml
@@ -33,6 +33,15 @@ class PythonMAECTests(EntityTestCase, unittest.TestCase):
         }
     }
 
+    def test_add_name_type(self):
+        maec_malware_instance = MAECInstance()
+        maec_malware_instance.add_name("Poison Ivy Variant v4392-acc")
+        maec_malware_instance.add_type("Remote Access Trojan")
+        maec_xml = text_type(maec_malware_instance.to_xml())
+        self.assertTrue("Poison Ivy Variant v4392-acc" in maec_xml)
+        self.assertTrue("Remote Access Trojan" in maec_xml)
+
+
 class PythonMAECEtreeTests(unittest.TestCase):
     XML = (
         """