import demistomock as demisto # noqa: F401 from CommonServerPython import * # noqa: F401 from CommonServerUserPython import * # noqa from MicrosoftApiModule import * # noqa: E402 import traceback from datetime import datetime, timedelta, UTC from typing import Any # ============================================================================ # Constants # ============================================================================ class Config: """Global static configuration.""" VENDOR = "microsoft" PRODUCT = "o365_message_trace" APP_NAME = "o365-message-trace" GRAPH_SCOPE = "https://graph.microsoft.com/.default" MESSAGE_TRACES_PATH = "v1.0/admin/exchange/tracing/messageTraces" DATE_FORMAT_FILTER = "%Y-%m-%dT%H:%M:%S.%fZ" DATE_FORMAT_EVENT = "%Y-%m-%dT%H:%M:%S.%fZ" DEFAULT_MAX_EVENTS = 50000 DEFAULT_PAGE_SIZE = 1000 # API default/maximum per page DEFAULT_FIRST_FETCH_MINUTES = 1 # Each fetch cycle only scans this many minutes starting from ``last_fetch``. # This keeps every run small and bounded even when the # integration is far behind, so a large backlog is drained oldest-first # across many runs instead of re-downloading days of events on every run. FETCH_WINDOW_MINUTES = 5 # Fixed backoff schedule (in seconds) applied between retries when the Graph # API responds with HTTP 429 (Too Many Requests). RATE_LIMIT_BACKOFFS = (30, 60, 90) # ============================================================================ # Microsoft client # ============================================================================ class O365MessageTraceClient(MicrosoftClient): """Thin subclass of :class:`MicrosoftClient` for the O365 Message Trace integration. It overrides :meth:`MicrosoftClient.http_request` while keeping the exact same logic as the parent implementation. All other behavior (authentication flows, token retrieval, metrics and error parsing) is inherited unchanged. """ def http_request( self, *args, resp_type="json", headers=None, return_empty_response=False, scope: str | None = None, resource: str = "", **kwargs, ): """ Overrides Base client request function, retrieves and adds to headers access token before sending the request. Args: resp_type: Type of response to return. will be ignored if `return_empty_response` is True. headers: Headers to add to the request. return_empty_response: Return the response itself if the return_code is 206. scope: A scope to request. Currently, will work only with self-deployed app. resource (str): The resource identifier for which the generated token will have access to. Returns: Response from api according to resp_type. The default is `json` (dict or list). """ token = self.get_access_token(resource=resource, scope=scope) default_headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json", "Accept": "application/json"} if headers: default_headers |= headers if self.timeout: kwargs["timeout"] = self.timeout kwargs["error_handler"] = self.handle_error_with_metrics response = super()._http_request( # type: ignore[misc] *args, resp_type="response", headers=default_headers, status_list_to_retry=[503, 429], backoff_factor=30, retries=3, ok_codes=(200, 201, 202, 204, 206, 404), **kwargs, ) MicrosoftClient.create_api_metrics(response.status_code) # 206 indicates Partial Content, reason will be in the warning header. # In that case, logs with the warning header will be written. if response.status_code == 206: demisto.debug(str(response.headers)) is_response_empty_and_successful = response.status_code == 204 if is_response_empty_and_successful and return_empty_response: return response # Handle 404 errors instead of raising them as exceptions: if response.status_code == 404: try: error_message = response.json() except Exception: error_message = "Not Found - 404 Response" raise NotFoundError(error_message) try: if resp_type == "json": return response.json() if resp_type == "text": return response.text if resp_type == "content": return response.content if resp_type == "xml": try: import defusedxml.ElementTree as defused_ET defused_ET.fromstring(response.text) except ImportError: demisto.debug("defused_ET is not supported, using ET instead.") ET.fromstring(response.text) return response except ValueError as exception: raise DemistoException(f"Failed to parse json object from response: {response.content}", exception) # ============================================================================ # Client # ============================================================================ class Client: """Microsoft Graph client for O365 Message Trace events. Composes an instance of :class:`MicrosoftClient` (from ``MicrosoftApiModule``) so that the integration supports all standard Microsoft authentication methods: client credentials, certificate (thumbprint + private key), authorization-code and Azure Managed Identities. """ def __init__( self, tenant_id: str, auth_id: str, enc_key: str | None, app_name: str, base_url: str, verify: bool, proxy: bool, certificate_thumbprint: str | None = None, private_key: str | None = None, auth_code: str | None = None, redirect_uri: str | None = None, managed_identities_client_id: str | None = None, azure_cloud: AzureCloud = AZURE_WORLDWIDE_CLOUD, ): grant_type = AUTHORIZATION_CODE if auth_code and redirect_uri else CLIENT_CREDENTIALS demisto.debug(f"[Auth] Using grant type: {grant_type}") client_args = { "tenant_id": tenant_id, "auth_id": auth_id, "enc_key": enc_key, "app_name": app_name, "base_url": base_url, "verify": verify, "proxy": proxy, "self_deployed": True, "certificate_thumbprint": certificate_thumbprint, "private_key": private_key, "auth_code": auth_code or "", "redirect_uri": redirect_uri, "grant_type": grant_type, "scope": Config.GRAPH_SCOPE, "resource": Resources.graph, "azure_cloud": azure_cloud, "azure_ad_endpoint": azure_cloud.endpoints.active_directory, "token_retrieval_url": urljoin(azure_cloud.endpoints.active_directory, f"/{tenant_id}/oauth2/v2.0/token"), "managed_identities_client_id": managed_identities_client_id, "managed_identities_resource_uri": Resources.graph, "command_prefix": Config.APP_NAME, "retry_on_rate_limit": True, "timeout": 60, } self.ms_client = O365MessageTraceClient(**client_args) # ------------------------------------------------------------------ # API calls # ------------------------------------------------------------------ def get_message_traces_page( self, start_date: str | None = None, end_date: str | None = None, next_link: str | None = None, page_size: int = Config.DEFAULT_PAGE_SIZE, ) -> dict[str, Any]: """Fetch a single page of message-trace records. When ``next_link`` is supplied it is used as-is (it already contains the required ``$skiptoken``); otherwise a fresh filtered request is issued using ``start_date`` and ``end_date``. """ if next_link: demisto.debug(f"[API] Following @odata.nextLink: {next_link}") return self.ms_client.http_request(method="GET", full_url=next_link, url_suffix="") params = { "$filter": f"receivedDateTime ge {start_date} and receivedDateTime le {end_date}", "$top": page_size, } demisto.debug(f"[API] First page request | params={params}") return self.ms_client.http_request(method="GET", url_suffix=Config.MESSAGE_TRACES_PATH, params=params) # ============================================================================ # Helpers # ============================================================================ def parse_datetime(value: str | None, default: datetime | None = None) -> datetime: """Parse a date string and always return a tz-aware UTC datetime.""" if not value: return default or datetime.now(UTC) parsed = arg_to_datetime(arg=value, is_utc=True) if not parsed: return default or datetime.now(UTC) if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=UTC) return parsed def format_datetime_for_filter(dt: datetime) -> str: """Format a datetime in the form expected by the Graph $filter clause.""" return dt.strftime(Config.DATE_FORMAT_FILTER) def deduplicate_events(events: list[dict], seen_ids: set[str]) -> list[dict]: """Filter out events whose IDs are already in ``seen_ids``.""" if not seen_ids: return events new_events: list[dict] = [] duplicates = 0 for event in events: event_id = event.get("_unique_id") if event_id and event_id in seen_ids: duplicates += 1 continue new_events.append(event) if duplicates: demisto.debug(f"[Dedup] Skipped {duplicates} duplicate events") return new_events def add_time_field(events: list[dict]) -> None: """Add the XSIAM-required ``_time`` field to each event.""" fallback_time = datetime.now(UTC).strftime(Config.DATE_FORMAT_FILTER) for event in events: received = event.get("receivedDateTime") event["_time"] = received if received else fallback_time def add_unique_id_field(events: list[dict]) -> None: """Add a ``_unique_id`` field to each event in the form ``|``. The new ``_unique_id`` field guarantees uniqueness across events that share the same underlying message id but were delivered to different recipients. """ for event in events: event_id = event.get("id") recipient = event.get("recipientAddress") if event_id and recipient: event["_unique_id"] = f"{event_id}|{recipient}" # ============================================================================ # Configuration parsing # ============================================================================ def parse_integration_params(params: dict[str, Any]) -> dict[str, Any]: """Parse and validate integration configuration parameters. Extracts authentication settings, connection settings, and validates the resulting authentication credentials from the raw ``demisto.params()`` dictionary. Supports client-credentials, certificate, authorization-code and Azure managed-identities flows. Args: params: Raw parameters from ``demisto.params()``. Returns: Validated configuration dictionary with the keys required to construct a :class:`Client` instance, plus ``max_events``. Raises: DemistoException: If required authentication credentials are missing or inconsistent for the resolved grant type. """ # ----- Tenant / Auth ID / Secret (support both creds objects and legacy plain params) ----- tenant_id = params.get("tenant_id", "") credentials_client_id = params.get("credentials_client_id") or {} client_id = credentials_client_id.get("password") credentials = params.get("credentials") or {} client_secret = credentials.get("password") or params.get("client_secret", "") # ----- Certificate auth ----- creds_certificate = params.get("creds_certificate") or {} certificate_thumbprint = creds_certificate.get("identifier") private_key_raw = creds_certificate.get("password") private_key = replace_spaces_in_credential(private_key_raw) if private_key_raw else None # ----- Authorization-code flow ----- auth_code_param = params.get("auth_code") or {} auth_code = auth_code_param.get("password") redirect_uri = params.get("redirect_uri") # ----- Managed Identities ----- managed_identities_client_id = get_azure_managed_identities_client_id(params) # ----- Common settings ----- azure_cloud = get_azure_cloud(params, "O365MessageTrace") base_url = (params.get("url") or urljoin(azure_cloud.endpoints.microsoft_graph_resource_id, "/")).rstrip("/") + "/" verify = not argToBoolean(params.get("insecure", False)) proxy = argToBoolean(params.get("proxy", False)) max_events = arg_to_number(params.get("max_fetch")) or Config.DEFAULT_MAX_EVENTS # ----- Validation ----- if not managed_identities_client_id: grant_type = AUTHORIZATION_CODE if auth_code and redirect_uri else CLIENT_CREDENTIALS if grant_type == AUTHORIZATION_CODE: if not tenant_id or not client_id or not client_secret or not auth_code or not redirect_uri: raise DemistoException( "Tenant ID, Client ID, Client Secret, Authorization code and Application redirect URI " "are required for the authorization code flow." ) elif grant_type == CLIENT_CREDENTIALS and (not tenant_id or not client_id or not client_secret): raise DemistoException("Tenant ID, Client ID and Client Secret are required for the client credentials flow.") if not client_secret and not (certificate_thumbprint and private_key) and not auth_code: raise DemistoException( "An authentication credential must be provided: Client Secret, " "Certificate Thumbprint + Private Key, or Authorization Code." ) return { "tenant_id": tenant_id, "auth_id": client_id or "", "enc_key": client_secret or None, "app_name": Config.APP_NAME, "base_url": base_url, "verify": verify, "proxy": proxy, "certificate_thumbprint": certificate_thumbprint, "private_key": private_key, "auth_code": auth_code, "redirect_uri": redirect_uri, "managed_identities_client_id": managed_identities_client_id, "azure_cloud": azure_cloud, "max_events": max_events, } # ============================================================================ # Core fetch logic # ============================================================================ def fetch_events_sequential( client: Client, start: datetime, end: datetime, max_events: int, ) -> list[dict]: """Fetch all message-trace pages for the [start, end] window sequentially. Iterates through ALL available pages using ``@odata.nextLink`` until no more pages remain, even when the running total already exceeds ``max_events``. The API returns the latest events first, so all events must be collected before they can be sorted ascending by ``receivedDateTime`` and truncated to the earliest ``max_events`` events. If the first page fails the exception is re-raised so the calling ``fetch_events`` cycle aborts and ``lastRun`` is not advanced (preventing data loss). If a later page fails we keep the events collected so far and log the failure - the next fetch cycle will resume from the high-water mark. """ if end <= start: demisto.debug(f"[Fetch] Empty time range ({start.isoformat()} -> {end.isoformat()}). Skipping.") return [] start_str = format_datetime_for_filter(start) end_str = format_datetime_for_filter(end) demisto.debug(f"[Fetch] Fetching window {start_str} -> {end_str} | max={max_events}") collected: list[dict] = [] next_link: str | None = None previous_link: str | None = None # Track to detect a non-advancing cursor. page_events: list[dict] = [{}] # Sentinel non-empty value to enter the loop. while page_events: try: response = client.get_message_traces_page( start_date=start_str, end_date=end_str, next_link=next_link, page_size=Config.DEFAULT_PAGE_SIZE, ) except Exception as e: demisto.error(f"[Fetch] Failed to fetch page for window {start_str} -> {end_str}: {e}\n{traceback.format_exc()}") # No events collected yet - propagate so lastRun is NOT updated # and we retry the same window on the next fetch cycle. if not collected: raise # We already have some events from previous pages - stop here and # let the caller persist what we have. break page_events = response.get("value", []) or [] collected.extend(page_events) demisto.debug( f"[Fetch] Window {start_str} -> {end_str}: page returned {len(page_events)} events " f"(running total: {len(collected)})" ) # Defensive stop #1: empty page means there is nothing more to read. if not page_events: demisto.debug(f"[Fetch] Window {start_str} -> {end_str}: empty page, stopping.") break next_link = response.get("@odata.nextLink") # Normal stop: no more pages. if not next_link: demisto.debug(f"[Fetch] Window {start_str} -> {end_str}: no more pages.") break # Defensive stop #2: the cursor did not advance (non-advancing/self-referential # @odata.nextLink). Without this, a misbehaving server could loop forever. if next_link == previous_link: demisto.error( f"[Fetch] Window {start_str} -> {end_str}: @odata.nextLink did not advance " f"({next_link}). Stopping to avoid an infinite loop." ) break previous_link = next_link # Sort all collected events ascending by receivedDateTime (parsed as datetime) so the # earliest event is first. collected.sort( key=lambda event: safe_strptime(event["receivedDateTime"], Config.DATE_FORMAT_EVENT) if event.get("receivedDateTime") else datetime.min ) if len(collected) > max_events: demisto.debug(f"[Fetch] Collected {len(collected)} events, truncating to max_events ({max_events}).") collected = collected[:max_events] return collected # ============================================================================ # Commands # ============================================================================ def test_module(client: Client) -> str: """Validate credentials and Graph connectivity by fetching a tiny window. Raises: DemistoException: If using the authorization code flow, since test-module cannot access the integration context required by that flow. The ``o365-message-trace-auth-test`` command should be used instead. """ demisto.debug("[Test] Starting test-module") if client.ms_client.grant_type == AUTHORIZATION_CODE: raise DemistoException( "Test module is not available for the authorization code flow. " "Use the o365-message-trace-auth-test command instead." ) try: end = datetime.now(UTC) start = end - timedelta(minutes=5) fetch_events_sequential(client, start, end, max_events=1) return "ok" except Exception as e: error_message = str(e) if "401" in error_message or "403" in error_message: return f"Authorization Error: verify Tenant ID, Client ID and authentication credentials. Details: {error_message}" raise def auth_test_command(client: Client) -> CommandResults: """Tests connectivity to Microsoft. Used to validate the authentication flow (especially the authorization-code flow) after the integration has been configured, since the standard test-module cannot access the integration context. """ demisto.debug("[Auth Test] Starting o365-message-trace-auth-test") try: end = datetime.now(UTC) start = end - timedelta(minutes=5) client.get_message_traces_page( start_date=format_datetime_for_filter(start), end_date=format_datetime_for_filter(end), page_size=1, ) except Exception as e: raise DemistoException(f"Authentication was not successful. Verify the configuration parameters. Error: {e}") from e return CommandResults(readable_output="Authentication was successful.") def get_events_command(client: Client, args: dict) -> CommandResults: """Manual command to retrieve events (and optionally push them to XSIAM).""" limit = arg_to_number(args.get("limit")) or 50 start_time = args.get("start_time") end_time = args.get("end_time") should_push_events = argToBoolean(args.get("should_push_events", False)) end_dt = parse_datetime(end_time, default=datetime.now(UTC)) start_dt = parse_datetime(start_time, default=end_dt - timedelta(minutes=Config.DEFAULT_FIRST_FETCH_MINUTES)) events = fetch_events_sequential(client, start_dt, end_dt, max_events=limit) add_unique_id_field(events) add_time_field(events) if should_push_events and events: send_events_to_xsiam(events=events, vendor=Config.VENDOR, product=Config.PRODUCT) readable = tableToMarkdown( "O365 Message Trace Events", events, removeNull=True, headerTransform=pascalToSpace, headers=["id", "receivedDateTime", "senderAddress", "recipientAddress", "subject", "status"], ) return CommandResults( readable_output=readable, outputs_prefix="O365MessageTrace.Event", outputs_key_field="id", outputs=events, ) def fetch_events(client: Client, max_events: int) -> None: """Scheduled fetch command - reads state, fetches, deduplicates, persists state.""" last_run = demisto.getLastRun() or {} demisto.debug(f"[Fetch] last_run={last_run}") last_fetch_str: str | None = last_run.get("last_fetch") seen_ids: list[str] = last_run.get("seen_ids", []) or [] now = datetime.now(UTC) if last_fetch_str: start_dt = parse_datetime(last_fetch_str) else: start_dt = now - timedelta(minutes=Config.DEFAULT_FIRST_FETCH_MINUTES) demisto.debug(f"[Fetch] First run - looking back {Config.DEFAULT_FIRST_FETCH_MINUTES} minutes from now") if now < start_dt + timedelta(minutes=Config.FETCH_WINDOW_MINUTES): window_end_dt = now demisto.debug(f"[Fetch] Window {start_dt.isoformat()} -> {window_end_dt.isoformat()} (now={now.isoformat()})") # Fetch all events in the window sequentially events = fetch_events_sequential(client, start_dt, window_end_dt, max_events=max_events) else: window_end_dt = min(start_dt + timedelta(minutes=Config.FETCH_WINDOW_MINUTES), now) all_events: list[dict] = [] # Walk fixed-size windows until we either catch up to ``now``, collect enough # events, or hit the safety cap on how many times ``fetch_events_sequential`` # may be called in a single run. while start_dt < now and len(all_events) < max_events: window_end_dt = min(start_dt + timedelta(minutes=Config.FETCH_WINDOW_MINUTES), now) try: events = fetch_events_sequential(client, start_dt, window_end_dt, max_events=max_events) except Exception as e: demisto.error(f"[Fetch] Error fetching events for window {start_dt} -> {window_end_dt}: {e}") break all_events.extend(events) # Guard: the window cursor MUST advance, or stop (mirrors the pagination guard). if window_end_dt <= start_dt: demisto.error(f"[Fetch] Window did not advance ({start_dt} -> {window_end_dt}); stopping.") break start_dt = window_end_dt window_end_dt = min(start_dt + timedelta(minutes=Config.FETCH_WINDOW_MINUTES), now) events = all_events if len(events) > max_events: events = events[:max_events] add_unique_id_field(events) add_time_field(events) demisto.debug(f"[Fetch] Fetched {len(events)} raw events before dedup") # Deduplicate against previous run's high-water-mark IDs new_events = deduplicate_events(events, set(seen_ids)) demisto.debug(f"[Fetch] {len(new_events)} new events after dedup") if new_events: send_events_to_xsiam(events=new_events, vendor=Config.VENDOR, product=Config.PRODUCT) demisto.debug(f"[Fetch] Sent {len(new_events)} events to XSIAM") # New high-water mark. With no events, advance to the window end and reset seen_ids. new_last_fetch = format_datetime_for_filter(window_end_dt) new_seen_ids: list[str] = [] # Use ALL fetched events (not just published ones): timestamps are second-granular, so # seen_ids must keep every ID at the boundary - including deduped-out ones - or the next # run (re-fetching at ``>= boundary``) would re-send already-sent events as duplicates. timed_events = [event for event in new_events if event.get("_time")] if timed_events: latest_time: str = max(event["_time"] for event in timed_events) new_last_fetch = latest_time ids_at_latest = [eid for event in timed_events if event.get("_time") == latest_time and (eid := event.get("_unique_id"))] # If the high-water mark hasn't moved, merge with the existing seen_ids if latest_time == last_fetch_str: new_seen_ids = list(set(seen_ids) | set(ids_at_latest)) else: new_seen_ids = ids_at_latest new_last_run = { "last_fetch": new_last_fetch, "seen_ids": new_seen_ids, } demisto.setLastRun(new_last_run) demisto.debug(f"[Fetch] Updated last_run={new_last_run}") # ============================================================================ # Main # ============================================================================ def main() -> None: # pragma: no cover params = demisto.params() args = demisto.args() command = demisto.command() demisto.debug(f"[Main] Command={command}") config = parse_integration_params(params) max_events = config.pop("max_events") try: client = Client(**config) # pylint: disable=E1123 if command == "test-module": return_results(test_module(client)) elif command == "o365-message-trace-auth-test": return_results(auth_test_command(client)) elif command == "o365-message-trace-auth-reset": return_results(reset_auth()) elif command == "o365-message-trace-get-events": return_results(get_events_command(client, args)) elif command == "fetch-events": fetch_events(client, max_events=max_events) elif command == "o365-message-trace-generate-login-url": return_results(generate_login_url(client.ms_client)) else: raise NotImplementedError(f"Command '{command}' is not implemented.") except Exception as e: error_msg = f"Failed to execute {command}. Error: {e!s}" demisto.error(f"{error_msg}\n{traceback.format_exc()}") return_error(error_msg) if __name__ in ("__main__", "__builtin__", "builtins"): main()