Make MAECInstance extension fully compatible with TypedFields. closes STIXProject#339

Author: emmanvg

stix/bindings/extensions/malware/maec_4_1.py

@@ -15,23 +15,18 @@
 from stix.bindings import register_extension
 import stix.bindings.ttp as ttp_binding
 
-try:
-    from maec.bindings.maec_package import PackageType
-    maec_installed = True
-except ImportError:
-    PackageType = None
-    maec_installed = False
 
 XML_NS = "http://stix.mitre.org/extensions/Malware#MAEC4.1-1"
 
 #
 # Data representation classes.
 #
 
+
 @register_extension
 class MAEC4_1InstanceType(ttp_binding.MalwareInstanceType):
     """The MAEC4.1InstanceType provides an extension to ttp_binding.MalwareInstanceType
-    which imports and leverages the MAEC 4.0.1 schema for structured
+    which imports and leverages the MAEC 4.1 schema for structured
     characterization of Malware."""
     subclass = None
     superclass = ttp_binding.MalwareInstanceType
@@ -89,11 +84,7 @@ def exportChildren(self, lwrite, level, nsmap, namespace_=XML_NS, name_='MAEC4.1
         else:
             eol_ = ''
         if self.MAEC is not None:
-            if maec_installed and isinstance(self.MAEC, PackageType):
-                self.MAEC.export(lwrite, level, namespace_='stix-maec:', name_='MAEC', pretty_print=pretty_print)
-            else:
-                showIndent(lwrite, level, pretty_print)
-                lwrite(etree_.tostring(self.MAEC, pretty_print=pretty_print).decode())
+            self.MAEC.export(lwrite, level, namespace_='stix-maec:', name_='MAEC', pretty_print=pretty_print)
     def build(self, node):
         self.__sourcenode__ = node
         already_processed = set()
@@ -105,12 +96,11 @@ def buildAttributes(self, node, attrs, already_processed):
         super(MAEC4_1InstanceType, self).buildAttributes(node, attrs, already_processed)
     def buildChildren(self, child_, node, nodeName_, fromsubclass_=False):
         if nodeName_ == 'MAEC':
-            if maec_installed:
-                obj_ = PackageType.factory()
-                obj_.build(child_)
-                self.set_MAEC(obj_)
-            else:
-                self.set_MAEC(child_)
+            # Fails hard if maec library is not installed in your Python environment.
+            from maec.bindings.maec_package import PackageType
+            obj_ = PackageType.factory()
+            obj_.build(child_)
+            self.set_MAEC(obj_)
         super(MAEC4_1InstanceType, self).buildChildren(child_, node, nodeName_, True)
 # end class MAEC4_1InstanceType
 

stix/extensions/malware/maec_4_1_malware.py

@@ -1,231 +1,37 @@
 # Copyright (c) 2017, The MITRE Corporation. All rights reserved.
 # See LICENSE.txt for complete terms.
 
-# stdlib
-from distutils.version import LooseVersion
-
 # external
-from lxml import etree
-import mixbox.xml
-from mixbox.vendor.six import BytesIO, iteritems, binary_type
+from mixbox import fields
+
 
 # internal
 import stix
-import stix.utils as utils
-import stix.ttp.malware_instance
+from stix.bindings.extensions.malware import maec_4_1 as maec_instance_binding
 from stix.ttp.malware_instance import MalwareInstance
-import stix.bindings.extensions.malware.maec_4_1 as ext_binding
-from mixbox import fields
-from stix.bindings.extensions.malware.maec_4_1 import maec_installed
-from lxml.etree import _ElementTree
-
-_MIN_PYTHON_MAEC_VERSION = '4.1.0.12'
-
-
-class UnsupportedVersion(Exception):
-    def __init__(self, message, expected, found):
-        super(UnsupportedVersion, self).__init__(message)
-        self.expected = expected
-        self.found = found
 
 
-def _check_maec_version():
-    """Checks that the installed python-maec has a version greater than or
-    equal to the minimum supported version.
-
-    Note:
-        We do this rather than having a python-maec dependency requirement
-        listed in setup.py because MAEC is used as an extension to STIX and
-        not a core component to STIX (like CybOX).
-
-    Raises:
-        ImportError: If python-maec is not installed.
-        UnsupportedVersion: If the python-maec installation does not satisfy
-            the version requirements.
-
+@stix.register_extension
+class MAECInstance(MalwareInstance):
     """
-    import maec
-
-    found = maec.__version__
-    expected = _MIN_PYTHON_MAEC_VERSION
-
-    if LooseVersion(found) >= LooseVersion(expected):
-        return
+    The MAECInstance object provides an extension to the MalwareInstanceType
+    which imports and leverages the MAEC 4.1 schema for structured
+    characterization of Malware.
 
-    fmt = ("Unsupported python-maec version installed: '%s'. Minimum version "
-           "is '%s'.")
-    error = fmt % (found, expected)
-    raise UnsupportedVersion(error, expected=expected, found=found)
-
-
-try:
-    # Check that the correct version of python-maec is installed.
-    _check_maec_version()
-
-    # Import maecPackage into global space
-    from maec.package.package import Package as maecPackage
-
-    _MAEC_INSTALLED = True
-except ImportError:
-    maecPackage, Package = None, None
-    _MAEC_INSTALLED = False
-
-
-def is_maec(obj):
-    """Checks if the input object is python-maec object.
-
-    Returns:
-        True if python-maec is ins
+    This class extension is automatically registered by the
+    MalwareInstanceFactory.
 
+    Warnings:
+        Interacting with the ``maec`` field will fail if the maec library is
+        not installed in your Python environment.
     """
-    if not _MAEC_INSTALLED:
-        return False
-
-    return isinstance(obj, maecPackage)
-
-def validate_maec_input(instance, value):
-    if value is None:
-        return
-    elif _MAEC_INSTALLED and is_maec(value):
-        return
-    elif mixbox.xml.is_element(value) or mixbox.xml.is_etree(value):
-        return
-    else:
-        error = (
-            "Cannot set maec to '{0}'. Expected 'lxml.etree._Element' or "
-            "'maec.package.package.Package'."
-        )
-        error = error.format(type(value))
-        raise ValueError(error)
-
-@stix.register_extension
-class MAECInstance(MalwareInstance):
-    _binding = ext_binding
+    _binding = maec_instance_binding
     _binding_class = _binding.MAEC4_1InstanceType
-    _namespace = 'http://stix.mitre.org/extensions/Malware#MAEC4.1-1'
-    _xml_ns_prefix = "stix-maec"
+    _namespace = "http://stix.mitre.org/extensions/Malware#MAEC4.1-1"
     _XSI_TYPE = "stix-maec:MAEC4.1InstanceType"
-    _TAG_MAEC = "{%s}MAEC" % _namespace
 
-    maec = fields.TypedField("MAEC", preset_hook=validate_maec_input)
+    maec = fields.TypedField("MAEC", type_="maec.package.package.Package")
 
     def __init__(self, maec=None):
         super(MAECInstance, self).__init__()
-        self.__input_namespaces__ = {}
-        self.__input_schemalocations__ = {}
         self.maec = maec
-
-    def _parse_etree(self, root):
-        node_tag = root.tag
-
-        if node_tag != self._TAG_MAEC:
-            self._cast_maec(root)
-
-        self._collect_namespaces(root)
-        self._collect_schemalocs(root)
-
-    def _cast_maec(self, node):
-        ns_maec = "http://maec.mitre.org/XMLSchema/maec-package-2"
-        node_ns = etree.QName(node).namespace
-
-        if node_ns == ns_maec:
-            etree.register_namespace(self._xml_ns_prefix, self._namespace)
-            node.tag = self._TAG_MAEC
-        else:
-            error = "Cannot set maec. Expected tag '{0}' found '{1}'."
-            error = error.format(self._TAG_MAEC, node.tag)
-            raise ValueError(error)
-
-    def _collect_schemalocs(self, node):
-        try:
-            schemaloc = mixbox.xml.get_schemaloc_pairs(node)
-            self.__input_schemalocations__ = dict(schemaloc)
-        except KeyError:
-            self.__input_schemalocations__ = {}
-
-    def _collect_namespaces(self, node):
-        self.__input_namespaces__ = dict(iteritems(node.nsmap))
-
-    @classmethod
-    def from_obj(cls, obj):
-        if not obj:
-            return None
-
-        return_obj = cls()
-
-        if _MAEC_INSTALLED:
-            obj.MAEC = maecPackage.from_obj(obj.MAEC)
-        else:
-            obj.MAEC = obj.MAEC
-
-        return_obj = super(MAECInstance, cls).from_obj(obj)
-
-        return return_obj
-
-    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)
-            root = mixbox.xml.get_etree_root(tree)
-            self._parse_etree(root)
-            self.maec = root
-
-        if _MAEC_INSTALLED and isinstance(self.maec, maecPackage):
-            return_obj.MAEC = self.maec.to_obj(ns_info=ns_info)
-        else:
-            return_obj.MAEC = self.maec
-
-        return return_obj
-
-    @classmethod
-    def _maec_from_dict(cls, d):
-        if _MAEC_INSTALLED:
-            return maecPackage.from_dict(d)
-
-        raise ValueError(
-            "Unable to parse 'maec' value in dictionary. Please "
-            "install python-maec to parse dictionary value."
-        )
-
-    @classmethod
-    def from_dict(cls, d, return_obj=None):
-        if not d:
-            return None
-
-        d = d.copy()
-
-        maec = d.get('maec')
-
-        if maec is None:
-            pass
-        elif isinstance(maec, dict):
-            d['maec'] = cls._maec_from_dict(maec)
-        elif isinstance(maec, binary_type):
-            d['maec'] = mixbox.xml.get_etree_root(BytesIO(maec))
-        else:
-            raise TypeError("Unknown type for 'maec' entry.")
-
-        return_obj = super(MAECInstance, cls).from_dict(d)
-
-        return return_obj
-
-    def to_dict(self):
-        d = super(MAECInstance, self).to_dict()
-
-        if self.maec is not None:
-            if mixbox.xml.is_element(self.maec) or mixbox.xml.is_etree(self.maec):
-                tree = mixbox.xml.get_etree(self.maec)
-                root = mixbox.xml.get_etree_root(tree)
-                self._parse_etree(root)
-                self.maec = root
-
-            if _MAEC_INSTALLED and isinstance(self.maec, maecPackage):
-                d['maec'] = self.maec.to_dict()
-            else:
-                d['maec'] = etree.tostring(self.maec)
-
-            if self._XSI_TYPE:
-                d['xsi:type'] = self._XSI_TYPE
-
-        return d