add get_ttp_urls to mimecast integration
Taylor committed
Jan 29, 2024 at 15:05 UTC
0edfa97e01f8c2b05faa11a96729a1439ba70319
3 files changed
+345
-5
backend/app/integrations/mimecast/routes/mimecast.py
+35
-2
@@ -11,8 +11,8 @@ from app.auth.utils import AuthHandler
11
from app.db.db_session import get_db
12
from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys
13
from app.integrations.mimecast.schema.mimecast import MimecastRequest
14
-from app.integrations.mimecast.schema.mimecast import MimecastResponse
15
-from app.integrations.mimecast.services.mimecast import invoke_mimecast
14
+from app.integrations.mimecast.schema.mimecast import MimecastResponse, MimecastTTPURLSRequest, MimecastHeaders
15
+from app.integrations.mimecast.services.mimecast import invoke_mimecast, get_ttp_urls
16
from app.integrations.routes import find_customer_integration
17
from app.integrations.routes import get_customer_integrations_by_customer_code
18
from app.integrations.schema import CustomerIntegrations
@@ -104,3 +104,36 @@ async def invoke_mimecast_route(mimecast_request: MimecastRequest, session: Asyn
104
auth_keys = MimecastAuthKeys(**mimecast_auth_keys)
105
106
return await invoke_mimecast(mimecast_request, auth_keys)
107
+
108
+@integration_mimecast_router.post(
109
+ "/ttp/urls",
110
+ response_model=MimecastResponse,
111
+ description="Pull down Mimecast TTP URLs for a given time range. "
112
+ "Link to docs: https://integrations.mimecast.com/documentation/endpoint-reference/logs-and-statistics/get-ttp-url-logs/ ",
113
+)
114
+async def mimecast_ttp_url_route(mimecast_request: MimecastRequest, session: AsyncSession = Depends(get_db)):
115
+ logger.info("Mimecast TTP URL request received")
116
+ customer_code=mimecast_request.customer_code
117
+ customer_integration_response = await get_customer_integration_response(mimecast_request.customer_code, session)
118
+
119
+ customer_integration = await find_customer_integration(
120
+ mimecast_request.customer_code,
121
+ mimecast_request.integration_name,
122
+ customer_integration_response,
123
+ )
124
+
125
+ mimecast_auth_keys = extract_mimecast_auth_keys(customer_integration)
126
+
127
+ auth_keys = MimecastAuthKeys(**mimecast_auth_keys)
128
+
129
+ mimecast_request = MimecastTTPURLSRequest(
130
+ ApplicationID=auth_keys.APP_ID,
131
+ ApplicationKey=auth_keys.APP_KEY,
132
+ AccessKey=auth_keys.ACCESS_KEY,
133
+ SecretKey=auth_keys.SECRET_KEY,
134
+ EmailAddress=auth_keys.EMAIL_ADDRESS,
135
+ time_range=mimecast_request.time_range,
136
+ )
137
+ logger.info(f"Mimecast TTP URL request: {mimecast_request}")
138
+
139
+ return await get_ttp_urls(mimecast_request, customer_code=customer_code)
backend/app/integrations/mimecast/schema/mimecast.py
+221
-2
@@ -1,9 +1,18 @@
1
from enum import Enum
2
from typing import List
3
+from datetime import datetime
4
+from datetime import timedelta
5
6
from pydantic import BaseModel
7
+from typing import Dict
8
from pydantic import Field
9
from pydantic import HttpUrl
10
+import base64
11
+import hashlib
12
+import hmac
13
+import uuid
14
+from pydantic import root_validator
15
+from typing import Optional
16
17
18
class PipelineRuleTitles(Enum):
@@ -29,6 +38,11 @@ class MimecastRequest(BaseModel):
38
description="The integration name.",
39
examples=["Office365"],
40
)
41
+ time_range: Optional[str] = Field(
42
+ "15m",
43
+ pattern="^[1-9][0-9]*[mhdw]$",
44
+ description="Time range for the query (1m, 1h, 1d, 1w)",
45
+ )
46
47
# # ensure the `integration_name` is always set to "Office365"
48
# @root_validator(pre=True)
@@ -53,8 +67,8 @@ class MimecastAuthKeys(BaseModel):
67
description="YOUR DEVELOPER APPLICATION KEY",
68
examples=["00002"],
69
)
56
- EMAIL_ADDRESS: str = Field(
57
- ...,
70
+ EMAIL_ADDRESS: Optional[str] = Field(
71
+ None,
72
description="EMAIL ADDRESS OF YOUR ADMINISTRATOR",
73
examples=["00002"],
74
)
@@ -109,3 +123,208 @@ class MimecastAPIEndpointResponse(BaseModel):
123
class MimecastScheduledResponse(BaseModel):
124
success: bool
125
message: str
126
+
127
+# ! MIMECAST TTP URLS ! #
128
+class MimecastHeaders(BaseModel):
129
+ Authorization: str = Field(
130
+ ...,
131
+ description="The Authorization header typically containing the access token.",
132
+ )
133
+ x_mc_app_id: str = Field(
134
+ ...,
135
+ alias="x-mc-app-id",
136
+ description="The Application ID for Mimecast.",
137
+ )
138
+ x_mc_date: str = Field(
139
+ ...,
140
+ alias="x-mc-date",
141
+ description="The date when the request was made.",
142
+ )
143
+ x_mc_req_id: str = Field(
144
+ ...,
145
+ alias="x-mc-req-id",
146
+ description="The unique request ID.",
147
+ )
148
+ Content_Type: str = Field(
149
+ ...,
150
+ alias="Content-Type",
151
+ description="The type of content, usually application/json.",
152
+ )
153
+
154
+ class Config:
155
+ allow_population_by_field_name = (
156
+ True # This allows field population by both alias and field name
157
+ )
158
+
159
+
160
+class MimecastTTPURLSRequest(BaseModel):
161
+ ApplicationID: str = Field(..., description="The ID of the Mimecast application.")
162
+ ApplicationKey: str = Field(
163
+ ...,
164
+ description="The key associated with the Mimecast application.",
165
+ )
166
+ AccessKey: str = Field(..., description="The access key for API authentication.")
167
+ SecretKey: str = Field(..., description="The secret key for API authentication.")
168
+ EmailAddress: str = Field(
169
+ ...,
170
+ description="The email address of the Mimecast administrator.",
171
+ )
172
+ BaseURL: Optional[str] = Field(None, description="The base URL for the Mimecast API.")
173
+ time_range: Optional[str] = Field(
174
+ "15m",
175
+ pattern="^[1-9][0-9]*[mhdw]$",
176
+ description="Time range for the query (1m, 1h, 1d, 1w)",
177
+ )
178
+ # headers: Dict[str, str] = Field(default_factory=dict) # default empty dictionary
179
+ headers: Optional[MimecastHeaders] = Field(
180
+ None,
181
+ description="The headers generated for the request.",
182
+ )
183
+ pagination_token: str = Field(None, description="Pagination token for API calls")
184
+
185
+ lower_bound: str = None
186
+ upper_bound: str = None
187
+
188
+ def __init__(self, *args, **kwargs):
189
+ super().__init__(*args, **kwargs)
190
+ self.generate_headers("/api/ttp/url/get-logs") # default URI
191
+
192
+ @root_validator(pre=True)
193
+ def set_time_bounds(cls, values):
194
+ time_range = values.get("time_range")
195
+ if time_range:
196
+ unit = time_range[-1]
197
+ amount = int(time_range[:-1])
198
+
199
+ now = datetime.utcnow()
200
+
201
+ if unit == "m":
202
+ lower_bound = now - timedelta(minutes=amount)
203
+ elif unit == "h":
204
+ lower_bound = now - timedelta(hours=amount)
205
+ elif unit == "d":
206
+ lower_bound = now - timedelta(days=amount)
207
+ elif unit == "w":
208
+ lower_bound = now - timedelta(weeks=amount)
209
+
210
+ values["lower_bound"] = (
211
+ lower_bound.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
212
+ )
213
+ values["upper_bound"] = now.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
214
+ return values
215
+
216
+ def generate_headers(self, uri: str) -> dict:
217
+ """Generate Mimecast request headers."""
218
+
219
+ request_id = str(uuid.uuid4())
220
+ hdr_date = datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S") + " UTC"
221
+
222
+ dataToSign = ":".join([hdr_date, request_id, uri, self.ApplicationKey])
223
+
224
+ hmac_sha1 = hmac.new(
225
+ base64.b64decode(self.SecretKey),
226
+ dataToSign.encode(),
227
+ digestmod=hashlib.sha1,
228
+ ).digest()
229
+ sig = base64.b64encode(hmac_sha1).rstrip()
230
+
231
+ headers_dict = {
232
+ "Authorization": "MC " + self.AccessKey + ":" + sig.decode(),
233
+ "x-mc-app-id": self.ApplicationID,
234
+ "x-mc-date": hdr_date,
235
+ "x-mc-req-id": request_id,
236
+ "Content-Type": "application/json",
237
+ }
238
+
239
+ self.headers = MimecastHeaders(
240
+ **headers_dict
241
+ ) # Create a new instance of MimecastHeaders and assign
242
+
243
+ return headers_dict
244
+##### ! SENDING REQUEST TO MIMECAST ! #####
245
+class DataItem(BaseModel):
246
+ oldestFirst: bool = Field(..., description="Ordering flag, oldest first if true.")
247
+ from_: datetime = Field(
248
+ ...,
249
+ alias="from",
250
+ description="Start date-time in ISO 8601 format.",
251
+ )
252
+ to: datetime = Field(..., description="End date-time in ISO 8601 format.")
253
+ route: str = Field(..., description="Routing information.")
254
+ scanResult: str = Field(..., description="Scan result.")
255
+
256
+ class Config:
257
+ allow_population_by_field_name = (
258
+ True # This allows field population by both alias and field name
259
+ )
260
+
261
+class RequestBody(BaseModel):
262
+ meta: Dict = Field({}, description="Meta information.")
263
+ data: List[DataItem] = Field(..., description="List of data items.")
264
+
265
+class TTPResponseClickLogs(BaseModel):
266
+ userEmailAddress: str
267
+ fromUserEmailAddress: str
268
+ url: str
269
+ ttpDefinition: str
270
+ subject: str
271
+ action: str
272
+ adminOverride: str
273
+ userOverride: str
274
+ scanResult: str
275
+ category: str
276
+ sendingIp: str
277
+ userAwarenessAction: str
278
+ date: str
279
+ actions: str
280
+ route: str
281
+ creationMethod: str
282
+ emailPartsDescription: List[str]
283
+ messageId: str
284
+
285
+
286
+class TTPResponseDataItem(BaseModel):
287
+ clickLogs: List[TTPResponseClickLogs]
288
+
289
+
290
+class TTPResponsePagination(BaseModel):
291
+ pageSize: int
292
+ totalCount: int
293
+ next: Optional[str]
294
+
295
+
296
+class ResponseMeta(BaseModel):
297
+ pagination: TTPResponsePagination
298
+ status: int
299
+
300
+
301
+class TtpURLResponseBody(BaseModel):
302
+ meta: ResponseMeta
303
+ data: List[TTPResponseDataItem]
304
+ fail: List[Dict] # Adjust this based on the actual structure of the "fail" field
305
+
306
+
307
+class ResponseAttachmentLogs(BaseModel):
308
+ senderAddress: str
309
+ recipientAddress: str
310
+ fileName: str
311
+ fileType: str
312
+ result: str
313
+ actionTriggered: str
314
+ date: str
315
+ details: str
316
+ route: str
317
+ messageId: str
318
+ subject: str
319
+ fileHash: str
320
+ definition: str
321
+
322
+
323
+class ResponseDataItemAttachment(BaseModel):
324
+ attachmentLogs: List[ResponseAttachmentLogs]
325
+
326
+
327
+class TtpURLAttachmentResponseBody(BaseModel):
328
+ meta: ResponseMeta
329
+ data: List[ResponseDataItemAttachment]
330
+ fail: List[Dict] # Adjust this based on the actual structure of the "fail" field
backend/app/integrations/mimecast/services/mimecast.py
+89
-1
@@ -18,7 +18,7 @@ from loguru import logger
18
from app.integrations.mimecast.schema.mimecast import MimecastAPIEndpointResponse
19
from app.integrations.mimecast.schema.mimecast import MimecastAuthKeys
20
from app.integrations.mimecast.schema.mimecast import MimecastRequest
21
-from app.integrations.mimecast.schema.mimecast import MimecastResponse
21
+from app.integrations.mimecast.schema.mimecast import MimecastResponse, MimecastTTPURLSRequest, TtpURLResponseBody, RequestBody, DataItem
22
from app.integrations.utils.collection import send_post_request
23
from app.integrations.utils.event_shipper import event_shipper
24
from app.integrations.utils.schema import EventShipperPayload
@@ -315,3 +315,91 @@ async def invoke_mimecast(mimecast_request: MimecastRequest, auth_keys: Mimecast
315
316
await delete_log_directory(log_file_path)
317
return MimecastResponse(success=True, message="Successfully invoked Mimecast integration.")
318
+
319
+
320
+# ! TTP URLS ! #
321
+async def custom_datetime_format(dt: datetime.datetime) -> str:
322
+ """Format a datetime object to a custom ISO-like string."""
323
+ return dt.strftime("%Y-%m-%dT%H:%M:%S%z").replace("+00:00", "+0000")
324
+
325
+async def create_ttp_request_body(
326
+ mimecast_request: MimecastTTPURLSRequest,
327
+) -> RequestBody:
328
+ """Create a request body for the Mimecast API call."""
329
+ meta_data = {}
330
+ if mimecast_request.pagination_token:
331
+ meta_data["pagination"] = {"pageToken": mimecast_request.pagination_token}
332
+ return RequestBody(
333
+ meta=meta_data,
334
+ data=[
335
+ DataItem(
336
+ oldestFirst=False,
337
+ from_=mimecast_request.lower_bound,
338
+ route="all",
339
+ to=mimecast_request.upper_bound,
340
+ scanResult="all",
341
+ ),
342
+ ],
343
+ )
344
+
345
+async def invoke_mimecast_api_ttp_urls(
346
+ mimecast_request: MimecastTTPURLSRequest,
347
+) -> TtpURLResponseBody:
348
+ """Invoke the Mimecast API call to get TTP URLs."""
349
+ logger.info("Mimecast TTP URL request received")
350
+ request_body = await create_ttp_request_body(mimecast_request)
351
+ request_dict = request_body.dict(by_alias=True)
352
+ logger.info(f"Request: {request_dict}")
353
+ for item in request_dict["data"]:
354
+ item["from"] = await custom_datetime_format(item["from"])
355
+ item["to"] = await custom_datetime_format(item["to"])
356
+ response = requests.post(
357
+ url=mimecast_request.BaseURL + "/api/ttp/url/get-logs",
358
+ headers=mimecast_request.headers.dict(by_alias=True),
359
+ data=str(request_dict),
360
+ )
361
+ return TtpURLResponseBody(**response.json())
362
+
363
+async def get_ttp_urls(mimecast_request: MimecastTTPURLSRequest, customer_code: str) -> MimecastResponse:
364
+ logger.info("Mimecast TTP URL request received")
365
+ # Get the BaseURL for the Mimecast integration
366
+ mimecast_base_url = await get_base_url(MimecastAuthKeys(
367
+ APP_ID=mimecast_request.ApplicationID,
368
+ APP_KEY=mimecast_request.ApplicationKey,
369
+ ACCESS_KEY=mimecast_request.AccessKey,
370
+ SECRET_KEY=mimecast_request.SecretKey,
371
+ EMAIL_ADDRESS=mimecast_request.EmailAddress,
372
+ URI="/api/login/discover-authentication",
373
+ )
374
+ )
375
+ # Add it to the request object
376
+ mimecast_request.BaseURL = mimecast_base_url.data.data[0].region.api
377
+
378
+ mimecast_request.pagination_token = None # Initialize pagination_token to None
379
+
380
+ while True:
381
+ response = await invoke_mimecast_api_ttp_urls(mimecast_request)
382
+ logger.info(f"Response: {response}")
383
+
384
+ for data in response.data[0].clickLogs:
385
+ message = EventShipperPayload(
386
+ customer_code=customer_code,
387
+ integration="mimecast",
388
+ version="1.0",
389
+ **data.dict(by_alias=True),
390
+ )
391
+ await event_shipper(message)
392
+
393
+ # Check if there is a "next" page token in the response
394
+ next_page_token = response.meta.pagination.next
395
+
396
+ if not next_page_token:
397
+ break # No more pages, break the loop
398
+
399
+ # Update the pagination_token for the next API call
400
+ mimecast_request.pagination_token = next_page_token
401
+
402
+ return MimecastResponse(
403
+ success=True,
404
+ message="Mimecast TTP URL request successful",
405
+ )