main
py 586 lines 22.8 KB
Raw
1 import json
2 import os
3 from datetime import datetime
4
5 import aiofiles
6 from fastapi import HTTPException
7 from loguru import logger
8 from sqlalchemy import and_
9 from sqlalchemy import update
10 from sqlalchemy.ext.asyncio import AsyncSession
11
12 from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest
13 from app.connectors.grafana.schema.dashboards import DefenderForEndpointDashboard
14 from app.connectors.grafana.services.dashboards import provision_dashboards
15 from app.connectors.grafana.utils.universal import create_grafana_client
16 from app.connectors.graylog.services.collector import (
17 get_content_pack_id_by_content_pack_name,
18 )
19 from app.connectors.graylog.services.collector import get_input_id_by_input_name
20 from app.connectors.graylog.services.collector import get_stream_id_by_stream_name
21 from app.connectors.graylog.services.streams import assign_stream_to_index
22 from app.connectors.graylog.utils.universal import send_post_request
23 from app.connectors.wazuh_indexer.services.monitoring import (
24 output_shard_number_to_be_set_based_on_nodes,
25 )
26 from app.customer_provisioning.schema.grafana import GrafanaDatasource
27 from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse
28 from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse
29 from app.customer_provisioning.schema.graylog import TimeBasedIndexSet
30 from app.customer_provisioning.schema.provision import ProvisionNewCustomer
31 from app.customer_provisioning.services.grafana import create_grafana_folder
32 from app.customer_provisioning.services.grafana import get_opensearch_version
33 from app.customers.routes.customers import get_customer_meta
34 from app.integrations.defender_for_endpoint.schema.provision import (
35 DefenderForEndpointCustomerDetails,
36 )
37 from app.integrations.defender_for_endpoint.schema.provision import (
38 ProvisionDefenderForEndpointAuthKeys,
39 )
40 from app.integrations.defender_for_endpoint.schema.provision import (
41 ProvisionDefenderForEndpointResponse,
42 )
43 from app.integrations.models.customer_integration_settings import CustomerIntegrations
44 from app.network_connectors.models.network_connectors import (
45 CustomerNetworkConnectorsMeta,
46 )
47 from app.stack_provisioning.graylog.schema.provision import ContentPackKeywords
48 from app.stack_provisioning.graylog.schema.provision import (
49 ProvisionNetworkContentPackRequest,
50 )
51 from app.stack_provisioning.graylog.services.provision import (
52 provision_content_pack_network_connector,
53 )
54 from app.utils import get_connector_attribute
55 from app.utils import get_customer_meta_attribute
56
57
58 #### ! GRAYLOG ! ####
59 async def build_index_set_config(request: DefenderForEndpointCustomerDetails) -> TimeBasedIndexSet:
60 """
61 Build the configuration for a time-based index set.
62
63 Args:
64 request (DefenderForEndpointCustomerDetails): The request object containing customer information.
65
66 Returns:
67 TimeBasedIndexSet: The configured time-based index set.
68 """
69 return TimeBasedIndexSet(
70 title=f"{request.customer_name} - DEFENDER FOR ENDPOINT EVENTS",
71 description=f"{request.customer_name} - DEFENDER FOR ENDPOINT EVENTS",
72 index_prefix=f"defender-for-endpoint-{request.customer_code}",
73 rotation_strategy_class="org.graylog2.indexer.rotation.strategies.TimeBasedRotationStrategy",
74 rotation_strategy={
75 "type": "org.graylog2.indexer.rotation.strategies.TimeBasedRotationStrategyConfig",
76 "rotation_period": "P1D",
77 "rotate_empty_index_set": False,
78 "max_rotation_period": None,
79 },
80 retention_strategy_class="org.graylog2.indexer.retention.strategies.DeletionRetentionStrategy",
81 retention_strategy={
82 "type": "org.graylog2.indexer.retention.strategies.DeletionRetentionStrategyConfig",
83 "max_number_of_indices": request.hot_data_retention,
84 },
85 creation_date=datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
86 index_analyzer="standard",
87 shards=await output_shard_number_to_be_set_based_on_nodes(),
88 replicas=request.index_replicas,
89 index_optimization_max_num_segments=1,
90 index_optimization_disabled=False,
91 writable=True,
92 field_type_refresh_interval=5000,
93 )
94
95
96 # Function to send the POST request and handle the response
97 async def send_index_set_creation_request(
98 index_set: TimeBasedIndexSet,
99 ) -> GraylogIndexSetCreationResponse:
100 """
101 Sends a request to create an index set in Graylog.
102
103 Args:
104 index_set (TimeBasedIndexSet): The index set to be created.
105
106 Returns:
107 GraylogIndexSetCreationResponse: The response from Graylog after creating the index set.
108 """
109 json_index_set = json.dumps(index_set.model_dump())
110 logger.info(f"json_index_set set: {json_index_set}")
111 response_json = await send_post_request(
112 endpoint="/api/system/indices/index_sets",
113 data=index_set.model_dump(),
114 )
115 return GraylogIndexSetCreationResponse(**response_json)
116
117
118 # Refactored create_index_set function
119 async def create_index_set(
120 request: ProvisionNewCustomer,
121 ) -> GraylogIndexSetCreationResponse:
122 """
123 Creates an index set for a new customer.
124
125 Args:
126 request (ProvisionNewCustomer): The request object containing the customer information.
127
128 Returns:
129 GraylogIndexSetCreationResponse: The response object containing the result of the index set creation.
130 """
131 logger.info(f"Creating index set for customer {request.customer_name}")
132 index_set_config = await build_index_set_config(request)
133 return await send_index_set_creation_request(index_set_config)
134
135
136 async def provision_content_pack(customer_details):
137 """
138 Provisions a content pack for a customer.
139
140 Args:
141 customer_details (CustomerDetails): The details of the customer.
142
143 Returns:
144 ContentPack: The provisioned content pack.
145 """
146 return await provision_content_pack_network_connector(
147 content_pack_request=ProvisionNetworkContentPackRequest(
148 content_pack_name="DEFENDER_FOR_ENDPOINT",
149 keywords=ContentPackKeywords(
150 customer_name=customer_details.customer_name,
151 customer_code=customer_details.customer_code,
152 protocol_type=customer_details.protocal_type,
153 syslog_port=customer_details.syslog_port,
154 ),
155 ),
156 )
157
158
159 async def get_stream_and_index_ids(customer_details):
160 """
161 Retrieves the stream ID and index ID for a given customer.
162
163 Args:
164 customer_details (CustomerDetails): The details of the customer.
165
166 Returns:
167 tuple: A tuple containing the stream ID and index ID.
168 """
169 stream_id = await get_stream_id_by_stream_name(stream_name=f"{customer_details.customer_name} - DEFENDER FOR ENDPOINT LOGS AND EVENTS")
170 index_id = (await create_index_set(request=customer_details)).data.id
171 content_pack_stream_id = await get_content_pack_id_by_content_pack_name(
172 content_pack_name=f"{customer_details.customer_name}_DEFENDER_FOR_ENDPOINT_STREAM",
173 )
174 if customer_details.protocal_type == "TCP":
175 content_pack_input_id = await get_content_pack_id_by_content_pack_name(
176 content_pack_name=f"{customer_details.customer_name}_DEFENDER_FOR_ENDPOINT_INPUT_TCP",
177 )
178 return stream_id, index_id, content_pack_stream_id, content_pack_input_id
179
180
181 #### ! GRAFANA ! ####
182 async def create_grafana_datasource(
183 customer_code: str,
184 session: AsyncSession,
185 ) -> GrafanaDataSourceCreationResponse:
186 """
187 Creates a Grafana datasource for the specified customer.
188
189 Args:
190 customer_code (str): The customer code.
191 session (AsyncSession): The async session.
192
193 Returns:
194 GrafanaDataSourceCreationResponse: The response containing the created datasource details.
195 """
196 logger.info("Creating Grafana datasource")
197 grafana_client = await create_grafana_client("Grafana")
198 # Switch to the newly created organization
199 grafana_client.user.switch_actual_user_organisation(
200 (await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id,
201 )
202 datasource_payload = GrafanaDatasource(
203 name="DEFENDER FOR ENDPOINT",
204 type="grafana-opensearch-datasource",
205 typeName="OpenSearch",
206 access="proxy",
207 url=await get_connector_attribute(
208 connector_id=1,
209 column_name="connector_url",
210 session=session,
211 ),
212 database=f"defender-for-endpoint-{customer_code}*",
213 basicAuth=True,
214 basicAuthUser=await get_connector_attribute(
215 connector_id=1,
216 column_name="connector_username",
217 session=session,
218 ),
219 secureJsonData={
220 "basicAuthPassword": await get_connector_attribute(
221 connector_id=1,
222 column_name="connector_password",
223 session=session,
224 ),
225 },
226 isDefault=False,
227 jsonData={
228 "database": f"defender-for-endpoint-{customer_code}*",
229 "flavor": "opensearch",
230 "includeFrozen": False,
231 "logLevelField": "severity",
232 "logMessageField": "summary",
233 "maxConcurrentShardRequests": 5,
234 "pplEnabled": True,
235 "timeField": "timestamp",
236 "tlsSkipVerify": True,
237 "version": await get_opensearch_version(),
238 },
239 readOnly=True,
240 )
241 results = grafana_client.datasource.create_datasource(
242 datasource=datasource_payload.model_dump(),
243 )
244 return GrafanaDataSourceCreationResponse(**results)
245
246
247 async def create_customer_network_connector_meta(
248 customer_details,
249 stream_id,
250 index_id,
251 content_pack_stream_id,
252 content_pack_input_id,
253 session,
254 ):
255 """
256 Create a CustomerNetworkConnectorsMeta object with the provided details.
257
258 Args:
259 customer_details (CustomerDetails): Details of the customer.
260 stream_id (int): ID of the Graylog stream.
261 index_id (int): ID of the Graylog index.
262 session (Session): Database session.
263
264 Returns:
265 CustomerNetworkConnectorsMeta: The created CustomerNetworkConnectorsMeta object.
266 """
267 return CustomerNetworkConnectorsMeta(
268 customer_code=customer_details.customer_code,
269 network_connector_name="DefenderForEndpoint",
270 graylog_stream_id=stream_id,
271 graylog_input_id=(await get_input_id_by_input_name(input_name=f"{customer_details.customer_name} - DEFENDER FOR ENDPOINT")),
272 graylog_pipeline_id="not_set",
273 graylog_content_pack_input_id=content_pack_input_id,
274 graylog_content_pack_stream_id=content_pack_stream_id,
275 grafana_org_id=(
276 await get_customer_meta_attribute(
277 session=session,
278 customer_code=customer_details.customer_code,
279 column_name="customer_meta_grafana_org_id",
280 )
281 ),
282 graylog_index_id=index_id,
283 grafana_dashboard_folder_id=None,
284 grafana_datasource_uid=None,
285 )
286
287
288 async def validate_grafana_organization_id(customer_code, session):
289 """
290 Validate the Grafana organization ID for the customer.
291
292 Args:
293 customer_code (str): The customer code.
294 session (Session): Database session.
295
296 Returns:
297 int: The Grafana organization ID.
298 """
299 return await get_customer_meta_attribute(session=session, customer_code=customer_code, column_name="customer_meta_grafana_org_id")
300
301
302 async def provision_defender_for_endpoint(
303 customer_details: DefenderForEndpointCustomerDetails,
304 keys: ProvisionDefenderForEndpointAuthKeys,
305 session: AsyncSession,
306 ) -> ProvisionDefenderForEndpointResponse:
307 """
308 Provisions a Defender For Endpoint customer by performing the following steps:
309 1. Provisions the content pack for the customer.
310 2. Retrieves the stream and index IDs for the customer.
311 3. Creates customer network connector metadata.
312 4. Assigns the stream to the index.
313 5. Inserts the customer network connector metadata into the database.
314 6. Creates a directory for the customer to store the docker compose and falconhose cfg.
315
316 Args:
317 customer_details (DefenderForEndpointCustomerDetails): The details of the DefenderForEndpoint customer.
318 keys (ProvisionDefenderForEndpointAuthKeys): The keys required for provisioning.
319 session (AsyncSession): The database session.
320
321 Returns:
322 None
323 """
324 # If customer name contains a space, replace it with a _
325 if " " in customer_details.customer_name:
326 customer_details.customer_name = customer_details.customer_name.replace(" ", "_")
327 if await validate_grafana_organization_id(customer_details.customer_code, session) is None:
328 raise HTTPException(status_code=404, detail="Grafana organization ID not found. Please provision Grafana for the customer first.")
329 logger.info(f"Provisioning Defender For Endpoint for customer {customer_details.customer_name}")
330 await provision_content_pack(customer_details)
331 stream_id, index_id, content_pack_stream_id, content_pack_input_id = await get_stream_and_index_ids(customer_details)
332 customer_network_connector_meta = await create_customer_network_connector_meta(
333 customer_details,
334 stream_id,
335 index_id,
336 content_pack_stream_id,
337 content_pack_input_id,
338 session,
339 )
340 await assign_stream_to_index(stream_id=stream_id, index_id=index_id)
341 # Grafana Deployment
342 customer_network_connector_meta.grafana_datasource_uid = (
343 await create_grafana_datasource(
344 customer_code=customer_details.customer_code,
345 session=session,
346 )
347 ).datasource.uid
348 grafana_folder = await create_grafana_folder(
349 organization_id=(
350 await get_customer_meta(
351 customer_details.customer_code,
352 session,
353 )
354 ).customer_meta.customer_meta_grafana_org_id,
355 folder_title="DEFENDER FOR ENDPOINT",
356 )
357 await provision_dashboards(
358 DashboardProvisionRequest(
359 dashboards=[dashboard.name for dashboard in DefenderForEndpointDashboard],
360 organizationId=(
361 await get_customer_meta(
362 customer_details.customer_code,
363 session,
364 )
365 ).customer_meta.customer_meta_grafana_org_id,
366 folderId=grafana_folder.id,
367 datasourceUid=customer_network_connector_meta.grafana_datasource_uid,
368 ),
369 )
370 customer_network_connector_meta.grafana_dashboard_folder_id = grafana_folder.id
371 await insert_into_customer_network_connectors_meta_table(
372 customer_network_connectors_meta=customer_network_connector_meta,
373 session=session,
374 )
375 await create_customer_directory_if_needed(customer_name=customer_details.customer_name)
376 await create_customer_data_directory(customer_name=customer_details.customer_name)
377 file = await load_and_replace_docker_compose(customer_name=customer_details.customer_name)
378 await save_uploaded_file(
379 file=file,
380 filename=f"{customer_details.customer_name}_docker-compose-defender-for-endpoint.yml",
381 customer_name=customer_details.customer_name,
382 )
383 await load_and_replace_filebeat_cfg(customer_details=customer_details, keys=keys, session=session)
384
385 await update_customer_integration_table(
386 customer_code=customer_details.customer_code,
387 session=session,
388 )
389
390 return ProvisionDefenderForEndpointResponse(
391 message="Defender For Endpoint for customer provisioned successfully",
392 success=True,
393 )
394
395
396 async def insert_into_customer_network_connectors_meta_table(
397 customer_network_connectors_meta: CustomerNetworkConnectorsMeta,
398 session: AsyncSession,
399 ) -> None:
400 """
401 Insert the customer network connectors meta into the database.
402
403 Args:
404 customer_network_connectors_meta (CustomerNetworkConnectorsMeta): The customer network connectors meta to insert.
405 session (AsyncSession): The async session object for database operations.
406
407 Returns:
408 None
409 """
410 logger.info("Inserting customer network connectors meta into the database")
411 session.add(customer_network_connectors_meta)
412 await session.commit()
413 logger.info("Customer network connectors meta inserted successfully")
414 return None
415
416
417 # ! Add the docker-compose.yml file to the `data` folder
418 project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
419 UPLOAD_FOLDER = os.path.join(project_root, "data")
420
421
422 async def create_customer_directory_if_needed(customer_name: str):
423 """
424 Create a directory for the customer in the UPLOAD_FOLDER if it doesn't exist.
425
426 Args:
427 customer_name (str): The name of the customer.
428 """
429 # Create the path to the customer's directory
430 # If customer name contains a space, replace it with a _
431 if " " in customer_name:
432 customer_name = customer_name.replace(" ", "_")
433 customer_directory = os.path.join(UPLOAD_FOLDER, customer_name)
434 # Check if the directory exists
435 if not os.path.exists(customer_directory):
436 # If it doesn't exist, create it
437 os.makedirs(customer_directory)
438
439
440 async def create_customer_data_directory(customer_name: str):
441 """
442 Create a data directory within the customer's directory in the UPLOAD_FOLDER
443
444 This function ensures there's a dedicated data directory for customer-specific files.
445 For example, if the customer directory is /opt/CoPilot/data/Customer1,
446 this will create /opt/CoPilot/data/Customer1/data.
447
448 Args:
449 customer_name (str): The name of the customer.
450
451 Returns:
452 str: The path to the created data directory
453 """
454 # Normalize customer name (replace spaces with underscores)
455 if " " in customer_name:
456 customer_name = customer_name.replace(" ", "_")
457
458 # Create the path to the customer's directory
459 customer_directory = os.path.join(UPLOAD_FOLDER, customer_name)
460
461 # Create the path to the data directory within the customer's directory
462 customer_data_directory = os.path.join(customer_directory, "data")
463
464 # Check if the customer directory exists, if not create it
465 if not os.path.exists(customer_directory):
466 os.makedirs(customer_directory)
467
468 # Check if the data directory exists, if not create it
469 if not os.path.exists(customer_data_directory):
470 os.makedirs(customer_data_directory)
471 logger.info(f"Created data directory for customer {customer_name}: {customer_data_directory}")
472 else:
473 logger.info(f"Data directory for customer {customer_name} already exists: {customer_data_directory}")
474
475 return customer_data_directory
476
477
478 async def load_and_replace_docker_compose(customer_name: str):
479 """
480 Load the docker-compose.yml file and replace the placeholder with the customer name.
481
482 Args:
483 customer_name (str): The name of the customer.
484
485 Returns:
486 str: The content of the docker-compose.yml file with the placeholder replaced.
487 """
488 # Get the current directory:
489 current_directory = os.path.dirname(os.path.abspath(__file__))
490 # Go up one level
491 parent_directory = os.path.dirname(current_directory)
492 # If customer name contains a space, replace it with a _
493 if " " in customer_name:
494 customer_name = customer_name.replace(" ", "_")
495 # Open the docker-compose.yml file and read the content
496 with open(os.path.join(parent_directory, "templates", "docker-compose.yml"), "r") as file:
497 data = file.read()
498 data = data.replace("CUSTOMER_NAME", customer_name)
499 return data
500
501
502 async def save_uploaded_file(file, filename, customer_name):
503 """
504 Save the uploaded file to the server.
505
506 Args:
507 file: The file to save.
508 filename: The name of the file.
509
510 Returns:
511 str: The path to the saved file.
512 """
513 # If customer name contains a space, replace it with a _
514 if " " in customer_name:
515 customer_name = customer_name.replace(" ", "_")
516 customer_upload_folder = os.path.join(UPLOAD_FOLDER, customer_name)
517 async with aiofiles.open(os.path.join(customer_upload_folder, filename), "wb") as f:
518 await f.write(file.encode())
519 return os.path.join(customer_upload_folder, filename)
520
521
522 async def load_and_replace_filebeat_cfg(
523 customer_details: DefenderForEndpointCustomerDetails,
524 keys: ProvisionDefenderForEndpointAuthKeys,
525 session: AsyncSession,
526 ):
527 """
528 Load the filebeat.yml file and replace the placeholders with the customer details.
529
530 Args:
531 customer_details (DefenderForEndpointCustomerDetails): The details of the customer.
532 keys (ProvisionDefenderForEndpointAuthKeys): The authentication keys for DefenderForEndpoint.
533
534 Returns:
535 str: The content of the filebeat.yml file with the placeholders replaced.
536 """
537 # Get the current directory:
538 current_directory = os.path.dirname(os.path.abspath(__file__))
539 # Go up one level
540 parent_directory = os.path.dirname(current_directory)
541 connector_url = str(await get_connector_attribute(connector_id=3, column_name="connector_url", session=session))
542 connector_url = connector_url.replace("https://", "").replace("http://", "").replace(":9000", "")
543 # Open the filebeat.yml file and read the content
544 with open(os.path.join(parent_directory, "templates", "filebeat.yml"), "r") as file:
545 data = file.read()
546 data = data.replace("REPLACE_TENANT_ID", keys.TENANT_ID)
547 data = data.replace("REPLACE_CLIENT_ID", keys.CLIENT_ID)
548 data = data.replace("REPLACE_CLIENT_SECRET", keys.CLIENT_SECRET)
549 data = data.replace("REPLACE_SYSLOG_HOST", connector_url)
550 data = data.replace("REPLACE_SYSLOG_PORT", keys.SYSLOG_PORT)
551 # Save the file
552 # If customer name contains a space, replace it with a _
553 if " " in customer_details.customer_name:
554 customer_details.customer_name = customer_details.customer_name.replace(" ", "_")
555 customer_upload_folder = os.path.join(UPLOAD_FOLDER, customer_details.customer_name)
556 async with aiofiles.open(os.path.join(customer_upload_folder, "filebeat.yml"), "w") as f:
557 await f.write(data)
558 return os.path.join(customer_upload_folder, "filebeat.yml")
559
560
561 async def update_customer_integration_table(
562 customer_code: str,
563 session: AsyncSession,
564 ) -> None:
565 """
566 Updates the `customer_integrations` table to set the `deployed` column to True where the `customer_code`
567 matches the given customer code and the `integration_service_name` is "DefenderForEndpoint".
568
569 Args:
570 customer_code (str): The customer code.
571 session (AsyncSession): The async session object for making HTTP requests.
572 """
573 logger.info(f"Updating customer integrations table for customer {customer_code}")
574 await session.execute(
575 update(CustomerIntegrations)
576 .where(
577 and_(
578 CustomerIntegrations.customer_code == customer_code,
579 CustomerIntegrations.integration_service_name == "DefenderForEndpoint",
580 ),
581 )
582 .values(deployed=True),
583 )
584 await session.commit()
585
586 return None