1# Copyright (c) 2017, The MITRE Corporation. All rights reserved. 2# See LICENSE.txt for complete terms. 3 4# external 5from lxml import etree 6import mixbox.xml 7from mixbox.fields import TypedField 8from mixbox.vendor.six import BytesIO, iteritems 9 10# internal 11import stix 12from stix.indicator.test_mechanism import _BaseTestMechanism 13import stix.bindings.extensions.test_mechanism.open_ioc_2010 as open_ioc_tm_binding 14 15 16@stix.register_extension 17class OpenIOCTestMechanism(_BaseTestMechanism): 18 _namespace = "http://stix.mitre.org/extensions/TestMechanism#OpenIOC2010-1" 19 _binding = open_ioc_tm_binding 20 _binding_class = _binding.OpenIOC2010TestMechanismType 21 _xml_ns_prefix = "stix-openioc" 22 _XSI_TYPE = "stix-openioc:OpenIOC2010TestMechanismType" 23 _TAG_IOC = "{%s}ioc" % _namespace 24 25 ioc = TypedField("ioc") 26 27 def __init__(self, id_=None, idref=None): 28 super(OpenIOCTestMechanism, self).__init__(id_=id_, idref=idref) 29 self.ioc = None 30 self.__input_namespaces__ = {} 31 self.__input_schemalocations__ = {} 32 33 def _collect_schemalocs(self, node): 34 try: 35 schemaloc = mixbox.xml.get_schemaloc_pairs(node) 36 self.__input_schemalocations__ = dict(schemaloc) 37 except KeyError: 38 self.__input_schemalocations__ = {} 39 40 def _collect_namespaces(self, node): 41 self.__input_namespaces__ = dict(iteritems(node.nsmap)) 42 43 def _cast_ioc(self, node): 44 ns_ioc = "http://schemas.mandiant.com/2010/ioc" 45 node_ns = etree.QName(node).namespace 46 47 if node_ns == ns_ioc: 48 etree.register_namespace(self._xml_ns_prefix, self._namespace) 49 node.tag = self._TAG_IOC 50 else: 51 error = "Cannot set ioc. Expected tag '{0}' found '{1}'." 52 error = error.format(self._TAG_IOC, node.tag) 53 raise ValueError(error) 54 55 def _processed_ioc(self): 56 if self.ioc is None: 57 return None 58 59 tree = mixbox.xml.get_etree(self.ioc) 60 root = mixbox.xml.get_etree_root(tree) 61 62 if root.tag != self._TAG_IOC: 63 self._cast_ioc(root) 64 65 self._collect_namespaces(root) 66 self._collect_schemalocs(root) 67 return tree 68 69 @classmethod 70 def from_obj(cls, obj): 71 if not obj: 72 return None 73 74 return_obj = super(OpenIOCTestMechanism, cls).from_obj(obj) 75 return_obj.ioc = obj.ioc 76 return return_obj 77 78 def to_obj(self, return_obj=None, ns_info=None): 79 if not return_obj: 80 return_obj = self._binding_class() 81 82 super(OpenIOCTestMechanism, self).to_obj(ns_info=ns_info) 83 return_obj.ioc = self._processed_ioc() 84 return return_obj 85 86 @classmethod 87 def from_dict(cls, d): 88 if not d: 89 return None 90 91 return_obj = super(OpenIOCTestMechanism, cls).from_dict(d) 92 93 if 'ioc' in d: 94 parser = mixbox.xml.get_xml_parser() 95 return_obj.ioc = etree.parse(BytesIO(d['ioc']), parser=parser) 96 97 return return_obj 98 99 def to_dict(self): 100 d = super(OpenIOCTestMechanism, self).to_dict() 101 102 if self.ioc: 103 d['ioc'] = etree.tostring(self._processed_ioc()) 104 105 return d 106 107