<commit_before>import stix
import stix.utils.parser
import stix.ttp.malware_instance
from stix.ttp.malware_instance import MalwareInstance
import stix.bindings.extensions.malware.maec_4_1 as ext_binding
from lxml import etree
from StringIO import StringIO

try:
    from maec.package.package import Package as maecPackage
    maec_installed = True
except ImportError:
    maec_installed = False

class MAECInstance(MalwareInstance):
    _binding = ext_binding
    _binding_class = _binding.MAEC4_1InstanceType
    _namespace = 'http://stix.mitre.org/extensions/Malware#MAEC4.1-1'
    _xml_ns_prefix = "stix-maec"
    _XSI_TYPE = "stix-maec:MAEC4.1InstanceType"

    def __init__(self, maec=None):
        super(MAECInstance, self).__init__()
        self.maec = maec
        
    @property
    def maec(self):
        return self._maec
    
    @maec.setter
    def maec(self, value):
        if value is None:
            self._maec = None
            return
        elif isinstance(value, etree._ElementTree):
            tree = value
        elif isinstance(value, etree._Element):
            tree = etree.ElementTree(value)
        elif maec_installed and isinstance(value, maecPackage):
            self._maec = value
        else:
            raise ValueError('maec must be instance of lxml.etree._Element '
                             'or lxml.etree._ElementTree'
                             'or maec.package.package.Package')
        
        if maec_installed and isinstance(self._maec, maecPackage):
            # Handle the namespace/schemaLocation processing for the MAEC content
            nsset = self._maec._get_namespaces()
            self.__input_namespaces__ = {}
            self.__input_schemalocations__ = {}
            for ns in nsset:
                if ns:
                    self.__input_namespaces__[ns.name] = ns.prefix
                    if ns.schema_location:
                        self.__input_schemalocations__[ns.name] = ns.schema_location
            # Add the MAEC Vocabs ns/schemalocation by default
            self.__input_namespaces__['http://maec.mitre.org/default_vocabularies-1'] = 'maecVocabs'
            self.__input_schemalocations__['http://maec.mitre.org/default_vocabularies-1'] = 'http://maec.mitre.org/language/version4.1/maec_default_vocabularies.xsd'
        else:
            root = tree.getroot()
            expected_node_tag = "{%s}MAEC" % (self._namespace)
            if root.tag != expected_node_tag:
                ns_maec = "http://maec.mitre.org/XMLSchema/maec-package-2"
                node_ns = etree.QName(root).namespace
                if node_ns == ns_maec:
                    # attempt to cast
                    etree.register_namespace(self._xml_ns_prefix, self._namespace)
                    root.tag = expected_node_tag
                else:
                    raise ValueError("Cannot set maec property. Expected tag %s found %s" 
                                     % (expected_node_tag, root.tag))
        
            self.__input_namespaces__ = {}
            for alias,ns in root.nsmap.iteritems():
                self.__input_namespaces__[ns] = alias

            try:
                schemaloc = stix.utils.parser.get_schemaloc_pairs(root)
                self.__input_schemalocations__ = dict(schemaloc)
            except KeyError:
                self.__input_schemalocations__ = {}
                
            self._maec = tree
        
    @classmethod
    def from_obj(cls, obj, return_obj=None):
        if not obj:
            return None
        if not return_obj:
            return_obj = cls()
        
        super(MAECInstance, cls).from_obj(obj, return_obj)

        if maec_installed:
            return_obj.maec = maecPackage.from_obj(obj.MAEC)
        else:
            return_obj.maec = obj.MAEC

        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(return_obj=return_obj, ns_info=ns_info)

        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 from_dict(cls, d, return_obj=None):
        if not d:
            return None
        if not return_obj:
            return_obj = cls()
            
        super(MAECInstance, cls).from_dict(d, return_obj)

        if 'maec' in d:
            maec = d['maec']

            if isinstance(maec, dict):
                if not maec_installed:
                    raise ValueError(
                        "Unable to parse 'maec' value in dictionary. Please "
                        "install python-maec to parse dictionary value."
                    )

                return_obj.maec = maecPackage.from_dict(maec)
            else:
                parser = stix.utils.parser.get_xml_parser()
                return_obj.maec = etree.parse(StringIO(maec), parser=parser)
        
        return return_obj


    def to_dict(self):
        d = super(MAECInstance, self).to_dict()

        if self.maec:
            if maec_installed and isinstance(self.maec, maecPackage):
                d['maec'] = self.maec.to_dict()
            else:
                d['maec'] = etree.tostring(self.maec)

        return d

stix.ttp.malware_instance.add_extension(MAECInstance)<commit_msg>Refactored maec extension code layout.<commit_after>import stix
import stix.ttp.malware_instance
from stix.ttp.malware_instance import MalwareInstance
import stix.bindings.extensions.malware.maec_4_1 as ext_binding
from lxml import etree
from StringIO import StringIO

