1# Copyright (c) 2017, The MITRE Corporation. All rights reserved.
2# See LICENSE.txt for complete terms.
3
4# external
5from mixbox import fields
6from mixbox import entities
7
8# internal
9import stix
10from stix.common import InformationSource
11
12# bindings
13import stix.bindings.data_marking as stix_data_marking_binding
14
15
16class MarkingStructureFactory(entities.EntityFactory):
17    @classmethod
18    def entity_class(cls, key):
19        import stix.extensions.marking.tlp  # noqa
20        import stix.extensions.marking.simple_marking  # noqa
21        import stix.extensions.marking.terms_of_use_marking  # noqa
22        return stix.lookup_extension(key, default=MarkingStructure)
23
24
25class MarkingStructure(stix.Entity):
26    _binding = stix_data_marking_binding
27    _binding_class = stix_data_marking_binding.MarkingStructureType
28    _namespace = 'http://data-marking.mitre.org/Marking-1'
29    _XSI_TYPE = None  # overridden by subclasses
30
31    id_ = fields.IdField("id")
32    idref = fields.IdrefField("idref")
33    marking_model_name = fields.TypedField("marking_model_name")
34    marking_model_ref = fields.TypedField("marking_model_ref")
35
36    def __init__(self):
37        super(MarkingStructure, self).__init__()
38
39        self.id_ = None
40        self.idref = None
41        self.marking_model_name = None
42        self.marking_model_ref = None
43
44    def to_dict(self):
45        d = super(MarkingStructure, self).to_dict()
46
47        if self._XSI_TYPE:
48            d['xsi:type'] = self._XSI_TYPE
49
50        return d
51
52    @staticmethod
53    def lookup_class(xsi_type):
54        return stix.lookup_extension(xsi_type, default=MarkingStructure)
55
56
57class MarkingSpecification(stix.Entity):
58    _binding = stix_data_marking_binding
59    _binding_class = stix_data_marking_binding.MarkingSpecificationType
60    _namespace = 'http://data-marking.mitre.org/Marking-1'
61
62    id_ = fields.IdField("id")
63    idref = fields.IdrefField("idref")
64    version = fields.TypedField("version")
65    controlled_structure = fields.TypedField("Controlled_Structure")
66    marking_structures = fields.TypedField("Marking_Structure", MarkingStructure, factory=MarkingStructureFactory, multiple=True, key_name="marking_structures")
67    information_source = fields.TypedField("Information_Source", InformationSource)
68
69    def __init__(self, controlled_structure=None, marking_structures=None):
70        super(MarkingSpecification, self).__init__()
71
72        self.id_ = None
73        self.idref = None
74        self.version = None
75        self.controlled_structure = controlled_structure
76        self.marking_structures = marking_structures
77        self.information_source = None
78
79
80class Marking(stix.EntityList):
81    _binding = stix_data_marking_binding
82    _binding_class = stix_data_marking_binding.MarkingType
83    _namespace = 'http://data-marking.mitre.org/Marking-1'
84    marking = fields.TypedField("Marking", MarkingSpecification, multiple=True)
85
86    def __init__(self, markings=None):
87        super(Marking, self).__init__(markings)
88
89    def add_marking(self, value):
90        self.marking.append(value)
91
92
93# Backwards compatibility
94add_extension = stix.add_extension
95