-
Notifications
You must be signed in to change notification settings - Fork 71
/
mongodb_filter.py
205 lines (179 loc) · 9.29 KB
/
mongodb_filter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
from medallion.filters.basic_filter import BasicFilter
from medallion.utils.common import convert_to_stix_datetime
class MongoDBFilter(BasicFilter):
def __init__(self, filter_args, basic_filter, allowed, start_index=0, end_index=None):
super(MongoDBFilter, self).__init__(filter_args)
self.basic_filter = basic_filter
self.full_query = self._query_parameters(allowed)
self.start_index = start_index
self.end_index = end_index
def _query_parameters(self, allowed):
parameters = self.basic_filter
if self.filter_args:
match_type = self.filter_args.get("match[type]")
if match_type and "type" in allowed:
types_ = match_type.split(",")
if len(types_) == 1:
parameters["_type"] = types_[0]
else:
parameters["_type"] = {"$in": types_}
match_id = self.filter_args.get("match[id]")
if match_id and "id" in allowed:
ids_ = match_id.split(",")
if len(ids_) == 1:
parameters["id"] = ids_[0]
else:
parameters["id"] = {"$in": ids_}
return parameters
def process_filter(self, data, allowed, manifest_info):
match_filter = {"$match": self.full_query}
pipeline = [match_filter]
# create added_after filter
added_after_date = self.filter_args.get("added_after")
if added_after_date:
added_after_timestamp = convert_to_stix_datetime(added_after_date)
date_filter = {"$match": {"date_added": {"$gt": added_after_timestamp}}}
pipeline.append(date_filter)
# need to handle marking-definitions differently as they are not versioned like SDO's
if self.filter_contains_marking_definition(pipeline):
# If we are finding marking-definitions from the objects collection
# we need to change the match criteria from "_type" to "type"
if data.name == "objects" and "_type" in pipeline[0]["$match"].keys():
pipeline[0]["$match"]["type"] = pipeline[0]["$match"].pop("_type")
# Calculate total number of matching documents
if data.name == "objects":
count = self.get_result_count(pipeline, manifest_info["mongodb_collection"])
else:
count = self.get_result_count(pipeline, data)
self.add_pagination_operations(pipeline)
cursor = data.aggregate(pipeline)
results = list(cursor)
return count, results
# create version filter
if "version" in allowed:
match_version = self.filter_args.get("match[version]")
if not match_version:
match_version = "last"
if "all" not in match_version:
actual_dates = [x for x in match_version.split(",") if (x != "first" and x != "last")]
# If specific dates have been selected, then we add these to the $match criteria
# created from the self.full_query at the beginning of this method. The reason we need
# to do this is because the $indexOfArray function below will return -1 if the date
# doesn't exist in the versions array. -1 will be interpreted by $arrayElemAt as the
# final element in the array and we will return the wrong result. i.e. not only will the
# version dates be incorrect, but we shouldn't have returned a result at all.
# if actual_dates:
if len(actual_dates) > 0:
pipeline.insert(1, {"$match": {"versions": {"$all": [",".join(actual_dates)]}}})
# The versions array in the mongodb document is ordered newest to oldest, so the 'last'
# (most recent date) is in first position in the list and the oldest 'first' is in
# the last position (equal to index -1 for $arrayElemAt)
version_selector = []
if "last" in match_version:
version_selector.append({"$arrayElemAt": ["$versions", 0]})
if "first" in match_version:
version_selector.append({"$arrayElemAt": ["$versions", -1]})
for d in actual_dates:
version_selector.append({"$arrayElemAt": ["$versions", {"$indexOfArray": ["$versions", d]}]})
version_filter = {
"$addFields": {
"versions": version_selector,
},
}
pipeline.append(version_filter)
if data.name == "manifests":
count = self.get_result_count(pipeline, data)
self.add_pagination_operations(pipeline)
cursor = data.aggregate(pipeline)
results = list(cursor)
else:
results = []
# Get the count of matching documents - need to unwind the versions selected to get accurate count.
count_pipeline = list(pipeline)
count_pipeline.append({"$unwind": "$versions"})
count = self.get_result_count(count_pipeline, manifest_info["mongodb_collection"])
# only bother doing the rest of the query if the start index is less than the total number of results.
if self.start_index < count:
# Join the filtered manifest(s) to the objects collection
join_objects = {
"$lookup": {
"from": "objects",
"localField": "id",
"foreignField": "id",
"as": "obj",
},
}
pipeline.append(join_objects)
# Copy the filtered version list to the embedded object document
add_versions = {
"$addFields": {"obj.versions": "$versions"},
}
pipeline.append(add_versions)
# denormalize the embedded objects and replace the document root
pipeline.append({"$unwind": "$obj"})
pipeline.append({"$replaceRoot": {"newRoot": "$obj"}})
# Redact the result set removing objects where the modified date is not in
# the versions array and the object isn't in the correct collection.
# The collection filter is required because the join between manifests and objects
# does not include collection_id
col_id = self.full_query["_collection_id"]
redact_objects = {
"$redact": {
"$cond": {
"if": {
"$and": [
{"$eq": ["$_collection_id", col_id]},
{
"$or": [
{"$eq": ["$type", "marking-definition"]},
{"$setIsSubset": [["$modified"], "$versions"]},
],
},
],
},
"then": "$$KEEP",
"else": "$$PRUNE",
},
},
}
pipeline.append(redact_objects)
# Project the final results
project_results = {
"$project": {
"versions": 0,
},
}
pipeline.append(project_results)
self.add_pagination_operations(pipeline)
cursor = manifest_info["mongodb_collection"].aggregate(pipeline)
results = list(cursor)
return count, results
def add_pagination_operations(self, pipeline):
if self.start_index is not None and self.end_index is not None:
pipeline.append({"$skip": self.start_index})
pipeline.append({"$limit": (self.end_index - self.start_index) + 1})
def get_result_count(self, pipeline, data):
count_pipeline = list(pipeline)
count_pipeline.append({"$count": "total_count"})
count_result = list(data.aggregate(count_pipeline))
if len(count_result) == 0:
# No results
return 0
count = count_result[0]["total_count"]
return count
def filter_contains_marking_definition(self, pipeline):
# If we are matching on id (either match[id]= or /{id}), then check if
# we are trying to find a marking definition. If so, we don't want do
# filter by version as marking-definition objects are not versioned.
if "id" in pipeline[0]["$match"].keys() and pipeline[0]["$match"]["id"].startswith("marking-definition"):
return True
if "_type" in pipeline[0]["$match"].keys():
if ((
isinstance(pipeline[0]["$match"]["_type"], dict) and
"$in" in pipeline[0]["$match"]["_type"].keys()
) and
("marking-definition" in pipeline[0]["$match"]["_type"]["$in"])):
return True
elif pipeline[0]["$match"]["_type"].startswith("marking-definition"):
return True
return False