| 1 | import json |
| 2 | from typing import List |
| 3 | |
| 4 | from fastapi import HTTPException |
| 5 | from loguru import logger |
| 6 | from sqlalchemy.ext.asyncio import AsyncSession |
| 7 | from sqlalchemy.future import select |
| 8 | from sqlalchemy.orm import selectinload |
| 9 | |
| 10 | from app.connectors.sublime.models.alerts import FlaggedRule |
| 11 | from app.connectors.sublime.models.alerts import Mailbox |
| 12 | from app.connectors.sublime.models.alerts import Recipient |
| 13 | from app.connectors.sublime.models.alerts import Sender |
| 14 | from app.connectors.sublime.models.alerts import SublimeAlerts |
| 15 | from app.connectors.sublime.models.alerts import TriggeredAction |
| 16 | from app.connectors.sublime.schema.alerts import AlertRequestBody |
| 17 | from app.connectors.sublime.schema.alerts import AlertResponseBody |
| 18 | from app.connectors.sublime.schema.alerts import SublimeAlertsResponse |
| 19 | from app.connectors.sublime.schema.alerts import SublimeAlertsSchema |
| 20 | from app.connectors.sublime.utils.universal import send_get_request |
| 21 | |
| 22 | |
| 23 | def create_sublime_alert(alert_request_body: AlertRequestBody) -> SublimeAlerts: |
| 24 | """ |
| 25 | Creates a SublimeAlerts object based on the provided AlertRequestBody. |
| 26 | |
| 27 | Args: |
| 28 | alert_request_body (AlertRequestBody): The request body containing the alert information. |
| 29 | |
| 30 | Returns: |
| 31 | SublimeAlerts: The created SublimeAlerts object. |
| 32 | """ |
| 33 | return SublimeAlerts( |
| 34 | api_version=alert_request_body.api_version, |
| 35 | created_at=alert_request_body.created_at, |
| 36 | event_id=alert_request_body.id, |
| 37 | type=alert_request_body.type, |
| 38 | message_id=alert_request_body.data.message.id, |
| 39 | canonical_id=alert_request_body.data.message.canonical_id, |
| 40 | external_id=alert_request_body.data.message.external_id, |
| 41 | message_source_id=alert_request_body.data.message.message_source_id, |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | def create_flagged_rules( |
| 46 | alert_request_body: AlertRequestBody, |
| 47 | sublime_alert_id: int, |
| 48 | ) -> List[FlaggedRule]: |
| 49 | """ |
| 50 | Create a list of flagged rules based on the given alert request body and sublime alert ID. |
| 51 | |
| 52 | Args: |
| 53 | alert_request_body (AlertRequestBody): The request body containing the alert data. |
| 54 | sublime_alert_id (int): The ID of the sublime alert. |
| 55 | |
| 56 | Returns: |
| 57 | List[FlaggedRule]: A list of flagged rules. |
| 58 | |
| 59 | """ |
| 60 | flagged_rules = [] |
| 61 | for rule in alert_request_body.data.flagged_rules: |
| 62 | tags_str = json.dumps(rule.tags) |
| 63 | flagged_rules.append( |
| 64 | FlaggedRule( |
| 65 | rule_id=rule.id, |
| 66 | name=rule.name, |
| 67 | severity=rule.severity, |
| 68 | tags=tags_str, |
| 69 | sublime_alert_id=sublime_alert_id, |
| 70 | ), |
| 71 | ) |
| 72 | return flagged_rules |
| 73 | |
| 74 | |
| 75 | def create_mailbox( |
| 76 | alert_request_body: AlertRequestBody, |
| 77 | sublime_alert_id: int, |
| 78 | ) -> Mailbox: |
| 79 | """ |
| 80 | Create a mailbox object based on the provided alert request body and sublime alert ID. |
| 81 | |
| 82 | Args: |
| 83 | alert_request_body (AlertRequestBody): The request body containing the alert data. |
| 84 | sublime_alert_id (int): The ID of the sublime alert. |
| 85 | |
| 86 | Returns: |
| 87 | Mailbox: The created mailbox object. |
| 88 | """ |
| 89 | return Mailbox( |
| 90 | external_id=alert_request_body.data.message.mailbox.external_id, |
| 91 | mailbox_id=alert_request_body.data.message.mailbox.id, |
| 92 | sublime_alert_id=sublime_alert_id, |
| 93 | ) |
| 94 | |
| 95 | |
| 96 | def create_triggered_actions( |
| 97 | alert_request_body: AlertRequestBody, |
| 98 | sublime_alert_id: int, |
| 99 | ) -> List[TriggeredAction]: |
| 100 | """ |
| 101 | Create a list of TriggeredAction objects based on the provided alert request body and sublime alert ID. |
| 102 | |
| 103 | Args: |
| 104 | alert_request_body (AlertRequestBody): The request body containing the data for the alert. |
| 105 | sublime_alert_id (int): The ID of the sublime alert. |
| 106 | |
| 107 | Returns: |
| 108 | List[TriggeredAction]: A list of TriggeredAction objects. |
| 109 | """ |
| 110 | triggered_actions = [] |
| 111 | for action in alert_request_body.data.triggered_actions: |
| 112 | triggered_actions.append( |
| 113 | TriggeredAction( |
| 114 | action_id=action.id, |
| 115 | name=action.name, |
| 116 | type=action.type, |
| 117 | sublime_alert_id=sublime_alert_id, |
| 118 | ), |
| 119 | ) |
| 120 | return triggered_actions |
| 121 | |
| 122 | |
| 123 | async def store_sublime_alert( |
| 124 | session: AsyncSession, |
| 125 | alert_request_body: AlertRequestBody, |
| 126 | ) -> AlertResponseBody: |
| 127 | """ |
| 128 | Stores a Sublime alert in the database. |
| 129 | |
| 130 | Args: |
| 131 | session (AsyncSession): The database session. |
| 132 | alert_request_body (AlertRequestBody): The request body containing the alert data. |
| 133 | |
| 134 | Returns: |
| 135 | AlertResponseBody: The response body indicating the success or failure of the operation. |
| 136 | """ |
| 137 | try: |
| 138 | sublime_alert = create_sublime_alert(alert_request_body) |
| 139 | session.add(sublime_alert) |
| 140 | await session.flush() # Flush to obtain the ID of the new alert |
| 141 | |
| 142 | flagged_rules = create_flagged_rules(alert_request_body, sublime_alert.id) |
| 143 | mailbox = create_mailbox(alert_request_body, sublime_alert.id) |
| 144 | triggered_actions = create_triggered_actions( |
| 145 | alert_request_body, |
| 146 | sublime_alert.id, |
| 147 | ) |
| 148 | sender = await create_sender(alert_request_body, sublime_alert.id) |
| 149 | recipient = await create_recipient(alert_request_body, sublime_alert.id) |
| 150 | |
| 151 | session.add_all(flagged_rules) |
| 152 | session.add(mailbox) |
| 153 | session.add_all(triggered_actions) |
| 154 | session.add(sender) |
| 155 | session.add(recipient) |
| 156 | |
| 157 | logger.info(f"Preparing to store: {sublime_alert}") |
| 158 | await session.commit() # Commit the changes asynchronously |
| 159 | logger.info(f"Alert {alert_request_body.id} stored in the database") |
| 160 | |
| 161 | return AlertResponseBody( |
| 162 | success=True, |
| 163 | message=f"Alert {alert_request_body.id} stored in the database", |
| 164 | ) |
| 165 | except Exception as e: |
| 166 | # Rollback in case of error |
| 167 | await session.rollback() |
| 168 | logger.error( |
| 169 | f"Failed to store alert {alert_request_body.id} in the database: {e}", |
| 170 | ) |
| 171 | raise HTTPException( |
| 172 | status_code=500, |
| 173 | detail=f"Failed to store alert {alert_request_body.id} in the database: {e}", |
| 174 | ) |
| 175 | |
| 176 | |
| 177 | async def create_sender( |
| 178 | alert_request_body: AlertRequestBody, |
| 179 | sublime_alert_id: int, |
| 180 | ) -> Sender: |
| 181 | """ |
| 182 | Create a Sender object based on the given alert request body and sublime alert ID. |
| 183 | |
| 184 | Args: |
| 185 | alert_request_body (AlertRequestBody): The request body containing the alert data. |
| 186 | sublime_alert_id (int): The ID of the sublime alert. |
| 187 | |
| 188 | Returns: |
| 189 | Sender: The created Sender object. |
| 190 | """ |
| 191 | return Sender( |
| 192 | email=await collect_sender(alert_request_body.data.message.id), |
| 193 | name="n/a", |
| 194 | sublime_alert_id=sublime_alert_id, |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | async def create_recipient( |
| 199 | alert_request_body: AlertRequestBody, |
| 200 | sublime_alert_id: int, |
| 201 | ) -> Recipient: |
| 202 | """ |
| 203 | Create a recipient for the given alert request body and sublime alert ID. |
| 204 | |
| 205 | Args: |
| 206 | alert_request_body (AlertRequestBody): The alert request body. |
| 207 | sublime_alert_id (int): The sublime alert ID. |
| 208 | |
| 209 | Returns: |
| 210 | Recipient: The created recipient. |
| 211 | """ |
| 212 | return Recipient( |
| 213 | email=await collect_recipient(alert_request_body.data.message.id), |
| 214 | name="n/a", |
| 215 | sublime_alert_id=sublime_alert_id, |
| 216 | ) |
| 217 | |
| 218 | |
| 219 | async def collect_sender(message_id: str) -> Sender: |
| 220 | """ |
| 221 | Get a single Sublime Alert from the database |
| 222 | |
| 223 | Args: |
| 224 | message_id (str): The ID of the message to retrieve the sender for. |
| 225 | |
| 226 | Returns: |
| 227 | str: The email address of the sender. |
| 228 | |
| 229 | Raises: |
| 230 | HTTPException: If there is an error retrieving the Sublime Alert. |
| 231 | """ |
| 232 | logger.info(f"Getting Sublime Alert with message_id {message_id}") |
| 233 | message_details = await send_get_request(f"/v0/messages/{message_id}") |
| 234 | if not message_details["success"]: |
| 235 | logger.error( |
| 236 | f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}", |
| 237 | ) |
| 238 | raise HTTPException( |
| 239 | status_code=500, |
| 240 | detail=f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}", |
| 241 | ) |
| 242 | logger.info(f"Successfully retrieved Sublime Alert with message_id {message_id}") |
| 243 | return message_details["data"]["sender"]["email"] |
| 244 | |
| 245 | |
| 246 | async def collect_recipient(message_id: str) -> Recipient: |
| 247 | """ |
| 248 | Get a single Sublime Alert recipient from the database. |
| 249 | |
| 250 | Args: |
| 251 | message_id (str): The ID of the Sublime Alert message. |
| 252 | |
| 253 | Returns: |
| 254 | str: The email address of the recipient. |
| 255 | |
| 256 | Raises: |
| 257 | HTTPException: If there is an error retrieving the Sublime Alert recipient. |
| 258 | """ |
| 259 | logger.info(f"Getting Sublime Alert with message_id {message_id}") |
| 260 | message_details = await send_get_request(f"/v0/messages/{message_id}") |
| 261 | if not message_details["success"]: |
| 262 | logger.error( |
| 263 | f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}", |
| 264 | ) |
| 265 | raise HTTPException( |
| 266 | status_code=500, |
| 267 | detail=f"Failed to get Sublime Alert with message_id {message_id}: {message_details['message']}", |
| 268 | ) |
| 269 | logger.info(f"Successfully retrieved Sublime Alert with message_id {message_id}") |
| 270 | return message_details["data"]["recipients"][0]["email"] |
| 271 | |
| 272 | |
| 273 | async def collect_alerts(session: AsyncSession) -> List[SublimeAlertsResponse]: |
| 274 | """ |
| 275 | Get all Sublime Alerts from the database asynchronously. |
| 276 | |
| 277 | Args: |
| 278 | session (AsyncSession): The database session. |
| 279 | |
| 280 | Returns: |
| 281 | List[SublimeAlertsResponse]: A list of SublimeAlertsResponse objects. |
| 282 | """ |
| 283 | logger.info("Getting all Sublime Alerts") |
| 284 | try: |
| 285 | # Asynchronous query to load all alerts and their related objects |
| 286 | stmt = select(SublimeAlerts).options( |
| 287 | selectinload(SublimeAlerts.flagged_rules), |
| 288 | selectinload(SublimeAlerts.mailbox), |
| 289 | selectinload(SublimeAlerts.triggered_actions), |
| 290 | selectinload(SublimeAlerts.sender), |
| 291 | selectinload(SublimeAlerts.recipients), |
| 292 | ) |
| 293 | result = await session.execute(stmt) |
| 294 | alerts = result.scalars().all() |
| 295 | |
| 296 | logger.info("Successfully retrieved all Sublime Alerts") |
| 297 | return SublimeAlertsResponse( |
| 298 | success=True, |
| 299 | message="Successfully retrieved all Sublime Alerts", |
| 300 | sublime_alerts=[SublimeAlertsSchema.from_orm(alert) for alert in alerts], |
| 301 | ) |
| 302 | except Exception as e: |
| 303 | logger.error(f"Failed to get all Sublime Alerts with error: {e}") |
| 304 | raise HTTPException( |
| 305 | status_code=500, |
| 306 | detail=f"Failed to get all Sublime Alerts with error: {e}", |
| 307 | ) |