main
py 1,130 lines 39.5 KB
Raw
1 from datetime import datetime
2 from datetime import timedelta
3 from enum import Enum
4 from typing import Any
5 from typing import Dict
6 from typing import List
7 from typing import Optional
8 from typing import Union
9
10 import requests
11 from fastapi import APIRouter
12 from fastapi import Depends
13 from fastapi import HTTPException
14 from fastapi import Request
15 from fastapi import Security
16 from fastapi.exceptions import RequestValidationError
17 from loguru import logger
18 from pydantic import BaseModel
19 from pydantic import Field
20 from pydantic import field_validator
21 from pydantic import model_validator
22 from sqlalchemy.ext.asyncio import AsyncSession
23 from sqlalchemy.future import select
24 from sqlalchemy.orm import joinedload
25
26 from app.auth.services.universal import find_user
27 from app.auth.utils import AuthHandler
28 from app.connectors.utils import get_connector_info_from_db
29 from app.customer_provisioning.models.default_settings import (
30 CustomerProvisioningDefaultSettings,
31 )
32 from app.db.all_models import Connectors
33 from app.db.db_session import get_db
34 from app.db.db_session import get_db_session
35 from app.db.db_session import get_session
36 from app.db.universal_models import CustomersMeta
37 from app.db.universal_models import LogEntry
38 from app.integrations.alert_creation_settings.models.alert_creation_settings import (
39 AlertCreationEventConfig,
40 )
41 from app.integrations.alert_creation_settings.models.alert_creation_settings import (
42 AlertCreationSettings,
43 )
44 from app.integrations.alert_creation_settings.models.alert_creation_settings import (
45 EventOrder,
46 )
47
48
49 ################## ! 422 VALIDATION ERROR TYPES FOR PYDANTIC VALUE ERROR RESPONSE ! ##################
50 class ErrorType(str, Enum):
51 # Legacy pydantic 1 codes — kept so any consumer reading these by string
52 # value continues to work. Pydantic 2 uses the *_V2 codes below.
53 PASSWORD_REGEX = "value_error.str.regex"
54 TIME_RANGE = "value_error.time_range"
55 JSON_INVALID = "json_invalid"
56 MIN_LENGTH = "value_error.any_str.min_length"
57 MAX_LENGTH = "value_error.any_str.max_length"
58 NOT_A_NUMBER = "value_error.number.not_a_number"
59 TOO_SMALL = "value_error.number.too_small"
60 TOO_LARGE = "value_error.number.too_large"
61 INVALID_DATETIME = "value_error.datetime"
62 INVALID_DATE = "value_error.date"
63 MIN_ITEMS = "value_error.list.min_items"
64 MAX_ITEMS = "value_error.list.max_items"
65 UNIQUE = "value_error.list.unique"
66 NONE_NOT_ALLOWED = "value_error.none.not_allowed"
67 MISSING = "value_error.missing"
68 GENERAL = "value_error"
69 INVALID_ENUM = "type_error.enum"
70 # Pydantic 2 codes (renamed in #849). Add more here as users surface them.
71 STRING_TOO_SHORT = "string_too_short"
72 STRING_TOO_LONG = "string_too_long"
73 STRING_PATTERN_MISMATCH = "string_pattern_mismatch"
74 INT_PARSING = "int_parsing"
75 GREATER_THAN = "greater_than"
76 GREATER_THAN_EQUAL = "greater_than_equal"
77 LESS_THAN = "less_than"
78 LESS_THAN_EQUAL = "less_than_equal"
79 DATETIME_PARSING = "datetime_parsing"
80 DATE_PARSING = "date_parsing"
81 ENUM = "enum"
82 MISSING_V2 = "missing"
83
84
85 class ValidationErrorItem(BaseModel):
86 field: str
87 error_type: ErrorType
88 message: str = None # Initialize as None
89
90 @model_validator(mode="after")
91 def set_message(self):
92 error_messages = {
93 ErrorType.PASSWORD_REGEX: "Password must be at least 8 characters long and contain at least one uppercase letter, one lowercase letter, one number, and one "
94 "special character.",
95 ErrorType.TIME_RANGE: "Invalid time range. Use 'h' for hours, 'd' for days, and 'w' for weeks.",
96 ErrorType.JSON_INVALID: "Invalid JSON. Please check your JSON syntax and try again.",
97 ErrorType.MIN_LENGTH: "Value is shorter than minimum length.",
98 ErrorType.MAX_LENGTH: "Value is longer than maximum length.",
99 ErrorType.NOT_A_NUMBER: "Input is not a number.",
100 ErrorType.TOO_SMALL: "Value is too small.",
101 ErrorType.TOO_LARGE: "Value is too large.",
102 ErrorType.INVALID_DATETIME: "Invalid datetime format.",
103 ErrorType.INVALID_DATE: "Invalid date format.",
104 ErrorType.MIN_ITEMS: "Number of items is less than minimum.",
105 ErrorType.MAX_ITEMS: "Number of items is more than maximum.",
106 ErrorType.UNIQUE: "Items are not unique.",
107 ErrorType.NONE_NOT_ALLOWED: "None is not an allowed value.",
108 ErrorType.MISSING: "Missing data for required field.",
109 ErrorType.GENERAL: "Invalid value.",
110 ErrorType.INVALID_ENUM: "Value is not a valid enumeration member.",
111 # Pydantic 2 codes — same human messages as their v1 equivalents.
112 ErrorType.STRING_TOO_SHORT: "Value is shorter than minimum length.",
113 ErrorType.STRING_TOO_LONG: "Value is longer than maximum length.",
114 ErrorType.STRING_PATTERN_MISMATCH: "Value does not match the required pattern.",
115 ErrorType.INT_PARSING: "Input is not a valid integer.",
116 ErrorType.GREATER_THAN: "Value is too small.",
117 ErrorType.GREATER_THAN_EQUAL: "Value is too small.",
118 ErrorType.LESS_THAN: "Value is too large.",
119 ErrorType.LESS_THAN_EQUAL: "Value is too large.",
120 ErrorType.DATETIME_PARSING: "Invalid datetime format.",
121 ErrorType.DATE_PARSING: "Invalid date format.",
122 ErrorType.ENUM: "Value is not a valid enumeration member.",
123 ErrorType.MISSING_V2: "Missing data for required field.",
124 }
125 if self.error_type in error_messages:
126 self.message = error_messages[self.error_type]
127 return self
128
129
130 class ValidationErrorResponse(BaseModel):
131 message: str
132 details: List[ValidationErrorItem]
133
134
135 ################## ! LOGGING TO `log_entry` table ! ##################
136 # #######! MODELS !########
137 class LogEntryModel(BaseModel):
138 event_type: str = Field(..., examples=["Info"], description="Event type")
139 user_id: Optional[int] = Field(None, examples=[1], description="User ID")
140 route: str = Field(..., examples=["/wazuh_indexer/health"], description="Route")
141 method: str = Field(..., examples=["GET"], description="Method")
142 status_code: int = Field(..., examples=[200], description="Status code")
143 message: str = Field(..., examples=["Route accessed"], description="Message")
144 additional_info: Optional[str] = Field(
145 None,
146 examples=["Additional details here"],
147 description="Additional info",
148 )
149
150
151 class LogRetrieveModel(LogEntryModel):
152 timestamp: datetime = Field(..., examples=[datetime.now()], description="Timestamp")
153
154
155 class LogsResponse(BaseModel):
156 logs: List[LogRetrieveModel]
157 success: bool
158 message: str
159
160
161 class EventType(str, Enum):
162 INFO = "Info"
163 ERROR = "Error"
164 # Add other event types as needed
165
166
167 class TimeRangeModel(BaseModel):
168 time_range: Union[str, int] = Field(
169 "1d",
170 description="Time range to fetch logs for, e.g., 1, 1h, 1d, 1w",
171 )
172
173 @field_validator("time_range")
174 @classmethod
175 def validate_time_range(cls, value):
176 """
177 Validate the time range value.
178
179 Args:
180 value (int or str): The time range value to be validated.
181
182 Returns:
183 str: The validated time range value.
184
185 Raises:
186 RequestValidationError: If the time range value is invalid.
187
188 """
189 try:
190 if isinstance(value, int):
191 if value < 1 or value > 7:
192 raise RequestValidationError(
193 [
194 {
195 "loc": ("time_range",),
196 "msg": "The integer part should be between 1 and 7.",
197 "type": "value_error.time_range",
198 },
199 ],
200 )
201 return f"{value}d" # convert integer to day representation
202
203 elif isinstance(value, str):
204 unit = value[-1]
205 int_part = int(value[:-1])
206
207 if unit not in ["h", "d", "w"]:
208 raise RequestValidationError(
209 [
210 {
211 "loc": ("time_range",),
212 "msg": "Invalid unit. Use 'h' for hours, 'd' for days, and 'w' for weeks.",
213 "type": "value_error.time_range",
214 },
215 ],
216 )
217
218 if int_part <= 0:
219 raise RequestValidationError(
220 [
221 {
222 "loc": ("time_range",),
223 "msg": "The integer part should be greater than 0.",
224 "type": "value_error.time_range",
225 },
226 ],
227 )
228
229 if unit == "w" and int_part > 1:
230 raise RequestValidationError(
231 [
232 {
233 "loc": ("time_range",),
234 "msg": "The maximum allowed time range is 1 week.",
235 "type": "value_error.time_range",
236 },
237 ],
238 )
239 return value
240
241 else:
242 raise RequestValidationError(
243 [
244 {
245 "loc": ("time_range",),
246 "msg": "Invalid type. Time range should be either an integer or a string.",
247 "type": "value_error.time_range",
248 },
249 ],
250 )
251
252 except ValueError:
253 raise RequestValidationError(
254 [
255 {
256 "loc": ("time_range",),
257 "msg": "Invalid format. Time range should be an integer followed by a unit (h, d, w).",
258 "type": "value_error.time_range",
259 },
260 ],
261 )
262
263
264 # ########! LOGGER CLASS !#########
265 class Logger:
266 def __init__(self, session: AsyncSession, auth_handler: AuthHandler):
267 self.session = session
268 self.auth_handler = auth_handler
269
270 async def get_user_id_from_request(self, request: Request):
271 """
272 Retrieves the user ID from the request object.
273
274 Args:
275 request (Request): The request object.
276
277 Returns:
278 int or None: The user ID if found, None otherwise.
279 """
280 auth_header = request.headers.get("Authorization")
281 if auth_header:
282 try:
283 token = auth_header.split(" ")[1] # Better split by space and take the second part
284 except IndexError:
285 raise HTTPException(status_code=401, detail="Invalid token")
286 username, _ = self.auth_handler.decode_token(token)
287 user = await find_user(username) # Correctly using await for an async call
288 if user:
289 return user.id
290 return None
291
292 async def insert_log_entry(self, log_entry_model: LogEntryModel):
293 """
294 Inserts a log entry into the database.
295
296 Args:
297 log_entry_model (LogEntryModel): The log entry model to be inserted.
298
299 Returns:
300 None
301 """
302 log_entry = LogEntry(**log_entry_model.model_dump())
303 self.session.add(log_entry)
304 await self.session.commit()
305
306 async def log_route_access(self, user_id, request: Request, response):
307 """
308 Logs the access of a route.
309
310 Args:
311 user_id (int): The ID of the user accessing the route.
312 request (Request): The request object containing information about the request.
313 response: The response object containing information about the response.
314
315 Returns:
316 None
317 """
318 log_entry_model = LogEntryModel(
319 event_type="Info",
320 user_id=user_id,
321 route=str(request.url),
322 method=request.method,
323 status_code=response.status_code,
324 message="Route accessed",
325 )
326 await self.insert_log_entry(log_entry_model)
327
328 async def log_error(
329 self,
330 user_id,
331 request: Request,
332 exception: Exception,
333 additional_info: Optional[str] = None,
334 ):
335 """
336 Logs an error event with the provided information.
337
338 Args:
339 user_id (int): The ID of the user associated with the error.
340 request (Request): The request object that triggered the error.
341 exception (Exception): The exception that occurred.
342 additional_info (Optional[str], optional): Additional information about the error. Defaults to None.
343 """
344 log_entry_model = LogEntryModel(
345 event_type="Error",
346 user_id=user_id,
347 route=str(request.url),
348 method=request.method,
349 status_code=500, # Internal Server Error
350 message=str(exception),
351 additional_info=additional_info,
352 )
353 await self.insert_log_entry(log_entry_model)
354
355 async def log_and_raise_http_error(
356 self,
357 user_id,
358 request: Request,
359 exception: Exception,
360 ):
361 """
362 Logs the error, including the user ID, request details, and the exception,
363 and raises an HTTPException with a status code of 500 (Internal Server Error).
364
365 Args:
366 user_id (int): The ID of the user.
367 request (Request): The request object.
368 exception (Exception): The exception that occurred.
369
370 Raises:
371 HTTPException: An HTTPException with a status code of 500 (Internal Server Error).
372 """
373 await self.log_error(user_id, request, exception)
374 raise HTTPException(status_code=500, detail="Internal Server Error")
375
376 async def fetch_all_logs(self):
377 """
378 Fetches all log entries asynchronously.
379
380 Returns:
381 A list of log entries.
382 """
383 result = await self.session.execute(select(LogEntry))
384 logs = result.scalars().all()
385 return logs
386
387
388 ################## ! RETRIEVE LOGS ROUTES ! ##################
389 logs_router = APIRouter()
390
391
392 @logs_router.get(
393 "",
394 response_model=LogsResponse,
395 description="Fetch all logs",
396 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
397 )
398 async def get_logs(session: AsyncSession = Depends(get_db)) -> LogsResponse:
399 """
400 Fetch all logs from the database.
401
402 This endpoint retrieves all the logs stored in the database and returns them
403 along with a success status and message.
404
405 Returns:
406 LogsResponse: A Pydantic model containing a list of logs and additional metadata.
407
408 Raises:
409 HTTPException: An exception with a 404 status code is raised if no logs are found.
410 """
411 auth_handler_instance = AuthHandler() # Initialize your AuthHandler
412 logger_instance = Logger(session, auth_handler_instance)
413
414 logs = await logger_instance.fetch_all_logs() # Assuming fetch_all_logs is an async function
415 if logs:
416 return LogsResponse(
417 logs=logs,
418 success=True,
419 message="Logs fetched successfully",
420 )
421 else:
422 raise HTTPException(status_code=404, detail="No logs found")
423
424
425 @logs_router.get(
426 "/{user_id}",
427 response_model=LogsResponse,
428 description="Fetch logs by user ID",
429 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
430 )
431 async def get_logs_by_user_id(
432 user_id: int,
433 session: AsyncSession = Depends(get_db),
434 ) -> LogsResponse:
435 """
436 Fetch all logs from the database where the user_id matches the provided user_id.
437
438 This endpoint retrieves all the logs stored in the database where the user_id matches the provided user_id
439 and returns them along with a success status and message.
440
441 Args:
442 user_id (int): The user_id to filter logs by.
443
444 Returns:
445 LogsResponse: A Pydantic model containing a list of logs and additional metadata.
446
447 Raises:
448 HTTPException: An exception with a 404 status code is raised if no logs are found.
449 """
450 result = await session.execute(select(LogEntry).filter(LogEntry.user_id == user_id))
451 logs = result.scalars().all()
452
453 if not logs:
454 raise HTTPException(
455 status_code=404,
456 detail=f"No logs found for user ID: {user_id}",
457 )
458
459 return LogsResponse(logs=logs, success=True, message="Logs fetched successfully")
460
461
462 @logs_router.post(
463 "/timerange",
464 response_model=LogsResponse,
465 description="Fetch logs by time range",
466 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
467 )
468 async def get_logs_by_time_range(
469 time_range: TimeRangeModel,
470 session: AsyncSession = Depends(get_db),
471 ) -> LogsResponse:
472 """
473 Fetch all logs from the database where the timestamp is within the provided time range.
474
475 This endpoint retrieves all the logs stored in the database where the timestamp is within the provided time range
476 and returns them along with a success status and message.
477
478 Args:
479 time_range (TimeRangeModel): The time range to filter logs by.
480
481 Returns:
482 LogsResponse: A Pydantic model containing a list of logs and additional metadata.
483
484 Raises:
485 HTTPException: An exception with a 404 status code is raised if no logs are found.
486 """
487 result = await session.execute(select(LogEntry))
488 logs = result.scalars().all()
489
490 if logs:
491 logs = [log for log in logs if log.timestamp >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))]
492 if logs != []:
493 return LogsResponse(
494 logs=logs,
495 success=True,
496 message="Logs fetched successfully",
497 )
498 else:
499 raise HTTPException(
500 status_code=404,
501 detail=f"No logs found for time range: {time_range.time_range}".format(
502 time_range=time_range.time_range,
503 ),
504 )
505 else:
506 raise HTTPException(status_code=404, detail="No logs found")
507
508
509 @logs_router.post(
510 "/{event_type}",
511 response_model=LogsResponse,
512 description="Fetch logs by event type",
513 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
514 )
515 async def get_logs_by_event_type(
516 event_type: EventType,
517 session: AsyncSession = Depends(get_db),
518 ) -> LogsResponse: # Update this line to use the new model
519 """
520 Fetch all logs from the database where the event_type matches the provided event_type.
521
522 This endpoint retrieves all the logs stored in the database where the event_type matches the provided event_type
523 and returns them along with a success status and message.
524
525 Args:
526 event_type (EventType): The event_type to filter logs by.
527
528 Returns:
529 LogsResponse: A Pydantic model containing a list of logs and additional metadata.
530
531 Raises:
532 HTTPException: An exception with a 404 status code is raised if no logs are found.
533 """
534 result = await session.execute(
535 select(LogEntry).filter(LogEntry.event_type == event_type),
536 )
537 logs = result.scalars().all()
538
539 if not logs:
540 raise HTTPException(
541 status_code=404,
542 detail=f"No logs found for event type: {event_type}",
543 )
544
545 return LogsResponse(logs=logs, success=True, message="Logs fetched successfully")
546
547
548 @logs_router.delete(
549 "",
550 response_model=LogsResponse,
551 description="Purge all logs",
552 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
553 )
554 async def purge_logs(
555 session: AsyncSession = Depends(get_db),
556 ) -> LogsResponse: # Update this line to use the new model
557 """
558 Purge all logs from the database.
559
560 This endpoint purges all the logs stored in the database and returns a success status and message.
561
562 Returns:
563 LogsResponse: A Pydantic model containing a list of logs and additional metadata.
564
565 Raises:
566 HTTPException: An exception with a 404 status code is raised if no logs are found.
567 """
568 result = await session.execute(select(LogEntry))
569 logs = result.scalars().all()
570
571 if logs:
572 for log in logs:
573 await session.delete(log)
574 await session.commit()
575 return LogsResponse(logs=[], success=True, message="Logs purged successfully")
576 else:
577 raise HTTPException(status_code=404, detail="No logs found")
578
579
580 @logs_router.delete(
581 "/timerange",
582 response_model=LogsResponse,
583 description="Purge logs by time range",
584 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
585 )
586 async def purge_logs_by_time_range(
587 time_range: TimeRangeModel,
588 session: AsyncSession = Depends(get_db),
589 ) -> LogsResponse:
590 """
591 Purge all logs from the database where the timestamp is within the provided time range.
592
593 This endpoint purges all the logs stored in the database where the timestamp is within the provided time range
594 and returns a success status and message.
595
596 Args:
597 time_range (TimeRangeModel): The time range to filter logs by.
598
599 Returns:
600 LogsResponse: A Pydantic model containing a list of logs and additional metadata.
601
602 Raises:
603 HTTPException: An exception with a 404 status code is raised if no logs are found.
604 """
605 result = await session.execute(select(LogEntry))
606 logs = result.scalars().all()
607
608 if logs:
609 logs = [log for log in logs if log.timestamp >= datetime.now() - timedelta(days=int(time_range.time_range[:-1]))]
610 if logs != []:
611 for log in logs:
612 await session.delete(log)
613 await session.commit()
614 return LogsResponse(
615 logs=[],
616 success=True,
617 message="Logs purged successfully",
618 )
619 else:
620 raise HTTPException(
621 status_code=404,
622 detail=f"No logs found for time range: {time_range.time_range}".format(
623 time_range=time_range.time_range,
624 ),
625 )
626 else:
627 raise HTTPException(status_code=404, detail="No logs found")
628
629
630 ################## ! ALLOWED FILES ! ##################
631 def allowed_file(filename):
632 """
633 Check if the given filename has an allowed extension.
634
635 Args:
636 filename (str): The name of the file to check.
637
638 Returns:
639 bool: True if the file has an allowed extension, False otherwise.
640 """
641 ALLOWED_EXTENSIONS = {"yaml", "txt"}
642 return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS
643
644
645 ################## ! DATABASE UTILS ! ##################
646 async def get_connector_attribute(
647 column_name: str,
648 connector_id: Optional[int] = None,
649 connector_name: Optional[str] = None,
650 session: AsyncSession = Depends(get_session),
651 ) -> Optional[Any]:
652 """
653 Retrieve the value of a specific column from a connector.
654
655 Args:
656 connector_id (Optional[int]): The ID of the connector.
657 connector_name (Optional[str]): The name of the connector.
658 column_name (str): The name of the column to retrieve.
659 session (AsyncSession, optional): The database session. Defaults to Depends(get_session).
660
661 Returns:
662 Optional[Any]: The value of the column, or None if the connector or column does not exist.
663 """
664 if not connector_id and not connector_name:
665 raise ValueError("Either connector_id or connector_name must be provided.")
666
667 query = select(Connectors)
668 if connector_id:
669 query = query.filter(Connectors.id == connector_id)
670 if connector_name:
671 query = query.filter(Connectors.connector_name == connector_name)
672
673 result = await session.execute(query)
674 connector = result.scalars().first()
675
676 if connector:
677 return getattr(connector, column_name, None)
678 return None
679
680
681 async def get_customer_meta_attribute(
682 customer_code: str,
683 column_name: str,
684 session: AsyncSession = Depends(get_session),
685 ) -> Optional[Any]:
686 """
687 Retrieve the value of a specific column from a customer.
688
689 Args:
690 customer_code (str): The code of the customer.
691 column_name (str): The name of the column to retrieve.
692 session (AsyncSession, optional): The database session. Defaults to Depends(get_session).
693
694 Returns:
695 Optional[Any]: The value of the column, or None if the customer or column does not exist.
696 """
697 result = await session.execute(
698 select(CustomersMeta).filter(CustomersMeta.customer_code == customer_code),
699 )
700 customer = result.scalars().first()
701
702 if customer:
703 return getattr(customer, column_name, None)
704 return None
705
706
707 async def get_customer_default_settings_attribute(
708 column_name: str,
709 session: AsyncSession = Depends(get_session),
710 ) -> Optional[Any]:
711 """
712 Retrieve the value of a specific column from a customer's default settings.
713
714 Args:
715 customer_code (str): The code of the customer.
716 column_name (str): The name of the column to retrieve.
717 session (AsyncSession, optional): The database session. Defaults to Depends(get_session).
718
719 Returns:
720 Optional[Any]: The value of the column, or None if the customer or column does not exist.
721 """
722 result = await session.execute(select(CustomerProvisioningDefaultSettings))
723 settings = result.scalars().first()
724
725 if settings:
726 return getattr(settings, column_name, None)
727 return None
728
729
730 async def get_customer_alert_settings(
731 customer_code: str,
732 session: AsyncSession,
733 ) -> Optional[AlertCreationSettings]:
734 """
735 Retrieve the alert creation settings for a specific customer.
736
737 Args:
738 customer_code (str): The code of the customer.
739 session (AsyncSession): The database session.
740
741 Returns:
742 Optional[AlertCreationSettings]: The alert creation settings for the customer, or None if not found.
743 """
744 result = await session.execute(
745 select(AlertCreationSettings).filter(
746 AlertCreationSettings.customer_code == customer_code,
747 ),
748 )
749 settings = result.scalars().first()
750
751 if not settings:
752 result = await session.execute(
753 select(AlertCreationSettings).filter(
754 AlertCreationSettings.office365_organization_id == customer_code,
755 ),
756 )
757 settings = result.scalars().first()
758
759 return settings
760
761
762 async def get_customer_alert_settings_office365(
763 office365_organization_id: str,
764 session: AsyncSession,
765 ) -> Optional[AlertCreationSettings]:
766 """
767 Retrieve the alert creation settings for a specific customer.
768
769 Args:
770 office365_organization_id (str): The Office365 Organization ID of the customer.
771 session (AsyncSession): The database session.
772
773 Returns:
774 Optional[AlertCreationSettings]: The alert creation settings for the customer, or None if not found.
775 """
776 result = await session.execute(
777 select(AlertCreationSettings).filter(
778 AlertCreationSettings.office365_organization_id == office365_organization_id,
779 ),
780 )
781 settings = result.scalars().first()
782
783 if settings:
784 return settings
785 return None
786
787
788 async def get_customer_alert_event_configs(
789 customer_code: str,
790 session: AsyncSession = Depends(get_session),
791 ) -> Optional[List[List[AlertCreationEventConfig]]]:
792 """
793 Retrieves the alert event configurations for a specific customer.
794
795 Args:
796 customer_code (str): The code of the customer.
797 session (AsyncSession, optional): The database session. Defaults to Depends(get_session).
798
799 Returns:
800 Optional[List[List[AlertCreationEventConfig]]]: A list of event configurations, or None if no settings found.
801 """
802 result = await session.execute(
803 select(AlertCreationSettings)
804 .options(
805 joinedload(AlertCreationSettings.event_orders).joinedload(
806 EventOrder.event_configs,
807 ),
808 )
809 .where(AlertCreationSettings.customer_code == customer_code),
810 )
811 settings = result.scalars().first()
812
813 if settings:
814 return [order.event_configs for order in settings.event_orders]
815 return None
816
817
818 ################## ! Wazuh Worker Provisioning App ! ##################
819 ################## ! https://github.com/socfortress/Customer-Provisioning-Worker ! ##################
820 # async def verify_wazuh_worker_provisioning_healtcheck(
821 # attributes: Dict[str, Any],
822 # ) -> Dict[str, Any]:
823 # """
824 # Verifies the connection to Wazuh Worker Provisioning service.
825
826 # Returns:
827 # dict: A dictionary containing 'connectionSuccessful' status.
828 # """
829 # logger.info(
830 # f"Verifying the wazuh-worker provisioning connection to {attributes['connector_url']}",
831 # )
832
833 # try:
834 # wazuh_worker_provisioning_healthcheck = requests.get(
835 # f"{attributes['connector_url']}/provision_worker/healthcheck",
836 # verify=False,
837 # )
838
839 # if wazuh_worker_provisioning_healthcheck.status_code == 200:
840 # return {
841 # "connectionSuccessful": True,
842 # "message": "Wazuh Worker Provisioning healthcheck successful",
843 # }
844 # else:
845 # logger.error(
846 # f"Connection to {attributes['connector_url']} failed with error: {wazuh_worker_provisioning_healthcheck.text}",
847 # )
848
849 # return {
850 # "connectionSuccessful": False,
851 # "message": f"Connection to {attributes['connector_url']} failed",
852 # }
853 # except Exception as e:
854 # logger.error(
855 # f"Connection to {attributes['connector_url']} failed with error: {e}",
856 # )
857
858 # return {
859 # "connectionSuccessful": False,
860 # "message": f"Connection to {attributes['connector_url']} failed with error.",
861 # }
862
863
864 async def verify_wazuh_worker_provisioning_healtcheck(
865 attributes: Dict[str, Any],
866 ) -> Dict[str, Any]:
867 """
868 Verifies the connection to Wazuh Worker Provisioning service.
869 Supports multiple hosts separated by commas.
870
871 Returns:
872 dict: A dictionary containing 'connectionSuccessful' status and details.
873 """
874 connector_url = attributes["connector_url"]
875 logger.info(f"Verifying the wazuh-worker provisioning connection to {connector_url}")
876
877 # Parse multiple hosts if comma-separated
878 hosts = [host.strip() for host in connector_url.split(",") if host.strip()]
879
880 if not hosts:
881 return {
882 "connectionSuccessful": False,
883 "message": "No valid hosts found in connector_url",
884 }
885
886 successful_hosts = []
887 failed_hosts = []
888 connection_details = []
889
890 # Test each host
891 for host in hosts:
892 try:
893 logger.info(f"Testing connection to host: {host}")
894
895 wazuh_worker_provisioning_healthcheck = requests.get(
896 f"{host}/provision_worker/healthcheck",
897 verify=False,
898 timeout=10, # Add timeout to prevent hanging
899 )
900
901 if wazuh_worker_provisioning_healthcheck.status_code == 200:
902 successful_hosts.append(host)
903 connection_details.append({"host": host, "status": "success", "status_code": 200, "message": "Connection successful"})
904 logger.info(f"Connection to {host} successful")
905 else:
906 failed_hosts.append(host)
907 connection_details.append(
908 {
909 "host": host,
910 "status": "failed",
911 "status_code": wazuh_worker_provisioning_healthcheck.status_code,
912 "message": f"HTTP {wazuh_worker_provisioning_healthcheck.status_code}: {wazuh_worker_provisioning_healthcheck.text}",
913 },
914 )
915 logger.error(f"Connection to {host} failed with status {wazuh_worker_provisioning_healthcheck.status_code}")
916
917 except Exception as e:
918 failed_hosts.append(host)
919 connection_details.append({"host": host, "status": "error", "status_code": None, "message": f"Connection error: {str(e)}"})
920 logger.error(f"Connection to {host} failed with error: {e}")
921
922 # Determine overall success
923 overall_success = len(successful_hosts) > 0
924
925 if overall_success:
926 if len(failed_hosts) == 0:
927 message = f"All {len(successful_hosts)} hosts connected successfully"
928 else:
929 message = f"{len(successful_hosts)} of {len(hosts)} hosts connected successfully"
930 else:
931 message = f"All {len(hosts)} hosts failed to connect"
932
933 return {
934 "connectionSuccessful": overall_success,
935 "message": message,
936 "total_hosts": len(hosts),
937 "successful_hosts": len(successful_hosts),
938 "failed_hosts": len(failed_hosts),
939 "successful_host_list": successful_hosts,
940 "failed_host_list": failed_hosts,
941 "connection_details": connection_details,
942 }
943
944
945 async def verify_wazuh_worker_provisioning_connection(connector_name: str) -> str:
946 """
947 Returns the status of the connection to Wazuh Worker Provisioning service.
948 """
949 async with get_db_session() as session: # This will correctly enter the context manager
950 attributes = await get_connector_info_from_db(connector_name, session)
951 if attributes is None:
952 logger.error("No Wazuh Worker Provisioning connector found in the database")
953 return None
954 return await verify_wazuh_worker_provisioning_healtcheck(attributes)
955
956
957 ################## ! HAPROXY Provisioning App ! ##################
958 ################## ! https://github.com/socfortress/Customer-Provisioning-Worker ! ##################
959 async def verify_haproxy_provisioning_healtcheck(
960 attributes: Dict[str, Any],
961 ) -> Dict[str, Any]:
962 """
963 Verifies the connection to HAPROXY Provisioning service.
964
965 Returns:
966 dict: A dictionary containing 'connectionSuccessful' status.
967 """
968 logger.info(
969 f"Verifying the HAPROXY provisioning connection to {attributes['connector_url']}",
970 )
971
972 try:
973 wazuh_worker_provisioning_healthcheck = requests.get(
974 f"{attributes['connector_url']}/provision_worker/healthcheck",
975 verify=False,
976 )
977
978 if wazuh_worker_provisioning_healthcheck.status_code == 200:
979 return {
980 "connectionSuccessful": True,
981 "message": "Wazuh Worker Provisioning healthcheck successful",
982 }
983 else:
984 logger.error(
985 f"Connection to {attributes['connector_url']} failed with error: {wazuh_worker_provisioning_healthcheck.text}",
986 )
987
988 return {
989 "connectionSuccessful": False,
990 "message": f"Connection to {attributes['connector_url']} failed",
991 }
992 except Exception as e:
993 logger.error(
994 f"Connection to {attributes['connector_url']} failed with error: {e}",
995 )
996
997 return {
998 "connectionSuccessful": False,
999 "message": f"Connection to {attributes['connector_url']} failed with error.",
1000 }
1001
1002
1003 async def verify_haproxy_provisioning_connection(connector_name: str) -> str:
1004 """
1005 Returns the status of the connection to HAPROXY Provisioning service.
1006 """
1007 async with get_db_session() as session: # This will correctly enter the context manager
1008 attributes = await get_connector_info_from_db(connector_name, session)
1009 if attributes is None:
1010 logger.error("No HAPROXY Provisioning connector found in the database")
1011 return None
1012 return await verify_haproxy_provisioning_healtcheck(attributes)
1013
1014
1015 ################## ! Alert Creation Provisioning App ! ##################
1016 ################## ! https://github.com/socfortress/Customer-Provisioning-Alert ! ##################
1017 async def verify_alert_creation_provisioning_healtcheck(
1018 attributes: Dict[str, Any],
1019 ) -> Dict[str, Any]:
1020 """
1021 Verifies the connection to Alert Creation Provisioning service.
1022
1023 Returns:
1024 dict: A dictionary containing 'connectionSuccessful' status.
1025 """
1026 logger.info(
1027 f"Verifying the Alert Creation provisioning connection to {attributes['connector_url']}",
1028 )
1029
1030 try:
1031 wazuh_worker_provisioning_healthcheck = requests.get(
1032 f"{attributes['connector_url']}/provision_alert/healthcheck",
1033 verify=False,
1034 )
1035
1036 if wazuh_worker_provisioning_healthcheck.status_code == 200:
1037 return {
1038 "connectionSuccessful": True,
1039 "message": "Alert Creation Provisioning healthcheck successful",
1040 }
1041 else:
1042 logger.error(
1043 f"Connection to {attributes['connector_url']} failed with error: {wazuh_worker_provisioning_healthcheck.text}",
1044 )
1045
1046 return {
1047 "connectionSuccessful": False,
1048 "message": f"Connection to {attributes['connector_url']} failed",
1049 }
1050 except Exception as e:
1051 logger.error(
1052 f"Connection to {attributes['connector_url']} failed with error: {e}",
1053 )
1054
1055 return {
1056 "connectionSuccessful": False,
1057 "message": f"Connection to {attributes['connector_url']} failed with error.",
1058 }
1059
1060
1061 async def verify_alert_creation_provisioning_connection(connector_name: str) -> str:
1062 """
1063 Returns the status of the connection to Alert Creation Provisioning service.
1064 """
1065 async with get_db_session() as session: # This will correctly enter the context manager
1066 attributes = await get_connector_info_from_db(connector_name, session)
1067 if attributes is None:
1068 logger.error("No Alert Creation Provisioning connector found in the database")
1069 return None
1070 return await verify_alert_creation_provisioning_healtcheck(attributes)
1071
1072
1073 ################## ! VIRUSTOTAL ! ##################
1074
1075
1076 async def verify_virustotal_healtcheck(
1077 attributes: Dict[str, Any],
1078 ) -> Dict[str, Any]:
1079 """
1080 Verifies the connection to VirusTotal service.
1081
1082 Returns:
1083 dict: A dictionary containing 'connectionSuccessful' status.
1084 """
1085 logger.info(
1086 f"Verifying the VirusTotal connection to {attributes['connector_url']} and API key {attributes['connector_api_key']}",
1087 )
1088
1089 try:
1090 virustotal_healthcheck = requests.get(
1091 f"{attributes['connector_url']}/files/99017f6eebbac24f351415dd410d522d",
1092 verify=False,
1093 headers={"x-apikey": attributes["connector_api_key"]},
1094 )
1095
1096 if virustotal_healthcheck.status_code == 200:
1097 return {
1098 "connectionSuccessful": True,
1099 "message": "VirusTotal healthcheck successful",
1100 }
1101 else:
1102 logger.error(
1103 f"Connection to {attributes['connector_url']} failed with error: {virustotal_healthcheck.text}",
1104 )
1105
1106 return {
1107 "connectionSuccessful": False,
1108 "message": f"Connection to {attributes['connector_url']} failed",
1109 }
1110 except Exception as e:
1111 logger.error(
1112 f"Connection to {attributes['connector_url']} failed with error: {e}",
1113 )
1114
1115 return {
1116 "connectionSuccessful": False,
1117 "message": f"Connection to {attributes['connector_url']} failed with error.",
1118 }
1119
1120
1121 async def verify_virustotal_connection(connector_name: str) -> str:
1122 """
1123 Returns the status of the connection to VirusTotal service.
1124 """
1125 async with get_db_session() as session: # This will correctly enter the context manager
1126 attributes = await get_connector_info_from_db(connector_name, session)
1127 if attributes is None:
1128 logger.error("No VirusTotal connector found in the database")
1129 return None
1130 return await verify_virustotal_healtcheck(attributes)