riposte-marketplace/integrations/sentinelone-sdl/scripts/get_alerts.py

import json, os, sys, time, urllib.request, urllib.parse, urllib.error
from datetime import datetime, timezone

GRAPHQL = "/web/api/v2.1/unifiedalerts/graphql"
_UNITS = {"second": 1, "minute": 60, "hour": 3600, "day": 86400, "week": 604800, "month": 2592000, "year": 31536000}

QUERY = """
query GetAlerts($first: Int!, $filters: [FilterInput!], $sorts: [SortInput!]) {
  alerts(first: $first, filters: $filters, sorts: $sorts, viewType: ALL) {
    totalCount
    edges { node {
      id externalId name description severity status classification confidenceLevel
      analystVerdict result detectedAt createdAt updatedAt firstSeenAt lastSeenAt
      storylineId attackSurfaces
      asset { id name osType osVersion agentVersion lastLoggedInUser }
      assignee { userId fullName email }
      detectionSource { vendor product engine }
      process { cmdLine parentName file { name path md5 sha1 sha256 } }
      analytics { category name type uid }
    } }
  }
}
"""


def _cfg():
    s = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
    base = str(s.get("url") or "").rstrip("/")
    headers = {"Authorization": "ApiToken " + s.get("api_token", ""), "Content-Type": "application/json", "Accept": "application/json"}
    return base, headers


def _inputs():
    return json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))


def _list(v):
    if v in (None, ""):
        return []
    if isinstance(v, list):
        return [str(x).strip().upper() for x in v if str(x).strip()]
    return [p.strip().upper() for p in str(v).split(",") if p.strip()]


def _rel_seconds(text):
    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 _to_ms(v):
    if not v:
        return None
    s = str(v).strip()
    secs = _rel_seconds(s)
    if secs is not None:
        return int((time.time() - secs) * 1000)
    if s.isdigit():
        n = int(s)
        return n if n >= 10 ** 12 else n * 1000
    try:
        dt = datetime.fromisoformat(s.replace("Z", "+00:00"))
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return int(dt.timestamp() * 1000)
    except Exception:
        return None


def _graphql(query, variables):
    base, headers = _cfg()
    data = json.dumps({"query": query, "variables": variables}).encode("utf-8")
    req = urllib.request.Request(base + GRAPHQL, data=data, headers=headers, method="POST")
    with urllib.request.urlopen(req, timeout=90) as r:
        raw = r.read()
        resp = json.loads(raw) if raw else {}
    if resp.get("errors"):
        msgs = [e.get("message", str(e)) for e in resp["errors"]]
        raise RuntimeError("GraphQL: " + "; ".join(msgs))
    return resp


def _filters(inp):
    f = []
    for field, key in [("severity", "severity"), ("status", "status"), ("classification", "classification"),
                       ("attack_surface", "attackSurfaces"), ("analyst_verdict", "analystVerdict")]:
        vals = _list(inp.get(field))
        if vals:
            f.append({"fieldId": key, "stringIn": {"values": vals}})
    os_types = _list(inp.get("os_type"))
    if os_types:
        f.append({"fieldId": "asset.osType", "stringIn": {"values": os_types}})
    if inp.get("asset_id"):
        f.append({"fieldId": "asset.id", "stringEqual": {"value": inp["asset_id"]}})
    if inp.get("external_id"):
        f.append({"fieldId": "externalId", "stringEqual": {"value": inp["external_id"]}})
    if inp.get("search_text"):
        f.append({"fieldId": "name", "match": {"values": [inp["search_text"]]}})
    if str(inp.get("unassigned_only") or "").lower() in ("1", "true", "yes"):
        f.append({"fieldId": "assignee.userId", "isNegated": True, "stringEqual": {"value": None}})
    if str(inp.get("unmitigated_only") or "").lower() in ("1", "true", "yes"):
        f.append({"fieldId": "result", "stringEqual": {"value": "UNMITIGATED"}})
    start = _to_ms(inp.get("start_time"))
    end = _to_ms(inp.get("end_time"))
    if start or end:
        rng = {}
        if start:
            rng["start"] = start
            rng["startInclusive"] = True
        if end:
            rng["end"] = end
            rng["endInclusive"] = True
        f.append({"fieldId": "detectedAt", "dateTimeRange": rng})
    return f


def run():
    inp = _inputs()
    limit = min(int(inp.get("limit") or 50), 200)
    variables = {"first": limit, "filters": _filters(inp), "sorts": [{"by": "detectedAt", "order": "ASC"}]}
    resp = _graphql(QUERY, variables)
    alerts = resp.get("data", {}).get("alerts", {})
    nodes = [e.get("node", {}) for e in alerts.get("edges", [])]
    print(json.dumps({"data": nodes, "totalCount": alerts.get("totalCount", 0)}))


try:
    run()
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)