1+ import json
2+ import uuid
3+
4+
5+ class OpenCTIStix2Splitter :
6+ def __init__ (self ):
7+ self .cache_index = {}
8+ self .cache_added = []
9+ self .entities = []
10+ self .relationships = []
11+
12+ def enlist_entity_element (self , item_id , raw_data ):
13+ nb_deps = 0
14+ item = raw_data [item_id ]
15+ is_marking = item ["id" ].startswith ('marking-definition--' )
16+ if "created_by_ref" in item and is_marking is False and self .cache_index .get (item ["created_by_ref" ]) is None :
17+ nb_deps += 1
18+ self .enlist_entity_element (item ["created_by_ref" ], raw_data )
19+
20+ if "object_refs" in item :
21+ for object_ref in item ["object_refs" ]:
22+ nb_deps += 1
23+ if self .cache_index .get (object_ref ) is None :
24+ self .enlist_entity_element (object_ref , raw_data )
25+
26+ if "object_marking_refs" in item :
27+ for object_marking_ref in item ["object_marking_refs" ]:
28+ nb_deps += 1
29+ if self .cache_index .get (object_marking_ref ) is None :
30+ self .enlist_entity_element (object_marking_ref , raw_data )
31+
32+ item ["nb_deps" ] = nb_deps
33+ self .entities .append (item )
34+ self .cache_index [item_id ] = item # Put in cache
35+
36+ def enlist_relation_element (self , item_id , raw_data ):
37+ nb_deps = 0
38+ item = raw_data [item_id ]
39+ source = item ["source_ref" ]
40+ target = item ["target_ref" ]
41+ if source .startswith ('relationship--' ):
42+ nb_deps += 1
43+ if self .cache_index .get (source ) is None :
44+ self .enlist_entity_element (target , raw_data )
45+ if target .startswith ('relationship--' ):
46+ nb_deps += 1
47+ if self .cache_index .get (target ) is None :
48+ self .enlist_entity_element (target , raw_data )
49+ item ["nb_deps" ] = nb_deps
50+ self .relationships .append (item )
51+ self .cache_index [item_id ] = item # Put in cache
52+
53+ def split_bundle (self , bundle ) -> list :
54+ """splits a valid stix2 bundle into a list of bundles
55+
56+ :param bundle: valid stix2 bundle
57+ :type bundle:
58+ :raises Exception: if data is not valid JSON
59+ :return: returns a list of bundles
60+ :rtype: list
61+ """
62+ try :
63+ bundle_data = json .loads (bundle )
64+ except :
65+ raise Exception ("File data is not a valid JSON" )
66+
67+ raw_data = {}
68+ for item in bundle_data ["objects" ]:
69+ raw_data [item ["id" ]] = item
70+
71+ for item in bundle_data ["objects" ]:
72+ is_entity = item ["type" ] != "relationship"
73+ if is_entity :
74+ self .enlist_entity_element (item ["id" ], raw_data )
75+
76+ for item in bundle_data ["objects" ]:
77+ is_relation = item ["type" ] == "relationship"
78+ if is_relation :
79+ self .enlist_relation_element (item ["id" ], raw_data )
80+
81+ bundles = []
82+ for entity in self .entities :
83+ bundles .append (self .stix2_create_bundle ([entity ]))
84+ for relationship in self .relationships :
85+ bundles .append (self .stix2_create_bundle ([relationship ]))
86+ return bundles
87+
88+ # @deprecated
89+ def split_stix2_bundle (self , bundle ) -> list :
90+ """splits a valid stix2 bundle into a list of bundles
91+
92+ :param bundle: valid stix2 bundle
93+ :type bundle:
94+ :raises Exception: if data is not valid JSON
95+ :return: returns a list of bundles
96+ :rtype: list
97+ """
98+ try :
99+ bundle_data = json .loads (bundle )
100+ except :
101+ raise Exception ("File data is not a valid JSON" )
102+
103+ # validation = validate_parsed_json(bundle_data)
104+ # if not validation.is_valid:
105+ # raise ValueError('The bundle is not a valid STIX2 JSON:' + bundle)
106+
107+ # Index all objects by id
108+ for item in bundle_data ["objects" ]:
109+ self .cache_index [item ["id" ]] = item
110+
111+ bundles = []
112+ # Reports must be handled because of object_refs
113+ for item in bundle_data ["objects" ]:
114+ if item ["type" ] == "report" :
115+ items_to_send = self .stix2_deduplicate_objects (
116+ self .stix2_get_report_objects (item )
117+ )
118+ for item_to_send in items_to_send :
119+ self .cache_added .append (item_to_send ["id" ])
120+ bundles .append (self .stix2_create_bundle (items_to_send ))
121+
122+ # Relationships not added in previous reports
123+ for item in bundle_data ["objects" ]:
124+ if item ["type" ] == "relationship" and item ["id" ] not in self .cache_added :
125+ items_to_send = self .stix2_deduplicate_objects (
126+ self .stix2_get_relationship_objects (item )
127+ )
128+ for item_to_send in items_to_send :
129+ self .cache_added .append (item_to_send ["id" ])
130+ bundles .append (self .stix2_create_bundle (items_to_send ))
131+
132+ # Entities not added in previous reports and relationships
133+ for item in bundle_data ["objects" ]:
134+ if item ["type" ] != "relationship" and item ["id" ] not in self .cache_added :
135+ items_to_send = self .stix2_deduplicate_objects (
136+ self .stix2_get_entity_objects (item )
137+ )
138+ for item_to_send in items_to_send :
139+ self .cache_added .append (item_to_send ["id" ])
140+ bundles .append (self .stix2_create_bundle (items_to_send ))
141+
142+ return bundles
143+
144+ @staticmethod
145+ def stix2_deduplicate_objects (items ) -> list :
146+ """deduplicate stix2 items
147+
148+ :param items: valid stix2 items
149+ :type items:
150+ :return: de-duplicated list of items
151+ :rtype: list
152+ """
153+
154+ ids = []
155+ final_items = []
156+ for item in items :
157+ if item ["id" ] not in ids :
158+ final_items .append (item )
159+ ids .append (item ["id" ])
160+ return final_items
161+
162+ def stix2_get_report_objects (self , report ) -> list :
163+ """get a list of items for a stix2 report object
164+
165+ :param report: valid stix2 report object
166+ :type report:
167+ :return: list of items for a stix2 report object
168+ :rtype: list
169+ """
170+
171+ items = [report ]
172+ # Add all object refs
173+ for object_ref in report ["object_refs" ]:
174+ items .append (self .cache_index [object_ref ])
175+ for item in items :
176+ if item ["type" ] == "relationship" :
177+ items = items + self .stix2_get_relationship_objects (item )
178+ else :
179+ items = items + self .stix2_get_entity_objects (item )
180+ return items
181+
182+ @staticmethod
183+ def stix2_create_bundle (items ):
184+ """create a stix2 bundle with items
185+
186+ :param items: valid stix2 items
187+ :type items:
188+ :return: JSON of the stix2 bundle
189+ :rtype:
190+ """
191+
192+ bundle = {
193+ "type" : "bundle" ,
194+ "id" : "bundle--" + str (uuid .uuid4 ()),
195+ "spec_version" : "2.0" ,
196+ "objects" : items ,
197+ }
198+ return json .dumps (bundle )
199+
200+ def stix2_get_relationship_objects (self , relationship ) -> list :
201+ """get a list of relations for a stix2 relationship object
202+
203+ :param relationship: valid stix2 relationship
204+ :type relationship:
205+ :return: list of relations objects
206+ :rtype: list
207+ """
208+
209+ items = [relationship ]
210+ # Get source ref
211+ if relationship ["source_ref" ] in self .cache_index :
212+ items .append (self .cache_index [relationship ["source_ref" ]])
213+
214+ # Get target ref
215+ if relationship ["target_ref" ] in self .cache_index :
216+ items .append (self .cache_index [relationship ["target_ref" ]])
217+
218+ # Get embedded objects
219+ embedded_objects = self .stix2_get_embedded_objects (relationship )
220+ # Add created by ref
221+ if embedded_objects ["created_by_ref" ] is not None :
222+ items .append (embedded_objects ["created_by_ref" ])
223+ # Add marking definitions
224+ if len (embedded_objects ["object_marking_refs" ]) > 0 :
225+ items = items + embedded_objects ["object_marking_refs" ]
226+
227+ return items
228+
229+ def stix2_get_embedded_objects (self , item ) -> dict :
230+ """gets created and marking refs for a stix2 item
231+
232+ :param item: valid stix2 item
233+ :type item:
234+ :return: returns a dict of created_by_ref of object_marking_refs
235+ :rtype: dict
236+ """
237+ # Marking definitions
238+ object_marking_refs = []
239+ if "object_marking_refs" in item :
240+ for object_marking_ref in item ["object_marking_refs" ]:
241+ if object_marking_ref in self .cache_index :
242+ object_marking_refs .append (self .cache_index [object_marking_ref ])
243+ # Created by ref
244+ created_by_ref = None
245+ if "created_by_ref" in item and item ["created_by_ref" ] in self .cache_index :
246+ created_by_ref = self .cache_index [item ["created_by_ref" ]]
247+
248+ return {
249+ "object_marking_refs" : object_marking_refs ,
250+ "created_by_ref" : created_by_ref ,
251+ }
252+
253+ def stix2_get_entity_objects (self , entity ) -> list :
254+ """process a stix2 entity
255+
256+ :param entity: valid stix2 entity
257+ :type entity:
258+ :return: entity objects as list
259+ :rtype: list
260+ """
261+
262+ items = [entity ]
263+ # Get embedded objects
264+ embedded_objects = self .stix2_get_embedded_objects (entity )
265+ # Add created by ref
266+ if embedded_objects ["created_by_ref" ] is not None :
267+ items .append (embedded_objects ["created_by_ref" ])
268+ # Add marking definitions
269+ if len (embedded_objects ["object_marking_refs" ]) > 0 :
270+ items = items + embedded_objects ["object_marking_refs" ]
271+
272+ return items
0 commit comments