riposte-marketplace/integrations/sentinelone/scripts/get_unified_alerts.py

import json, os, sys, urllib.request, urllib.parse, urllib.error

# Unified (UAM) alerts live behind a GraphQL endpoint, not the REST v2.1 tree the
# rest of this package uses. The console-wide alert model is a different object
# from /cloud-detection/alerts (see get_alerts): one alert per detection across
# endpoint, cloud, identity and third-party surfaces.
VIEW_TYPES = ("ALL", "CLOUD", "CUSTOM_ALERTS", "ENDPOINT", "IDENTITY", "STAR", "THIRD_PARTY")

# Fields requested per alert node. Kept in one place so the mapper and this
# query cannot drift apart.
NODE_FIELDS = """
  id name description severity status result analystVerdict attackSurfaces
  classification confidenceLevel externalId createdAt updatedAt firstSeenAt
  lastSeenAt detectedAt noteExists dataSources storylineId ticketId
  fileName fileHash
  analytics { uid name type typeValue category }
  assignee { userId fullName email }
  detectionSource { product vendor engine }
  asset {
    id name agentUuid agentVersion assetTypeClassifier category subcategory type
    connectivityToConsole osType osVersion pendingReboot lastLoggedInUser policy
  }
  detectionTime {
    asset {
      agentVersion consoleIpAddress domain ipV4 ipV6 lastLoggedInUser
      osName osRevision osType policy subscriptionTime
    }
    attacker { host ip }
    cloud { accountId cloudProvider image instanceId instanceSize location network tags }
    kubernetes {
      clusterName namespaceName nodeName podName containerId containerImageName
      containerName containerNetworkStatus controllerName controllerType
    }
    scope { accountId accountName groupName siteName }
    targetUser { name emailAddress domain }
  }
  process { cmdLine parentName file { certSubject md5 name path sha1 sha256 } }
  realTime { scope { account { id name } group { id name } site { id name } } }
"""


def request(method, url, headers, body=None):
    data = json.dumps(body).encode("utf-8") if body is not None else None
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=60) as resp:
        raw = resp.read()
        return json.loads(raw) if raw else {}


def epoch_ms(v):
    """Normalise a watermark to epoch milliseconds.

    The alert-rule scheduler injects the watermark raw, in whichever format the
    operator picked for the rule (iso8601 | unix | unix_ms), so accept all three
    rather than forcing one and failing silently on the others.
    """
    s = str(v or "").strip()
    if not s:
        return None
    if s.isdigit():
        n = int(s)
        return n * 1000 if n < 100000000000 else n  # seconds vs milliseconds
    import datetime
    t = s.replace("Z", "+00:00")
    try:
        dt = datetime.datetime.fromisoformat(t)
    except ValueError:
        raise ValueError("cannot read timestamp %r: expected ISO8601 or an epoch" % s)
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=datetime.timezone.utc)
    return int(dt.timestamp() * 1000)


def gql_string(v):
    """Quote a value for inline use in the query (json.dumps escapes to GraphQL's rules)."""
    return json.dumps(str(v))


def build_filters(inputs):
    out = []
    start = epoch_ms(inputs.get("created_after"))
    end = epoch_ms(inputs.get("created_until"))
    if start is not None or end is not None:
        parts = []
        if start is not None:
            parts.append("start: %d" % start)
        if end is not None:
            parts.append("end: %d" % end)
        out.append('{ fieldId: "createdAt" dateTimeRange: { %s } }' % " ".join(parts))
    updated = epoch_ms(inputs.get("updated_after"))
    if updated is not None:
        out.append('{ fieldId: "updatedAt" dateTimeRange: { start: %d } }' % updated)
    for field, key in (("status", "status"), ("severity", "severity"),
                       ("analystVerdict", "analyst_verdict"), ("id", "alert_id")):
        v = str(inputs.get(key) or "").strip()
        if v:
            out.append('{ fieldId: "%s" stringEqual: { value: %s } }' % (field, gql_string(v)))
    return out


def main():
    secrets = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
    inputs = json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))
    base = secrets.get("url", "").rstrip("/") + "/web/api/v2.1"
    headers = {
        "Authorization": "ApiToken " + secrets.get("api_token", ""),
        "Accept": "application/json",
        "Content-Type": "application/json",
    }

    view_type = str(inputs.get("view_type") or "ALL").strip().upper()
    if view_type not in VIEW_TYPES:
        raise ValueError("view_type must be one of %s" % ", ".join(VIEW_TYPES))

    limit = int(inputs.get("limit") or 100)
    sort_by = "updatedAt" if str(inputs.get("updated_after") or "").strip() else "createdAt"
    filters = build_filters(inputs)
    filter_clause = "filters: [ %s ]" % " ".join(filters) if filters else ""

    # Ascending, so a truncated run leaves the watermark on the oldest unseen
    # alert rather than skipping everything between it and the newest.
    nodes, cursor, page_info = [], str(inputs.get("cursor") or "").strip(), {}
    total = None
    while len(nodes) < limit:
        page = min(limit - len(nodes), 100)
        query = """query Alerts {
  alerts(
    first: %d
    viewType: %s
    after: %s
    sort: { by: "%s", order: ASC }
    %s
  ) {
    totalCount
    edges { node { %s } }
    pageInfo { hasNextPage hasPreviousPage startCursor endCursor }
  }
}""" % (page, view_type, gql_string(cursor) if cursor else "null", sort_by, filter_clause, NODE_FIELDS)

        resp = request("POST", base + "/unifiedalerts/graphql", headers, {"query": query})
        # GraphQL answers 200 with an errors array; without this an unusable
        # query would look like an endless run of empty, successful fetches.
        if resp.get("errors"):
            print(json.dumps({"error": "GraphQL error", "detail": resp["errors"]}))
            sys.exit(1)

        alerts = (resp.get("data") or {}).get("alerts") or {}
        if total is None:
            total = alerts.get("totalCount")
        edges = alerts.get("edges") or []
        for edge in edges:
            node = edge.get("node")
            if node:
                nodes.append(node)
        page_info = alerts.get("pageInfo") or {}
        cursor = page_info.get("endCursor")
        # An empty page ends the walk even when the console still claims a next
        # one: without this, a console answering hasNextPage on no edges would
        # keep the fetch spinning until the run is killed.
        if not edges or not page_info.get("hasNextPage") or not cursor:
            break

    # Flattened to {data: [...]} so the alert rule reads a plain array of alerts
    # (results_path: data) and dedups on a top-level id, matching the REST
    # commands in this package. The edges/node envelope carries no id of its own.
    print(json.dumps({
        "data": nodes,
        "pagination": {
            "totalCount": total,
            "returned": len(nodes),
            "hasNextPage": bool(page_info.get("hasNextPage")),
            "endCursor": page_info.get("endCursor"),
        },
    }))


try:
    main()
except urllib.error.HTTPError as e:
    print(json.dumps({"error": "HTTP " + str(e.code), "detail": e.read().decode("utf-8", "replace")}))
    sys.exit(1)
except Exception as e:
    print(json.dumps({"error": str(e)}))
    sys.exit(1)