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

import json, os, sys, time, 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")

# Walking the cursor is bounded: filtering client-side (see the attempt ladder) a
# page can be discarded almost whole, and without a cap a narrow filter on a busy
# console would page until the run is killed.
MAX_PAGES = 50

# The console narrows to platform-rule alerts with analyticsCategory=Platform
# Rule, but the fieldIds this endpoint takes are a mix of flat names (createdAt,
# externalId) and dotted paths (asset.osType, assignee.userId), so both spellings
# of the category are tried before falling back to filtering what was fetched.
CATEGORY_FIELD_IDS = ("analytics.category", "analyticsCategory")

# Same reason: the sort argument is spelled one way or the other depending on the
# console's schema version, and a wrong guess fails the whole query rather than
# degrading, so both are tried.
SORT_FORMS = ("sort", "sorts")

# 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 } } }
"""

_UNITS = {"second": 1, "minute": 60, "hour": 3600, "day": 86400,
          "week": 604800, "month": 2592000, "year": 31536000}


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 rel_seconds(text):
    """Read a relative window like '24 hours' or 'last 7 days' into seconds."""
    num = unit = None
    for t in str(text).lower().replace("last", "").split():
        if t.isdigit():
            num = int(t)
        elif t.rstrip("s") in _UNITS:
            unit = t.rstrip("s")
    return num * _UNITS[unit] if (num is not None and unit) else None


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. A relative window
    ('24 hours') is accepted too, for a rule that reads a rolling window instead
    of a watermark.
    """
    s = str(v or "").strip()
    if not s:
        return None
    secs = rel_seconds(s)
    if secs is not None:
        return int((time.time() - secs) * 1000)
    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, an epoch or a window like '24 hours'" % 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 csv_list(v, upper=False):
    """Split a comma-separated input into values, dropping the empties."""
    if v in (None, ""):
        return []
    items = v if isinstance(v, list) else str(v).split(",")
    out = [str(x).strip() for x in items if str(x).strip()]
    return [x.upper() for x in out] if upper else out


def norm(v):
    """Fold a category to letters/digits, so "Platform Rule", "platform rule" and
    PLATFORM_RULE all compare equal (used only by the client-side fallback)."""
    return "".join(ch for ch in str(v or "").lower() if ch.isalnum())


def category_matches(node, categories):
    return norm(((node.get("analytics") or {}).get("category"))) in {norm(c) for c in categories}


def string_filter(field, values):
    """Equality on one value, membership on several — an unfiltered field is the
    caller's job to skip. stringEqual is kept for the single-value case because
    that is what every console here is known to accept."""
    if len(values) == 1:
        return '{ fieldId: "%s" stringEqual: { value: %s } }' % (field, gql_string(values[0]))
    return '{ fieldId: "%s" stringIn: { values: [%s] } }' % (
        field, ", ".join(gql_string(v) for v in values))


def date_range_filter(field, start, end):
    parts = []
    if start is not None:
        parts.append("start: %d" % start)
    if end is not None:
        parts.append("end: %d" % end)
    return '{ fieldId: "%s" dateTimeRange: { %s } }' % (field, " ".join(parts))


def build_filters(inputs, category_field):
    """Assemble the filter list. category_field is the fieldId to filter the
    analytic category on, or None to leave it to the client-side pass."""
    out = []

    # ── Time windows ──────────────────────────────────────────────────────
    for field, after_key, until_key in (("createdAt", "created_after", "created_until"),
                                        ("detectedAt", "detected_after", "detected_until")):
        start = epoch_ms(inputs.get(after_key))
        end = epoch_ms(inputs.get(until_key))
        if start is not None or end is not None:
            out.append(date_range_filter(field, start, end))
    updated = epoch_ms(inputs.get("updated_after"))
    if updated is not None:
        out.append(date_range_filter("updatedAt", updated, None))

    # ── Enum-valued fields: comma-separated, upper-cased ──────────────────
    for field, key in (("status", "status"),
                       ("severity", "severity"),
                       ("analystVerdict", "analyst_verdict"),
                       ("classification", "classification"),
                       ("attackSurfaces", "attack_surface"),
                       ("asset.osType", "os_type")):
        values = csv_list(inputs.get(key), upper=True)
        if values:
            out.append(string_filter(field, values))

    # ── Identifiers, taken as written ─────────────────────────────────────
    for field, key in (("id", "alert_id"),
                       ("asset.id", "asset_id"),
                       ("externalId", "external_id")):
        values = csv_list(inputs.get(key))
        if values:
            out.append(string_filter(field, values))

    # ── The detecting analytic's category (e.g. Platform Rule) ────────────
    if category_field:
        values = csv_list(inputs.get("analytics_category"))
        if values:
            out.append(string_filter(category_field, values))

    # ── Free text and the two console toggles ─────────────────────────────
    search = str(inputs.get("search_text") or "").strip()
    if search:
        out.append('{ fieldId: "name" match: { values: [%s] } }' % gql_string(search))
    if str(inputs.get("unassigned_only") or "").strip().lower() in ("1", "true", "yes"):
        out.append('{ fieldId: "assignee.userId" isNegated: true stringEqual: { value: null } }')
    if str(inputs.get("unmitigated_only") or "").strip().lower() in ("1", "true", "yes"):
        out.append('{ fieldId: "result" stringEqual: { value: "UNMITIGATED" } }')
    return out


