1# Copyright (c) 2017, The MITRE Corporation. All rights reserved.
2# See LICENSE.txt for complete terms.
3
4# external
5from mixbox import entities
6from mixbox import fields
7
8# internal
9import stix
10from stix.bindings import course_of_action as coa_binding
11
12
13class StructuredCOAFactory(entities.EntityFactory):
14    @classmethod
15    def entity_class(cls, key):
16        import stix.extensions.structured_coa.generic_structured_coa  # noqa
17        return stix.lookup_extension(key)
18
19
20class _BaseStructuredCOA(stix.Entity):
21    _namespace = "http://stix.mitre.org/CourseOfAction-1"
22    _binding = coa_binding
23    _binding_class = coa_binding.StructuredCOAType
24
25    id_ = fields.IdField("id")
26    idref = fields.IdrefField("idref")
27
28    def __init__(self, id_=None, idref=None):
29        super(_BaseStructuredCOA, self).__init__()
30        self.id_ = id_
31        self.idref = idref
32
33    def to_dict(self):
34        d = super(_BaseStructuredCOA, self).to_dict()
35        d['xsi:type'] = self._XSI_TYPE
36        return d
37
38    def to_obj(self, ns_info=None):
39        obj = super(_BaseStructuredCOA, self).to_obj(ns_info=ns_info)
40        obj.xsi_type = self._XSI_TYPE
41        return obj
42
43# Backwards compatibility
44add_extension = stix.add_extension
45