create iris alert more modular and docstrings (#52)
taylor_socfortress committed
Jul 18, 2023 at 14:46 UTC
54c75fcd3ff3ef43f8cafd6ec7baa4cf023d4c15
1 file changed
+245
-146
backend/app/services/DFIR_IRIS/alerts.py
+245
-146
@@ -1,12 +1,8 @@
1
-# Standard library imports
1
from typing import Any
2
from typing import Dict
3
from typing import Set
4
6
-# Local application imports
5
from dfir_iris_client.alert import Alert
8
-
9
-# Third-party library imports
6
from loguru import logger
7
8
from app.models.agents import agent_metadata_schema
@@ -21,6 +17,40 @@ from app.services.DFIR_IRIS.universal import UniversalService
17
class IRISAlertsService:
18
"""
19
A service class that encapsulates the logic for pulling alerts from DFIR-IRIS.
20
+
21
+ Attributes
22
+ ----------
23
+ universal_service : UniversalService
24
+ UniversalService object for interacting with "DFIR-IRIS"
25
+ iris_session : Optional[dict]
26
+ Session object for interacting with "DFIR-IRIS". None if session creation failed.
27
+
28
+ Methods
29
+ -------
30
+ _create_iris_session() -> Optional[dict]:
31
+ Create a session with the universal service.
32
+ list_alerts() -> Dict[str, Any]:
33
+ List all alerts from DFIR-IRIS.
34
+ create_alert_general(alert_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
35
+ Create an alert with the provided data.
36
+ get_agent_data(agent_id: str) -> Dict[str, Any]:
37
+ Get agent data based on agent_id.
38
+ get_asset_type_id(os: str) -> int:
39
+ Use AssetTypeResolver to determine the asset type ID.
40
+ validate_ioc_type(ioc_value: str) -> str:
41
+ Validate IoC type using validators.
42
+ create_alert_with_payload(alert_payload: Dict[str, Any]) -> Dict[str, Any]:
43
+ Create an alert using given alert_payload.
44
+ field_exists_ioc(alert_data: Dict[str, Any]) -> Dict[str, Any]:
45
+ Check if an IoC field exists in the alert data.
46
+ valid_ioc_fields() -> Set[str]:
47
+ Get the set of valid IoC fields.
48
+ create_general_payload(alert_data: Dict[str, Any], agent_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
49
+ Craft the general alert payload when it does not contain an IoC.
50
+ create_general_ioc_payload(alert_data: Dict[str, Any], agent_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
51
+ Craft the general alert payload when it does contain an IoC.
52
+ create_base_payload(alert_data: Dict[str, Any], agent_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
53
+ Craft the base alert payload.
54
"""
55
56
def __init__(self):
@@ -74,72 +104,116 @@ class IRISAlertsService:
104
105
def create_alert_general(self, alert_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
106
"""
77
- Create an alert with the provided data.
78
- """
79
- # Get agent data
80
- agent_id = alert_data["agent_id"]
81
- service = AgentService()
82
- agent = service.get_agent(agent_id)
83
- agent_data = agent_metadata_schema.dump(agent)
107
+ Create an alert within DFIR-IRIS with the provided data.
108
85
- # Use AssetTypeResolver to determine the asset type ID
86
- asset_resolver = AssetTypeResolver(agent_data["os"])
87
- agent_asset_type = asset_resolver.get_asset_type_id()
109
+ Parameters
110
+ ----------
111
+ alert_data : Dict[str, Any]
112
+ The alert data used to create the alert.
113
+ alert_id : str
114
+ The ID of the alert.
115
+ index : str
116
+ The index.
117
89
- # Append the asset type ID to the alert data
90
- alert_data["asset_type_id"] = agent_asset_type
118
+ Returns
119
+ -------
120
+ Dict[str, Any]
121
+ The result of the alert creation. Contains information on whether the alert creation was successful,
122
+ an associated message, and the resulting data.
123
+ """
124
+ agent_data = self.get_agent_data(alert_data["agent_id"])
125
+ alert_data["asset_type_id"] = self.get_asset_type_id(agent_data["os"])
126
92
- # Check if IoC field exists
127
ioc_field_present = self.field_exists_ioc(alert_data)
128
if ioc_field_present["success"]:
129
logger.info(f"Found IoC field: {ioc_field_present}")
130
alert_data["ioc_value"] = ioc_field_present["field_value"]
131
+ alert_data["ioc_type"] = self.validate_ioc_type(ioc_field_present["field_value"])
132
+ alert_payload = self.create_general_ioc_payload(alert_data, agent_data, alert_id, index)
133
+ else:
134
+ alert_payload = self.create_general_payload(alert_data, agent_data, alert_id, index)
135
98
- # Define the validator classes to be used
99
- validators = [IPv4AddressValidator, HashValidator, DomainValidator]
100
- ioc_type = None
136
+ return self.create_alert_with_payload(alert_payload)
137
102
- # Iterate over each validator class
103
- for Validator in validators:
104
- validator = Validator(ioc_field_present["field_value"])
105
- result = validator.validate()
138
+ def get_agent_data(self, agent_id: str) -> Dict[str, Any]:
139
+ """
140
+ Get agent data based on agent_id from the `agent_metadata` table.
141
107
- # If the validation is successful, store the ioc_type and break the loop
108
- if result["success"]:
109
- ioc_type = result["ioc_type"]
110
- break
142
+ Parameters
143
+ ----------
144
+ agent_id : str
145
+ The ID of the agent.
146
112
- if ioc_type is not None:
113
- alert_data["ioc_type"] = ioc_type
114
- else:
115
- logger.error("Failed to validate IoC value.")
147
+ Returns
148
+ -------
149
+ Dict[str, Any]
150
+ The agent data corresponding to the given agent_id.
151
+ """
152
+ service = AgentService()
153
+ agent = service.get_agent(agent_id)
154
+ return agent_metadata_schema.dump(agent)
155
117
- alert_payload = self.create_general_ioc_payload(alert_data=alert_data, agent_data=agent_data, alert_id=alert_id, index=index)
156
+ def get_asset_type_id(self, os: str) -> int:
157
+ """
158
+ Use AssetTypeResolver to determine the asset type ID to set within DFIR-IRIS.
159
119
- # Create an alert
120
- alert = Alert(session=self.iris_session)
121
- result = self.universal_service.fetch_and_parse_data(
122
- self.iris_session,
123
- alert.add_alert,
124
- alert_payload,
125
- )
160
+ Parameters
161
+ ----------
162
+ os : str
163
+ The operating system (OS) string used to resolve the asset type ID.
164
127
- if not result["success"]:
128
- return {
129
- "success": False,
130
- "message": "Failed to create alert in DFIR-IRIS",
131
- }
165
+ Returns
166
+ -------
167
+ int
168
+ The ID corresponding to the asset type.
169
+ """
170
+ asset_resolver = AssetTypeResolver(os)
171
+ return asset_resolver.get_asset_type_id()
172
133
- return {
134
- "success": True,
135
- "message": "Successfully created alert in DFIR-IRIS",
136
- "results": result["data"],
137
- }
173
+ def validate_ioc_type(self, ioc_value: str) -> str:
174
+ """
175
+ Validate IoC type using validators.
176
+
177
+ Parameters
178
+ ----------
179
+ ioc_value : str
180
+ The value to validate the IoC type.
181
+
182
+ Returns
183
+ -------
184
+ str
185
+ The type of the IoC. Returns None if validation fails.
186
+ """
187
+ validators = [IPv4AddressValidator, HashValidator, DomainValidator]
188
+ ioc_type = None
189
+
190
+ for Validator in validators:
191
+ validator = Validator(ioc_value)
192
+ result = validator.validate()
193
+
194
+ if result["success"]:
195
+ ioc_type = result["ioc_type"]
196
+ break
197
+
198
+ if ioc_type is None:
199
+ logger.error("Failed to validate IoC value.")
200
+ return ioc_type
201
139
- # Create alert payload
140
- alert_payload = self.create_general_payload(alert_data=alert_data, agent_data=agent_data, alert_id=alert_id, index=index)
202
+ def create_alert_with_payload(self, alert_payload: Dict[str, Any]) -> Dict[str, Any]:
203
+ """
204
+ Create the alert payload to be sent to DFIR-IRIS.
205
+
206
+ Parameters
207
+ ----------
208
+ alert_payload : Dict[str, Any]
209
+ The payload used to create the alert.
210
142
- # Create an alert
211
+ Returns
212
+ -------
213
+ Dict[str, Any]
214
+ The result of the alert creation. Contains information on whether the alert creation was successful,
215
+ an associated message, and the resulting data.
216
+ """
217
alert = Alert(session=self.iris_session)
218
result = self.universal_service.fetch_and_parse_data(
219
self.iris_session,
@@ -161,7 +235,20 @@ class IRISAlertsService:
235
236
def field_exists_ioc(self, alert_data: Dict[str, Any]) -> Dict[str, Any]:
237
"""
164
- Check if an IoC field exists in the alert data.
238
+ Checks if an IoC field exists in the alert data.
239
+
240
+ Parameters
241
+ ----------
242
+ alert_data : Dict[str, Any]
243
+ The alert data to check for the presence of an IoC field.
244
+
245
+ Returns
246
+ -------
247
+ Dict[str, Any]
248
+ If an IoC field exists, returns a dictionary with 'success' set to True,
249
+ 'field_name' as the name of the field, and 'field_value' as the value of the field.
250
+ If an IoC field does not exist, returns a dictionary with 'success' set to False,
251
+ and 'field_name' set to None.
252
"""
253
for field_name in self.valid_ioc_fields:
254
if field_name in alert_data:
@@ -175,56 +262,39 @@ class IRISAlertsService:
262
@property
263
def valid_ioc_fields(self) -> Set[str]:
264
"""
178
- Get the set of valid IoC fields.
265
+ Getter for the set of valid IoC fields.
266
+
267
+ Returns
268
+ -------
269
+ Set[str]
270
+ The set of valid IoC fields.
271
"""
272
return {"misp_value", "opencti_value", "threat_intel_value"}
273
274
def create_general_payload(self, alert_data: Dict[str, Any], agent_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
275
"""
184
- Craft the general alert payload when it does not contain an IoC.
276
+ Crafts the general alert payload when it does not contain an IoC.
277
+
278
+ Parameters
279
+ ----------
280
+ alert_data : Dict[str, Any]
281
+ The alert data.
282
+ agent_data : Dict[str, Any]
283
+ The agent data.
284
+ alert_id : str
285
+ The ID of the alert.
286
+ index : str
287
+ The index.
288
+
289
+ Returns
290
+ -------
291
+ Dict[str, Any]
292
+ The crafted alert payload. If an error occurs during crafting,
293
+ a dictionary with 'success' set to False and an error message is returned.
294
"""
295
try:
187
- payload = {
188
- "alert_title": alert_data["rule_description"],
189
- "alert_description": alert_data["rule_description"],
190
- "alert_source": "Wazuh",
191
- "assets": [
192
- {
193
- "asset_name": agent_data["hostname"],
194
- "asset_ip": agent_data["ip_address"],
195
- "asset_description": agent_data["os"],
196
- "asset_type_id": alert_data["asset_type_id"],
197
- },
198
- ],
199
- # "alert_source_link": f"{self.grafana_url}/explore?left=%5B%22now-6h%22,%22now%22,%22WAZUH%22,%7B%22refId%22"
200
- # ":%22A%22,%22query%22:%22process_id:%5C%22"
201
- # f"{alert_data['process_id']}%5C%22%20AND%20"
202
- # f"agent_name:%5C%22{alert_data['agent_name']}%5C%22%22,%22alias%22"
203
- # ":%22%22,%22metrics%22:%5B%7B%22id%22:%221%22,%22type%22:%22logs%22,%22settings%22:%7B%22limit%22:%22500%22"
204
- # "%7D%7D%5D,%22bucketAggs%22:%5B%5D,%22timeField%22:%22timestamp%22%7D%5D",
205
- "alert_status_id": 3,
206
- "alert_severity_id": 5,
207
- # "alert_customer_id": customer_code_details["customer_code_iris_id"],
208
- "alert_customer_id": 1,
209
- "alert_source_content": alert_data,
210
- "alert_context": {
211
- # "customer_id": f"{alert_data['alert_payload']['alert_details']['_source']['agent_labels_customer']},"
212
- # f"{customer_code_details['customer_code_iris_index']}",
213
- "alert_id": alert_id,
214
- "alert_name": alert_data["rule_description"],
215
- "alert_level": alert_data["rule_level"],
216
- "rule_id": alert_data["rule_id"],
217
- "asset_name": agent_data["hostname"],
218
- "asset_ip": agent_data["ip_address"],
219
- "asset_type": alert_data["asset_type_id"],
220
- "process_id": alert_data["process_id"],
221
- # If the `rule_mitre_id` field exists in the alert_details, add it to the payload
222
- "rule_mitre_id": alert_data.get("rule_mitre_id", "n/a"),
223
- "rule_mitre_tactic": alert_data.get("rule_mitre_tactic", "n/a"),
224
- "rule_mitre_technique": alert_data.get("rule_mitre_technique", "n/a"),
225
- },
226
- "alert_note": alert_data.get("ask_socfortress", "Ask SOCFortress not enabled"),
227
- }
296
+ payload = self.create_base_payload(alert_data, agent_data, alert_id, index)
297
+ payload["alert_note"] = alert_data.get("ask_socfortress", "Ask SOCFortress not enabled")
298
return payload
299
except Exception as e:
300
logger.error(f"Error creating general alert payload: {e}")
@@ -238,59 +308,88 @@ class IRISAlertsService:
308
index: str,
309
) -> Dict[str, Any]:
310
"""
241
- Craft the general alert payload when it does contain an IoC.
311
+ Crafts the general alert payload when it does contain an IoC.
312
+
313
+ Parameters
314
+ ----------
315
+ alert_data : Dict[str, Any]
316
+ The alert data.
317
+ agent_data : Dict[str, Any]
318
+ The agent data.
319
+ alert_id : str
320
+ The ID of the alert.
321
+ index : str
322
+ The index.
323
+
324
+ Returns
325
+ -------
326
+ Dict[str, Any]
327
+ The crafted alert payload with IoC. If an error occurs during crafting,
328
+ a dictionary with 'success' set to False and an error message is returned.
329
"""
330
try:
244
- payload = {
245
- "alert_title": alert_data["rule_description"],
246
- "alert_description": alert_data["rule_description"],
247
- "alert_source": "Wazuh",
248
- "assets": [
249
- {
250
- "asset_name": agent_data["hostname"],
251
- "asset_ip": agent_data["ip_address"],
252
- "asset_description": agent_data["os"],
253
- "asset_type_id": alert_data["asset_type_id"],
254
- },
255
- ],
256
- # "alert_source_link": f"{self.grafana_url}/explore?left=%5B%22now-6h%22,%22now%22,%22WAZUH%22,%7B%22refId%22"
257
- # ":%22A%22,%22query%22:%22process_id:%5C%22"
258
- # f"{alert_data['process_id']}%5C%22%20AND%20"
259
- # f"agent_name:%5C%22{alert_data['agent_name']}%5C%22%22,%22alias%22"
260
- # ":%22%22,%22metrics%22:%5B%7B%22id%22:%221%22,%22type%22:%22logs%22,%22settings%22:%7B%22limit%22:%22500%22"
261
- # "%7D%7D%5D,%22bucketAggs%22:%5B%5D,%22timeField%22:%22timestamp%22%7D%5D",
262
- "alert_status_id": 3,
263
- "alert_severity_id": 5,
264
- # "alert_customer_id": customer_code_details["customer_code_iris_id"],
265
- "alert_customer_id": 1,
266
- "alert_source_content": alert_data,
267
- "alert_context": {
268
- # "customer_id": f"{alert_data['alert_payload']['alert_details']['_source']['agent_labels_customer']},"
269
- # f"{customer_code_details['customer_code_iris_index']}",
270
- "alert_id": alert_id,
271
- "alert_name": alert_data["rule_description"],
272
- "alert_level": alert_data["rule_level"],
273
- "rule_id": alert_data["rule_id"],
274
- "asset_name": agent_data["hostname"],
275
- "asset_ip": agent_data["ip_address"],
276
- "asset_type": alert_data["asset_type_id"],
277
- "process_id": alert_data["process_id"],
278
- # If the `rule_mitre_id` field exists in the alert_details, add it to the payload
279
- "rule_mitre_id": alert_data.get("rule_mitre_id", "n/a"),
280
- "rule_mitre_tactic": alert_data.get("rule_mitre_tactic", "n/a"),
281
- "rule_mitre_technique": alert_data.get("rule_mitre_technique", "n/a"),
331
+ payload = self.create_base_payload(alert_data, agent_data, alert_id, index)
332
+ payload["alert_note"] = alert_data.get("ask_socfortress", "Ask SOCFortress not enabled")
333
+ payload["alert_iocs"] = [
334
+ {
335
+ "ioc_value": alert_data["ioc_value"],
336
+ "ioc_description": "IoC found in the alert",
337
+ "ioc_tlp_id": 1,
338
+ "ioc_type_id": alert_data["ioc_type"],
339
},
283
- "alert_note": alert_data.get("ask_socfortress", "Ask SOCFortress not enabled"),
284
- "alert_iocs": [
285
- {
286
- "ioc_value": alert_data["ioc_value"],
287
- "ioc_description": "IoC found in the alert",
288
- "ioc_tlp_id": 1,
289
- "ioc_type_id": alert_data["ioc_type"],
290
- },
291
- ],
292
- }
340
+ ]
341
return payload
342
except Exception as e:
343
logger.error(f"Error creating general alert payload with IoC: {e}")
344
return {"success": False, "message": f"Error creating general alert payload with IoC: {e}"}
345
+
346
+ def create_base_payload(self, alert_data: Dict[str, Any], agent_data: Dict[str, Any], alert_id: str, index: str) -> Dict[str, Any]:
347
+ """
348
+ Crafts the base alert payload.
349
+
350
+ Parameters
351
+ ----------
352
+ alert_data : Dict[str, Any]
353
+ The alert data.
354
+ agent_data : Dict[str, Any]
355
+ The agent data.
356
+ alert_id : str
357
+ The ID of the alert.
358
+ index : str
359
+ The index.
360
+
361
+ Returns
362
+ -------
363
+ Dict[str, Any]
364
+ The crafted base alert payload.
365
+ """
366
+ return {
367
+ "alert_title": alert_data["rule_description"],
368
+ "alert_description": alert_data["rule_description"],
369
+ "alert_source": "Wazuh",
370
+ "assets": [
371
+ {
372
+ "asset_name": agent_data["hostname"],
373
+ "asset_ip": agent_data["ip_address"],
374
+ "asset_description": agent_data["os"],
375
+ "asset_type_id": alert_data["asset_type_id"],
376
+ },
377
+ ],
378
+ "alert_status_id": 3,
379
+ "alert_severity_id": 5,
380
+ "alert_customer_id": 1,
381
+ "alert_source_content": alert_data,
382
+ "alert_context": {
383
+ "alert_id": alert_id,
384
+ "alert_name": alert_data["rule_description"],
385
+ "alert_level": alert_data["rule_level"],
386
+ "rule_id": alert_data["rule_id"],
387
+ "asset_name": agent_data["hostname"],
388
+ "asset_ip": agent_data["ip_address"],
389
+ "asset_type": alert_data["asset_type_id"],
390
+ "process_id": alert_data["process_id"],
391
+ "rule_mitre_id": alert_data.get("rule_mitre_id", "n/a"),
392
+ "rule_mitre_tactic": alert_data.get("rule_mitre_tactic", "n/a"),
393
+ "rule_mitre_technique": alert_data.get("rule_mitre_technique", "n/a"),
394
+ },
395
+ }