try:
    from maec.package.package import Package as maecPackage
    maec_installed = True
except ImportError:
    maec_installed = False

class MAECInstance(MalwareInstance):
    _binding = ext_binding
    _binding_class = _binding.MAEC4_1InstanceType
    _namespace = 'http://stix.mitre.org/extensions/Malware#MAEC4.1-1'
    _xml_ns_prefix = "stix-maec"
    _XSI_TYPE = "stix-maec:MAEC4.1InstanceType"

    def __init__(self, maec=None):
        super(MAECInstance, self).__init__()
        self.maec = maec
        
    @property
    def maec(self):
        return self._maec
    
    @maec.setter
    def maec(self, value):
        if value is None:
            self._maec = None
        elif maec_installed and isinstance(value, maecPackage):
            self._parse_maec_package(value)
            self._maec = value
        elif isinstance(value, (etree._ElementTree, etree._Element)):
            root = stix.utils.parser.get_etree_root(value)
            self._parse_etree(root)
            self._maec = etree.ElementTree(root)
        else:
            raise ValueError(
                'maec must be instance of lxml.etree._Element or '
                'lxml.etree._ElementTree or maec.package.package.Package'
            )

    def _parse_etree(self, root):
        ns_maec = "http://maec.mitre.org/XMLSchema/maec-package-2"
        node_ns = etree.QName(root).namespace
        expected_node_tag = "{%s}MAEC" % (self._namespace)
        node_tag = root.tag

        if node_tag != expected_node_tag:
            if node_ns != ns_maec:
                raise ValueError(
                    "Cannot set maec property. Expected tag %s but found %s" %
                    (expected_node_tag, node_tag)
                )

            # attempt to cast
            etree.register_namespace(self._xml_ns_prefix, self._namespace)
            root.tag = expected_node_tag

        self.__input_namespaces__ = dict(
            (ns, alias) for alias, ns in root.nsmap.iteritems()
        )

        try:
            schemaloc = stix.utils.parser.get_schemaloc_pairs(root)
            self.__input_schemalocations__ = dict(schemaloc)
        except KeyError:
            self.__input_schemalocations__ = {}


    def _parse_maec_package(self, package):
        # Handle the namespace/schemaLocation processing for the MAEC content
        nsset = package._get_namespaces()

        ns_dict = dict(
            (ns.name, ns.prefix) for ns in nsset if ns
        )
        schemaloc_dict = dict(
            (ns.name, ns.schema_location) for ns in nsset if ns and ns.schema_location
        )

        # Add the MAEC Vocabs ns/schemalocation by default
        vocab_ns = 'http://maec.mitre.org/default_vocabularies-1'
        vocab_schemaloc = 'http://maec.mitre.org/language/version4.1/maec_default_vocabularies.xsd'

        ns_dict[vocab_ns] = 'maecVocabs'
        schemaloc_dict[vocab_ns] = vocab_schemaloc

        self.__input_namespaces__ = ns_dict
        self.__input_schemalocations__ = schemaloc_dict

    @classmethod
    def from_obj(cls, obj, return_obj=None):
        if not obj:
            return None
        if not return_obj:
            return_obj = cls()
        
        super(MAECInstance, cls).from_obj(obj, return_obj)

        if maec_installed:
            return_obj.maec = maecPackage.from_obj(obj.MAEC)
        else:
            return_obj.maec = obj.MAEC

        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(return_obj=return_obj, ns_info=ns_info)

        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 from_dict(cls, d, return_obj=None):
        if not d:
            return None
        if not return_obj:
            return_obj = cls()
            
        super(MAECInstance, cls).from_dict(d, return_obj)

        if 'maec' in d:
            maec = d['maec']

            if isinstance(maec, dict):
                if not maec_installed:
                    raise ValueError(
                        "Unable to parse 'maec' value in dictionary. Please "
                        "install python-maec to parse dictionary value."
                    )

                return_obj.maec = maecPackage.from_dict(maec)
            else:
                parser = stix.utils.parser.get_xml_parser()
                return_obj.maec = etree.parse(StringIO(maec), parser=parser)
        
        return return_obj


    def to_dict(self):
        d = super(MAECInstance, self).to_dict()

        if self.maec:
            if maec_installed and isinstance(self.maec, maecPackage):
                d['maec'] = self.maec.to_dict()
            else:
                d['maec'] = etree.tostring(self.maec)

        return d

stix.ttp.malware_instance.add_extension(MAECInstance)
{'fraction_non_alphanumeric': 0.08315584776368182, 'fraction_numerical': 0.0025002315029169367, 'mean_word_length': 1.906350914962325, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 12}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2640186', 'n_tokens_mistral': 3325, 'n_tokens_neox': 3238, 'n_words': 657}