| 1 | import base64 |
| 2 | import datetime |
| 3 | import hashlib |
| 4 | import hmac |
| 5 | import io |
| 6 | import json |
| 7 | import os |
| 8 | import shutil |
| 9 | import time |
| 10 | import uuid |
| 11 | from zipfile import ZipFile |
| 12 | |
| 13 | import aiofiles |
| 14 | import requests |
| 15 | from fastapi import HTTPException |
| 16 | from loguru import logger |
| 17 | |
| 18 | from app.integrations.mimecast.schema.mimecast import DataItem |
| 19 | from app.integrations.mimecast.schema.mimecast import MimecastAPIEndpointResponse |
| 20 | from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys |
| 21 | from app.integrations.mimecast.schema.mimecast import MimecastRequest |
| 22 | from app.integrations.mimecast.schema.mimecast import MimecastResponse |
| 23 | from app.integrations.mimecast.schema.mimecast import MimecastTTPURLSRequest |
| 24 | from app.integrations.mimecast.schema.mimecast import RequestBody |
| 25 | from app.integrations.mimecast.schema.mimecast import TtpURLResponseBody |
| 26 | from app.integrations.utils.collection import send_post_request |
| 27 | from app.integrations.utils.event_shipper import event_shipper |
| 28 | from app.integrations.utils.schema import EventShipperPayload |
| 29 | |
| 30 | |
| 31 | async def get_checkpoint_filename(customer_code: str): |
| 32 | """ |
| 33 | Retrieves the checkpoint filename for the Mimecast integration. |
| 34 | If the checkpoint file does not exist, it will be created asynchronously. |
| 35 | """ |
| 36 | # Relative path from the current script to the checkpoint directory |
| 37 | checkpoint_directory = os.path.join(os.path.dirname(__file__), "..", "checkpoint") |
| 38 | checkpoint_filename = os.path.join( |
| 39 | checkpoint_directory, |
| 40 | f"mimecast_{customer_code}.checkpoint", |
| 41 | ) |
| 42 | |
| 43 | # Normalize the path to remove relative path components |
| 44 | checkpoint_filename = os.path.normpath(checkpoint_filename) |
| 45 | |
| 46 | # Create the checkpoint directory if it does not exist |
| 47 | if not os.path.exists(checkpoint_directory): |
| 48 | os.makedirs(checkpoint_directory) |
| 49 | |
| 50 | # Create the checkpoint file if it does not exist |
| 51 | if not os.path.exists(checkpoint_filename): |
| 52 | async with aiofiles.open(checkpoint_filename, "w") as f: |
| 53 | await f.write("") |
| 54 | |
| 55 | return checkpoint_filename |
| 56 | |
| 57 | |
| 58 | async def get_log_file_path(customer_code: str): |
| 59 | """ |
| 60 | Retrieves the log file path for the Mimecast integration and customer. |
| 61 | |
| 62 | Args: |
| 63 | customer_code (str): The code of the customer. |
| 64 | |
| 65 | Returns: |
| 66 | str: The log file path for the Mimecast integration and customer. |
| 67 | """ |
| 68 | # Relative path from the current script to the log directory |
| 69 | log_directory = os.path.join(os.path.dirname(__file__), "..", "logs") |
| 70 | # Normalize the path to remove relative path components |
| 71 | log_directory = os.path.abspath(log_directory) |
| 72 | |
| 73 | # Create a directory for the customer if it does not exist |
| 74 | customer_log_directory = os.path.join(log_directory, customer_code) |
| 75 | if not os.path.exists(customer_log_directory): |
| 76 | os.makedirs(customer_log_directory) |
| 77 | |
| 78 | # Return the customer directory |
| 79 | return customer_log_directory |
| 80 | |
| 81 | |
| 82 | async def read_file(filename: str): |
| 83 | """ |
| 84 | Reads the contents of the given file. |
| 85 | """ |
| 86 | async with aiofiles.open(filename, "r") as f: |
| 87 | return await f.read() |
| 88 | |
| 89 | |
| 90 | async def get_hdr_date(): |
| 91 | return datetime.datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S UTC") |
| 92 | |
| 93 | |
| 94 | async def get_base_url( |
| 95 | mimecast_auth_keys: MimecastAuthKeys, |
| 96 | ) -> MimecastAPIEndpointResponse: |
| 97 | """ |
| 98 | Retrieves the base URL for the Mimecast integration. |
| 99 | """ |
| 100 | post_body = dict() |
| 101 | post_body["data"] = [{}] |
| 102 | post_body["data"][0]["emailAddress"] = mimecast_auth_keys.EMAIL_ADDRESS |
| 103 | |
| 104 | # Create variables required for request headers |
| 105 | request_id = str(uuid.uuid4()) |
| 106 | request_date = await get_hdr_date() |
| 107 | headers = { |
| 108 | "x-mc-app-id": mimecast_auth_keys.APP_ID, |
| 109 | "x-mc-req-id": request_id, |
| 110 | "x-mc-date": request_date, |
| 111 | } |
| 112 | try: |
| 113 | response = await send_post_request( |
| 114 | endpoint="https://api.mimecast.com/api/login/discover-authentication", |
| 115 | headers=headers, |
| 116 | data=post_body, |
| 117 | ) |
| 118 | if response["success"] is True: |
| 119 | logger.info( |
| 120 | f"Successfully retrieved base URL for Mimecast integration. Response: {response}", |
| 121 | ) |
| 122 | return MimecastAPIEndpointResponse(**response) |
| 123 | else: |
| 124 | logger.error( |
| 125 | f"Unable to retrieve base URL for Mimecast integration. Response: {response}", |
| 126 | ) |
| 127 | raise HTTPException( |
| 128 | status_code=400, |
| 129 | detail="Unable to retrieve base URL for Mimecast integration.", |
| 130 | ) |
| 131 | except Exception as e: |
| 132 | logger.error( |
| 133 | f"Unable to retrieve base URL for Mimecast integration. Exception: {e}", |
| 134 | ) |
| 135 | raise HTTPException( |
| 136 | status_code=400, |
| 137 | detail="Unable to retrieve base URL for Mimecast integration.", |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | async def get_mta_siem_logs( |
| 142 | checkpoint_filename: str, |
| 143 | base_url: str, |
| 144 | auth_keys: MimecastAuthKeys, |
| 145 | ): |
| 146 | """ |
| 147 | Retrieves the MTA SIEM logs from the Mimecast integration. |
| 148 | """ |
| 149 | # Build post body for request |
| 150 | post_body = dict() |
| 151 | post_body["data"] = [{}] |
| 152 | post_body["data"][0]["type"] = "MTA" |
| 153 | post_body["data"][0]["compress"] = True |
| 154 | post_body["data"][0]["token"] = await read_file(checkpoint_filename) |
| 155 | |
| 156 | # Create variables required for request headers |
| 157 | request_id = str(uuid.uuid4()) |
| 158 | request_date = await get_hdr_date() |
| 159 | |
| 160 | unsigned_auth_header = "{date}:{req_id}:{uri}:{app_key}".format( |
| 161 | date=request_date, |
| 162 | req_id=request_id, |
| 163 | uri=auth_keys.URI, |
| 164 | app_key=auth_keys.APP_KEY, |
| 165 | ) |
| 166 | hmac_sha1 = hmac.new( |
| 167 | base64.b64decode(auth_keys.SECRET_KEY), |
| 168 | unsigned_auth_header.encode(), |
| 169 | digestmod=hashlib.sha1, |
| 170 | ).digest() |
| 171 | sig = base64.encodebytes(hmac_sha1).rstrip() |
| 172 | headers = { |
| 173 | "Authorization": "MC " + auth_keys.ACCESS_KEY + ":" + sig.decode(), |
| 174 | "x-mc-app-id": auth_keys.APP_ID, |
| 175 | "x-mc-date": request_date, |
| 176 | "x-mc-req-id": request_id, |
| 177 | "Content-Type": "application/json", |
| 178 | } |
| 179 | try: |
| 180 | response = requests.post( |
| 181 | url=base_url + auth_keys.URI, |
| 182 | headers=headers, |
| 183 | data=json.dumps(post_body), |
| 184 | ) |
| 185 | return response.content, response.headers |
| 186 | except Exception as e: |
| 187 | logger.error( |
| 188 | f"Unable to retrieve MTA SIEM logs from Mimecast integration. Exception: {e}", |
| 189 | ) |
| 190 | raise HTTPException( |
| 191 | status_code=400, |
| 192 | detail="Unable to retrieve MTA SIEM logs from Mimecast integration.", |
| 193 | ) |
| 194 | |
| 195 | |
| 196 | async def process_response(response, checkpoint_filename: str, log_file_path: str): |
| 197 | """ |
| 198 | Processes the response body from the Mimecast integration. |
| 199 | """ |
| 200 | if response != "error": |
| 201 | resp_body = response[0] |
| 202 | resp_headers = response[1] |
| 203 | content_type = resp_headers["Content-Type"] |
| 204 | |
| 205 | # End if response is JSON as there is no log file to download |
| 206 | if content_type == "application/json": |
| 207 | logger.info("No more logs available") |
| 208 | return False |
| 209 | # Process log file |
| 210 | elif content_type == "application/octet-stream": |
| 211 | logger.info("Content-Type: application/octet-stream") |
| 212 | file_name = resp_headers["Content-Disposition"].split('="') |
| 213 | file_name = file_name[1][:-1] |
| 214 | |
| 215 | # Save mc-siem-token page token to check point directory |
| 216 | await write_checkpoint_file( |
| 217 | checkpoint_filename, |
| 218 | resp_headers["mc-siem-token"], |
| 219 | ) |
| 220 | log_filename = os.path.join(log_file_path, file_name) |
| 221 | await write_log_file(log_filename, resp_body) |
| 222 | return None |
| 223 | |
| 224 | |
| 225 | async def write_checkpoint_file(filename: str, data: str): |
| 226 | """ |
| 227 | Writes the given data to the given file. |
| 228 | """ |
| 229 | async with aiofiles.open(filename, "w") as f: |
| 230 | await f.write(data) |
| 231 | |
| 232 | |
| 233 | async def write_log_file(filename: str, resp_body): |
| 234 | """ |
| 235 | Writes the given data to the given file. |
| 236 | """ |
| 237 | if ".zip" in filename: |
| 238 | try: |
| 239 | byte_content = io.BytesIO(resp_body) |
| 240 | zip_file = ZipFile(byte_content) |
| 241 | zip_file.extractall(filename) |
| 242 | except Exception as e: |
| 243 | logger.error(f"Unable to extract zip file. Exception: {e}") |
| 244 | raise HTTPException(status_code=400, detail="Unable to extract zip file.") |
| 245 | else: |
| 246 | async with aiofiles.open(filename, "w") as f: |
| 247 | await f.write(resp_body) |
| 248 | |
| 249 | |
| 250 | async def process_log_file( |
| 251 | filename: str, |
| 252 | filename2: str, |
| 253 | log_file_path: str, |
| 254 | customer_code: str, |
| 255 | ): |
| 256 | """ |
| 257 | Process a log file by reading its contents and shipping events. |
| 258 | """ |
| 259 | log_file_full_path = build_log_file_path(log_file_path, filename, filename2) |
| 260 | file_creation_time = get_file_creation_time(log_file_full_path) |
| 261 | logger.info(f"File creation time: {file_creation_time} and filename: {filename2}") |
| 262 | |
| 263 | await read_and_ship_log_file(log_file_full_path, customer_code) |
| 264 | await safely_delete_file(log_file_full_path) |
| 265 | |
| 266 | |
| 267 | def build_log_file_path(log_file_path: str, filename: str, filename2: str) -> str: |
| 268 | """ |
| 269 | Constructs the full path for a log file. |
| 270 | """ |
| 271 | return os.path.join(log_file_path, filename, filename2) |
| 272 | |
| 273 | |
| 274 | def get_file_creation_time(file_path: str) -> str: |
| 275 | """ |
| 276 | Returns the creation time of a file. |
| 277 | """ |
| 278 | return time.ctime(os.path.getctime(file_path)) |
| 279 | |
| 280 | |
| 281 | async def read_and_ship_log_file(file_path: str, customer_code: str): |
| 282 | """ |
| 283 | Reads a log file line by line, converts each line to JSON, and ships the event. |
| 284 | """ |
| 285 | with open(file_path, "r", encoding="utf-8") as file: |
| 286 | for line in file: |
| 287 | log_entry = convert_to_json(line) |
| 288 | message = EventShipperPayload( |
| 289 | customer_code=customer_code, |
| 290 | integration="mimecast", |
| 291 | version="1.0", |
| 292 | **log_entry, |
| 293 | ) |
| 294 | await event_shipper(message) |
| 295 | |
| 296 | |
| 297 | async def safely_delete_file(file_path: str): |
| 298 | """ |
| 299 | Attempts to delete a file and logs the outcome. |
| 300 | """ |
| 301 | try: |
| 302 | os.remove(file_path) |
| 303 | logger.info(f"Successfully deleted the file: {file_path}") |
| 304 | except OSError as e: |
| 305 | logger.error(f"Error: {e.strerror}. File: {file_path}") |
| 306 | |
| 307 | |
| 308 | def convert_to_json(log_line: str) -> dict: |
| 309 | """ |
| 310 | Converts a log line to a JSON object. |
| 311 | """ |
| 312 | log_dict = {} |
| 313 | for pair in log_line.split("|"): |
| 314 | if "=" in pair: |
| 315 | key, value = pair.split("=", 1) |
| 316 | log_dict[key.strip()] = value.strip() |
| 317 | return log_dict |
| 318 | |
| 319 | |
| 320 | async def delete_log_directory(log_file_path: str): |
| 321 | """ |
| 322 | Deletes the log directory for the given customer. |
| 323 | """ |
| 324 | try: |
| 325 | shutil.rmtree(log_file_path) |
| 326 | logger.info(f"Successfully deleted the directory: {log_file_path}") |
| 327 | except OSError as e: |
| 328 | raise HTTPException( |
| 329 | status_code=400, |
| 330 | detail=f"Error: {e.strerror}. Directory: {log_file_path}", |
| 331 | ) |
| 332 | |
| 333 | |
| 334 | async def invoke_mimecast( |
| 335 | mimecast_request: MimecastRequest, |
| 336 | auth_keys: MimecastAuthKeys, |
| 337 | ) -> MimecastResponse: |
| 338 | """ |
| 339 | Invokes the Mimecast integration. |
| 340 | """ |
| 341 | mimecast_base_url = await get_base_url(auth_keys) |
| 342 | try: |
| 343 | logger.info(f"mimecast_base_url: {mimecast_base_url.data.data[0].region.api}") |
| 344 | except Exception as e: |
| 345 | logger.error( |
| 346 | f"Unable to retrieve base URL for Mimecast integration. Exception: {e}", |
| 347 | ) |
| 348 | raise HTTPException( |
| 349 | status_code=400, |
| 350 | detail="Unable to retrieve base URL for Mimecast integration.", |
| 351 | ) |
| 352 | checkpoint_filename = await get_checkpoint_filename(mimecast_request.customer_code) |
| 353 | log_file_path = await get_log_file_path(mimecast_request.customer_code) |
| 354 | response = await get_mta_siem_logs( |
| 355 | checkpoint_filename, |
| 356 | mimecast_base_url.data.data[0].region.api, |
| 357 | auth_keys, |
| 358 | ) |
| 359 | |
| 360 | await process_response(response, checkpoint_filename, log_file_path) |
| 361 | for filename in os.listdir(log_file_path): |
| 362 | if os.path.isdir(os.path.join(log_file_path, filename)): |
| 363 | for filename2 in os.listdir(os.path.join(log_file_path, filename)): |
| 364 | await process_log_file( |
| 365 | filename, |
| 366 | filename2, |
| 367 | log_file_path, |
| 368 | customer_code=mimecast_request.customer_code, |
| 369 | ) |
| 370 | logger.info(f"Log file path: {log_file_path}") |
| 371 | else: |
| 372 | await process_log_file(filename, filename2, log_file_path) |
| 373 | |
| 374 | await delete_log_directory(log_file_path) |
| 375 | return MimecastResponse( |
| 376 | success=True, |
| 377 | message="Successfully invoked Mimecast integration.", |
| 378 | ) |
| 379 | |
| 380 | |
| 381 | # ! TTP URLS ! # |
| 382 | async def custom_datetime_format(dt: datetime.datetime) -> str: |
| 383 | """Format a datetime object to a custom ISO-like string.""" |
| 384 | return dt.strftime("%Y-%m-%dT%H:%M:%S%z").replace("+00:00", "+0000") |
| 385 | |
| 386 | |
| 387 | async def create_ttp_request_body( |
| 388 | mimecast_request: MimecastTTPURLSRequest, |
| 389 | ) -> RequestBody: |
| 390 | """Create a request body for the Mimecast API call.""" |
| 391 | meta_data = {} |
| 392 | if mimecast_request.pagination_token: |
| 393 | meta_data["pagination"] = {"pageToken": mimecast_request.pagination_token} |
| 394 | return RequestBody( |
| 395 | meta=meta_data, |
| 396 | data=[ |
| 397 | DataItem( |
| 398 | oldestFirst=False, |
| 399 | from_=mimecast_request.lower_bound, |
| 400 | route="all", |
| 401 | to=mimecast_request.upper_bound, |
| 402 | scanResult="all", |
| 403 | ), |
| 404 | ], |
| 405 | ) |
| 406 | |
| 407 | |
| 408 | async def invoke_mimecast_api_ttp_urls( |
| 409 | mimecast_request: MimecastTTPURLSRequest, |
| 410 | ) -> TtpURLResponseBody: |
| 411 | """Invoke the Mimecast API call to get TTP URLs.""" |
| 412 | logger.info("Mimecast TTP URL request received") |
| 413 | request_body = await create_ttp_request_body(mimecast_request) |
| 414 | request_dict = request_body.model_dump(by_alias=True) |
| 415 | logger.info(f"Request: {request_dict}") |
| 416 | for item in request_dict["data"]: |
| 417 | item["from"] = await custom_datetime_format(item["from"]) |
| 418 | item["to"] = await custom_datetime_format(item["to"]) |
| 419 | response = requests.post( |
| 420 | url=mimecast_request.BaseURL + "/api/ttp/url/get-logs", |
| 421 | headers=mimecast_request.headers.model_dump(by_alias=True), |
| 422 | data=str(request_dict), |
| 423 | ) |
| 424 | return TtpURLResponseBody(**response.json()) |
| 425 | |
| 426 | |
| 427 | async def get_ttp_urls( |
| 428 | mimecast_request: MimecastTTPURLSRequest, |
| 429 | customer_code: str, |
| 430 | ) -> MimecastResponse: |
| 431 | logger.info("Mimecast TTP URL request received") |
| 432 | # Get the BaseURL for the Mimecast integration |
| 433 | mimecast_base_url = await get_base_url( |
| 434 | MimecastAuthKeys( |
| 435 | APP_ID=mimecast_request.ApplicationID, |
| 436 | APP_KEY=mimecast_request.ApplicationKey, |
| 437 | ACCESS_KEY=mimecast_request.AccessKey, |
| 438 | SECRET_KEY=mimecast_request.SecretKey, |
| 439 | EMAIL_ADDRESS=mimecast_request.EmailAddress, |
| 440 | URI="/api/login/discover-authentication", |
| 441 | ), |
| 442 | ) |
| 443 | # Add it to the request object |
| 444 | mimecast_request.BaseURL = mimecast_base_url.data.data[0].region.api |
| 445 | |
| 446 | mimecast_request.pagination_token = None # Initialize pagination_token to None |
| 447 | |
| 448 | while True: |
| 449 | response = await invoke_mimecast_api_ttp_urls(mimecast_request) |
| 450 | logger.info(f"Response: {response}") |
| 451 | |
| 452 | for data in response.data[0].clickLogs: |
| 453 | message = EventShipperPayload( |
| 454 | customer_code=customer_code, |
| 455 | integration="mimecast", |
| 456 | version="1.0", |
| 457 | **data.model_dump(by_alias=True), |
| 458 | ) |
| 459 | await event_shipper(message) |
| 460 | |
| 461 | # Check if there is a "next" page token in the response |
| 462 | next_page_token = response.meta.pagination.next |
| 463 | |
| 464 | if not next_page_token: |
| 465 | break # No more pages, break the loop |
| 466 | |
| 467 | # Update the pagination_token for the next API call |
| 468 | mimecast_request.pagination_token = next_page_token |
| 469 | |
| 470 | return MimecastResponse( |
| 471 | success=True, |
| 472 | message="Mimecast TTP URL request successful", |
| 473 | ) |