def sort_clause(form, by):
    """ASC either way, so a truncated run leaves the watermark on the oldest
    unseen alert rather than skipping everything between it and the newest."""
    if form == "sorts":
        return 'sorts: [ { by: "%s", order: ASC } ]' % by
    return 'sort: { by: "%s", order: ASC }' % by


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"
    categories = csv_list(inputs.get("analytics_category"))

    # Query shapes to try, in order, until the console accepts one. Without a
    # category asked for there is nothing to degrade, so only the sort varies.
    category_modes = list(CATEGORY_FIELD_IDS) + [None] if categories else [None]
    attempts = [(s, c) for s in SORT_FORMS for c in category_modes]
    attempt = 0
    settled = False

    nodes, cursor, page_info = [], str(inputs.get("cursor") or "").strip(), {}
    total = None
    pages = 0
    stopped_mid_page = None
    while len(nodes) < limit and pages < MAX_PAGES:
        sort_form, category_field = attempts[attempt]
        client_side_category = bool(categories) and category_field is None
        # Server-side, every alert on a page is one we keep, so ask for exactly
        # what is missing. Filtered client-side, most of a page is discarded, so
        # read full pages rather than walking one alert at a time.
        page = 100 if client_side_category else min(limit - len(nodes), 100)
        page_cursor = cursor
        filters = build_filters(inputs, category_field)
        filter_clause = "filters: [ %s ]" % " ".join(filters) if filters else ""
        query = """query Alerts {
  alerts(
    first: %d
    viewType: %s
    after: %s
    %s
    %s
  ) {
    totalCount
    edges { node { %s } }
    pageInfo { hasNextPage hasPreviousPage startCursor endCursor }
  }
}""" % (page, view_type, gql_string(cursor) if cursor else "null",
        sort_clause(sort_form, 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"):
            # A console that does not know a fieldId or the sort spelling answers
            # with an errors array rather than ignoring it, so the next shape is
            # tried on the same cursor. Once a page has come back the shape is
            # settled and a later error is a real one.
            if not settled and attempt + 1 < len(attempts):
                attempt += 1
                continue
            print(json.dumps({"error": "GraphQL error", "detail": resp["errors"]}))
            sys.exit(1)

        settled = True
        pages += 1
        alerts = (resp.get("data") or {}).get("alerts") or {}
        if total is None:
            total = alerts.get("totalCount")
        edges = alerts.get("edges") or []
        for i, edge in enumerate(edges):
            node = edge.get("node")
            if not node:
                continue
            if client_side_category and not category_matches(node, categories):
                continue
            nodes.append(node)
            # Client-side, a page holds more alerts than the limit leaves room
            # for. A cursor only addresses page boundaries, so stopping here
            # reports the cursor this page was read with: the next run re-reads
            # the page and the rule dedups on id, rather than the tail of the
            # page being skipped for good.
            if len(nodes) >= limit and i + 1 < len(edges):
                stopped_mid_page = page_cursor
                break
        page_info = alerts.get("pageInfo") or {}
        cursor = page_info.get("endCursor")
        if stopped_mid_page is not None:
            break
        # 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

    sort_form, category_field = attempts[attempt]
    # 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 is the console's count for the filters it applied, so
            # with a client-side category it counts more than data[] holds.
            "totalCount": total,
            "returned": len(nodes),
            "hasNextPage": True if stopped_mid_page is not None else bool(page_info.get("hasNextPage")),
            "endCursor": (stopped_mid_page or None) if stopped_mid_page is not None else page_info.get("endCursor"),
            "pagesRead": pages,
            # Which query shape the console accepted — worth reading once per
            # tenant: 'client' means the category was filtered after the fetch.
            "categoryFilter": (None if not categories
                               else ("server:" + category_field if category_field else "client")),
            "sortForm": sort_form,
        },
    }))


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)