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