4
from fastapi import HTTPException
5
from loguru import logger
6
from sqlalchemy.ext.asyncio import AsyncSession
7
+from sqlalchemy.future import select
8
9
+# from app.integrations.alert_escalation.utils.universal import get_agent_data
10
+from app.agents.routes.agents import get_agent
11
+from app.agents.schema.agents import AgentsResponse
12
from app.connectors.dfir_iris.utils.universal import fetch_and_validate_data
13
from app.connectors.dfir_iris.utils.universal import initialize_client_and_alert
14
from app.connectors.utils import get_connector_info_from_db
15
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
16
+from app.integrations.alert_creation.models.alert_settings import AlertCreationSettings
17
from app.integrations.alert_escalation.schema.general_alert import CreateAlertRequest
18
from app.integrations.alert_escalation.schema.general_alert import CreateAlertResponse
19
from app.integrations.alert_escalation.schema.general_alert import GenericAlertModel
23
from app.integrations.alert_escalation.schema.general_alert import IrisAsset
24
from app.integrations.alert_escalation.schema.general_alert import IrisIoc
25
from app.integrations.alert_escalation.schema.general_alert import ValidIocFields
21
-from app.integrations.alert_escalation.utils.universal import get_agent_data
22
-from app.integrations.alert_escalation.utils.universal import get_asset_type_id
23
-from app.integrations.alert_escalation.utils.universal import validate_ioc_type
26
+from app.integrations.utils.alerts import get_asset_type_id
27
+from app.integrations.utils.alerts import validate_ioc_type
28
+from app.utils import get_customer_alert_settings
29
+
30
+
31
+async def is_customer_code_valid(customer_code: str, session: AsyncSession) -> bool:
32
+ logger.info(f"Checking if customer_code: {customer_code} is valid.")
33
+
34
+ result = await session.execute(
35
+ select(AlertCreationSettings).where(AlertCreationSettings.customer_code == customer_code),
36
+ )
37
+ settings = result.scalars().first()
38
+
39
+ if settings:
40
+ return True
41
+
42
+ return False
43
44
45
def valid_ioc_fields() -> Set[str]:
53
return {field.value for field in ValidIocFields}
54
55
56
+async def construct_alert_source_link(alert_details: GenericAlertModel, session: AsyncSession) -> str:
57
+ """
58
+ Construct the alert source link for the alert details.
59
+ Parameters
60
+ ----------
61
+ alert_details: CreateAlertRequest
62
+ The alert details.
63
+ Returns
64
+ -------
65
+ str
66
+ The alert source link.
67
+ """
68
+ # Check if the alert has a process id and that it is not "No process ID found"
69
+ if hasattr(alert_details, "process_id") and alert_details._source.process_id != "No process ID found":
70
+ query_string = f"%22query%22:%22process_id:%5C%22{alert_details._source.process_id}%5C%22%20AND%20"
71
+ else:
72
+ query_string = f"%22query%22:%22_id:%5C%22{alert_details._id}%5C%22%20AND%20"
73
+
74
+ grafana_url = (
75
+ await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)
76
+ ).grafana_url
77
+
78
+ return (
79
+ f"{grafana_url}/explore?left=%5B%22now-6h%22,%22now%22,%22WAZUH%22,%7B%22refId%22:%22A%22,"
80
+ f"{query_string}"
81
+ f"agent_name:%5C%22{alert_details._source.agent_name}%5C%22%22,"
82
+ "%22alias%22:%22%22,%22metrics%22:%5B%7B%22id%22:%221%22,%22type%22:%22logs%22,%22settings%22:%7B%22limit%22:%22500%22%7D%7D%5D,"
83
+ "%22bucketAggs%22:%5B%5D,%22timeField%22:%22timestamp%22%7D%5D"
84
+ )
85
+
86
+
87
async def get_single_alert_details(alert_details: CreateAlertRequest) -> GenericAlertModel:
88
logger.info(f"Fetching alert details for alert {alert_details.alert_id} in index {alert_details.index_name}")
89
es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
96
raise HTTPException(status_code=400, detail=f"Failed to collect alert details: {e}")
97
98
49
-def build_ioc_payload(alert_details: GenericAlertModel) -> Optional[IrisIoc]:
99
+async def build_ioc_payload(alert_details: GenericAlertModel) -> Optional[IrisIoc]:
100
for field in valid_ioc_fields():
101
if hasattr(alert_details._source, field):
102
ioc_value = getattr(alert_details._source, field)
53
- ioc_type = validate_ioc_type(ioc_value=ioc_value)
103
+ ioc_type = await validate_ioc_type(ioc_value=ioc_value)
104
return IrisIoc(ioc_value=ioc_value, ioc_description="IoC found in alert", ioc_tlp_id=1, ioc_type_id=ioc_type)
105
return None
106
107
58
-def build_asset_payload(agent_data, alert_details) -> IrisAsset:
59
- return IrisAsset(
60
- asset_name=agent_data.hostname,
61
- asset_ip=agent_data.ip_address,
62
- asset_description=agent_data.os,
63
- asset_type_id=alert_details.asset_type_id,
64
- )
108
+async def build_asset_payload(agent_data: AgentsResponse, alert_details) -> IrisAsset:
109
+ if agent_data.success:
110
+ return IrisAsset(
111
+ asset_name=agent_data.agents[0].hostname,
112
+ asset_ip=agent_data.agents[0].ip_address,
113
+ asset_description=agent_data.agents[0].os,
114
+ asset_type_id=await get_asset_type_id(agent_data.agents[0].os),
115
+ )
116
+ return IrisAsset()
117
118
67
-def build_alert_context_payload(alert_details: GenericAlertModel, agent_data) -> IrisAlertContext:
119
+async def build_alert_context_payload(
120
+ alert_details: GenericAlertModel,
121
+ agent_data: AgentsResponse,
122
+ session: AsyncSession,
123
+) -> IrisAlertContext:
124
return IrisAlertContext(
125
+ customer_iris_id=(
126
+ await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)
127
+ ).iris_customer_id,
128
+ customer_name=(
129
+ await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)
130
+ ).customer_name,
131
+ customer_cases_index=(
132
+ await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)
133
+ ).iris_index,
134
alert_id=alert_details._id,
135
alert_name=alert_details._source.rule_description,
136
alert_level=alert_details._source.rule_level,
137
rule_id=alert_details._source.rule_id,
73
- asset_name=agent_data.hostname,
74
- asset_ip=agent_data.ip_address,
75
- asset_type=alert_details.asset_type_id,
138
+ asset_name=agent_data.agents[0].hostname,
139
+ asset_ip=agent_data.agents[0].ip_address,
140
+ asset_type=await get_asset_type_id(agent_data.agents[0].os),
141
process_id=getattr(alert_details._source, "process_id", "No process id found"),
142
rule_mitre_id=getattr(alert_details._source, "rule_mitre_id", "No rule mitre id found"),
143
rule_mitre_tactic=getattr(alert_details._source, "rule_mitre_tactic", "No rule mitre tactic found"),
145
)
146
147
83
-def build_alert_payload(alert_details: GenericAlertModel, agent_data, ioc_payload: Optional[IrisIoc]) -> IrisAlertPayload:
84
- asset_payload = build_asset_payload(agent_data, alert_details)
85
- context_payload = build_alert_context_payload(alert_details, agent_data)
86
- if ioc_payload:
87
- logger.info(f"Alert has IoC: {ioc_payload}")
88
- return IrisAlertPayload(
89
- alert_title=alert_details._source.rule_description,
90
- alert_description=alert_details._source.rule_description,
91
- alert_source="CoPilot",
92
- assets=[asset_payload],
93
- alert_status_id=3,
94
- alert_severity_id=5,
95
- alert_customer_id=1,
96
- alert_source_content=alert_details._source,
97
- alert_context=context_payload,
98
- alert_iocs=[ioc_payload],
99
- )
100
- else:
101
- logger.info("Alert does not have IoC")
102
- return IrisAlertPayload(
103
- alert_title=alert_details._source.rule_description,
104
- alert_description=alert_details._source.rule_description,
105
- alert_source="CoPilot",
106
- assets=[asset_payload],
107
- alert_status_id=3,
108
- alert_severity_id=5,
109
- alert_customer_id=1,
110
- alert_source_content=alert_details._source,
111
- alert_context=context_payload,
112
- )
148
+async def build_alert_payload(
149
+ alert_details: GenericAlertModel,
150
+ agent_data,
151
+ ioc_payload: Optional[IrisIoc],
152
+ session: AsyncSession,
153
+) -> IrisAlertPayload:
154
+ asset_payload = await build_asset_payload(agent_data, alert_details)
155
+ context_payload = await build_alert_context_payload(alert_details=alert_details, agent_data=agent_data, session=session)
156
+ timefield = (await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)).timefield
157
+ # Get the timefield value from the alert_details
158
+ if hasattr(alert_details, timefield):
159
+ alert_details.time_field = getattr(alert_details, timefield)
160
+ logger.info(f"Alert has context: {context_payload}")
161
+ try:
162
+ if ioc_payload:
163
+ logger.info(f"Alert has IoC: {ioc_payload}")
164
+ return IrisAlertPayload(
165
+ alert_title=alert_details._source.rule_description,
166
+ alert_source_link=await construct_alert_source_link(alert_details, session=session),
167
+ alert_description=alert_details._source.rule_description,
168
+ alert_source="CoPilot",
169
+ assets=[asset_payload],
170
+ alert_status_id=3,
171
+ alert_severity_id=5,
172
+ alert_customer_id=(
173
+ await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)
174
+ ).iris_customer_id,
175
+ alert_source_content=alert_details._source,
176
+ alert_context=context_payload,
177
+ alert_iocs=[ioc_payload],
178
+ alert_source_event_time=alert_details.time_field,
179
+ )
180
+ else:
181
+ logger.info("Alert does not have IoC")
182
+ return IrisAlertPayload(
183
+ alert_title=alert_details._source.rule_description,
184
+ alert_source_link=await construct_alert_source_link(alert_details, session=session),
185
+ alert_description=alert_details._source.rule_description,
186
+ alert_source="CoPilot",
187
+ assets=[asset_payload],
188
+ alert_status_id=3,
189
+ alert_severity_id=5,
190
+ alert_customer_id=(
191
+ await get_customer_alert_settings(customer_code=alert_details._source.agent_labels_customer, session=session)
192
+ ).iris_customer_id,
193
+ alert_source_content=alert_details._source,
194
+ alert_context=context_payload,
195
+ alert_source_event_time=alert_details.time_field,
196
+ )
197
+ except Exception as e:
198
+ logger.error(f"Failed to build alert payload: {e}")
199
+ raise HTTPException(status_code=500, detail=f"Failed to build alert payload: {e}")
200
201
202
async def construct_soc_alert_url(root_url: str, soc_alert_id: int) -> str:
248
async def create_alert(alert: CreateAlertRequest, session: AsyncSession) -> CreateAlertResponse:
249
logger.info(f"Creating alert {alert.alert_id} in IRIS")
250
alert_details = await get_single_alert_details(alert_details=alert)
164
- agent_data = await get_agent_data(session, agent_id=alert_details._source.agent_id)
165
- alert_details.asset_type_id = get_asset_type_id(os=agent_data.os)
166
- ioc_payload = build_ioc_payload(alert_details)
167
- iris_alert_payload = build_alert_payload(alert_details, agent_data, ioc_payload)
251
+ logger.info(f"Alert details: {alert_details}")
252
+ if await is_customer_code_valid(customer_code=alert_details._source.agent_labels_customer, session=session) is False:
253
+ logger.info(f"Invalid customer_code: {alert_details._source.agent_labels_customer}")
254
+ raise HTTPException(status_code=200, detail="Invalid customer_code, or the customer is not configured for alert creation.")
255
+ agent_data = await get_agent(agent_id=alert_details._source.agent_id, db=session)
256
+ ioc_payload = await build_ioc_payload(alert_details=alert_details)
257
+ iris_alert_payload = await build_alert_payload(
258
+ alert_details=alert_details,
259
+ agent_data=agent_data,
260
+ ioc_payload=ioc_payload,
261
+ session=session,
262
+ )
263
client, alert_client = await initialize_client_and_alert("DFIR-IRIS")
264
result = await fetch_and_validate_data(client, alert_client.add_alert, iris_alert_payload.to_dict())
265
es_client = await create_wazuh_indexer_client("Wazuh-Indexer")