riposte-marketplace/integrations/rapid7-insightidr/scripts/get_investigations_full.py

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

# The Alert Triage API lives under /idr/at, a base of its own — not /idr/v1 or
# /idr/v2 like the investigation endpoints. Same host, same API key.
AT = "idr/at/"


def _cfg():
    s = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
    region = str(s.get("region") or "us").strip().lower()
    base = "https://" + region + ".api.insight.rapid7.com/"
    headers = {"X-Api-Key": s.get("api_key", ""), "Content-Type": "application/json", "Accept": "application/json"}
    multi = str(s.get("is_multi_customer") or "").strip().lower() in ("1", "true", "yes")
    return base, headers, multi


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() for x in v if str(x).strip()]
    return [p.strip() for p in str(v).split(",") if p.strip()]


def request(method, path, params=None, body=None):
    base, headers, _ = _cfg()
    url = base + path.lstrip("/")
    if params:
        clean = {k: v for k, v in params.items() if v not in (None, "")}
        if clean:
            url += ("&" if "?" in url else "?") + urllib.parse.urlencode(clean, doseq=True)
    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=90) as r:
        raw = r.read()
        return json.loads(raw) if raw else {}

import time
from datetime import datetime, timedelta, timezone

ISO = "%Y-%m-%dT%H:%M:%SZ"
RRN_BATCH = 100
DEFAULT_BUDGET_S = 60

# idr/v2/investigations answers with 13 fields, three of them null: a title, a
# priority, a status, an assignee and timestamps. No host, no account, no
# process, no MITRE code — and not one of the alerts the investigation
# aggregates. An incident built from it indexes zero entities, so it cannot be
# pivoted on, cannot be correlated, and gives an analyst nothing to read.
#
# The detection content is two hops away: the investigation's alert list gives
# RRNs, and the Alert Triage API turns those into alerts carrying
# rule_keys_of_interest (the host, the process, its command line, its parent,
# the hashes) and rule.mitre_tcodes. This command walks both hops and returns
# the investigation with its alerts inlined, so the INVESTIGATION stays the unit
# that becomes an incident — results_path = data, dedup on rrn, exactly as the
# plain investigation feed. Only the command changes.


def _alert_rrn(a):
    """The alert's RRN, or None when this alert has none.

    The investigation alert list is not consistent about where the RRN sits.
    Some entries carry it in `id` (rrn:alerts:eu:<org>:alert:1:<hex>); others put
    a plain UUID there and the RRN, if any, in `rrn`. `rrn` is therefore read
    first, and anything that is not RRN-shaped is dropped rather than sent on.

    That last part is not defensive habit. /alerts/ops/rrns validates the whole
    list and rejects the entire request when one entry is malformed, so a single
    UUID slipping through costs every alert in the batch — which is exactly how
    this failed: HTTP 400, eight invalid RRNs, zero alerts fetched, the whole run
    lost. Dropping them here means the investigation still arrives, with its
    unretrievable alerts counted rather than silently missing.
    """
    for v in (a.get("rrn"), a.get("id")):
        if isinstance(v, str) and v.startswith("rrn:"):
            return v
    return None


def run():
    _, _, multi = _cfg()
    inp = _inputs()
    limit = int(inp.get("limit") or 50)
    max_alerts = int(inp.get("max_alerts") or 500)
    budget = float(inp.get("time_budget_seconds") or DEFAULT_BUDGET_S)
    started = time.monotonic()

    start_time = inp.get("start_time")
    if not start_time:
        days = int(inp.get("default_days") or 1)
        start_time = (datetime.now(timezone.utc) - timedelta(days=days)).strftime(ISO)

    inv_page = request("GET", "idr/v2/investigations", params={
        "index": inp.get("index") or "0",
        "size": inp.get("page_size") or limit,
        "statuses": inp.get("statuses"),
        "priorities": inp.get("priorities"),
        "start_time": start_time,
        "end_time": inp.get("end_time"),
        "sort_field": "Created time",
        "sort_direction": "ASC",
        "multi-customer": "true" if multi else None,
    })
    investigations = (inv_page.get("data") or [])[:limit]

    # Alert RRNs per investigation. One call each — there is no bulk form — so
    # this is the step that grows with the window, and the one the time budget
    # guards. An investigation whose alerts cannot be read keeps its place in the
    # output with an alerts_error: losing the whole run over one of them, or
    # dropping it silently, would both be worse.
    rrns_by_inv, order, truncated, skipped = {}, [], False, 0
    for inv in investigations:
        inv_rrn = inv.get("rrn") or ""
        if not inv_rrn:
            continue
        if time.monotonic() - started > budget or len(order) >= max_alerts:
            truncated = True
            break
        try:
            page = request("GET", "idr/v2/investigations/" + urllib.parse.quote(inv_rrn, safe="") + "/alerts",
                           params={"multi-customer": "true" if multi else None})
        except urllib.error.HTTPError as e:
            inv["alerts_error"] = "HTTP " + str(e.code)
            continue
        got = []
        for a in (page.get("data") or []):
            rrn = _alert_rrn(a)
            if rrn is None:
                skipped += 1
                continue
            got.append(rrn)
            if rrn not in order:
                order.append(rrn)
        rrns_by_inv[inv_rrn] = got

    # The alerts themselves, batched, then indexed so each investigation can be
    # handed its own.
    field_ids = _list(inp.get("field_ids"))
    by_rrn = {}
    for i in range(0, len(order), RRN_BATCH):
        body = {"rrns": order[i:i + RRN_BATCH]}
        if field_ids:
            body["field_ids"] = field_ids
        out = request("POST", AT + "alerts/ops/rrns", body=body)
        for a in (out.get("alerts") or []):
            if a.get("rrn"):
                by_rrn[a["rrn"]] = a

    fetched = 0
    for inv in investigations:
        rrns = rrns_by_inv.get(inv.get("rrn") or "", [])
        alerts = [by_rrn[r] for r in rrns if r in by_rrn]
        inv["alerts"] = alerts
        inv["alert_count"] = len(rrns)
        # Say it when the alert list is a sample: alert_count is what the
        # investigation actually has, len(alerts) what this run retrieved.
        if len(alerts) < len(rrns):
            inv["alerts_truncated"] = True
        fetched += len(alerts)
        # Lift a few first-alert values to the top level. The incident-field
        # mapper resolves dotted paths only — it cannot index an array — so
        # `alerts[0].ai_suggested_disposition` in a field_mappings entry reads as
        # nothing, silently, and the custom field arrives empty with no error
        # anywhere. The OCSF entries are evaluated by full JSONata and do not
        # need this; only field_mappings does.
        if alerts:
            a = alerts[0]
            log = (a.get("log_details") or [{}])[0]
            inv.update({
                "first_alert_rrn": a.get("rrn"),
                "first_alert_type": a.get("type"),
                "first_alert_source": a.get("external_source"),
                "first_alert_status": a.get("status"),
                "first_alert_disposition": a.get("disposition"),
                "first_alert_ai_disposition": a.get("ai_suggested_disposition"),
                "first_alert_event_type": a.get("triggering_event_type"),
                "first_alert_rule_rrn": (a.get("rule") or {}).get("rrn"),
                "first_alert_log_id": log.get("log_id"),
                "first_alert_logset_id": log.get("logset_id"),
                "first_alert_log_entry_id": log.get("log_entry_id"),
            })

    print(json.dumps({
        "data": investigations,
        "metadata": {
            "investigations_returned": len(investigations),
            "alert_rrns_collected": len(order),
            "alerts_fetched": fetched,
            "alert_ids_without_rrn": skipped,
            "truncated": truncated,
            "elapsed_seconds": round(time.monotonic() - started, 1),
        },
    }))


try:
    run()
except urllib.error.HTTPError as e:
    detail = e.read().decode("utf-8", "replace")
    if e.code in (401, 403):
        detail += (" — the Alert Triage API (/idr/at) is only available to Managed Detection and"
                   " Response customers with access to the Alerts experience; a key that works on"
                   " the investigation endpoints can still be refused here")
    print(json.dumps({"error": "HTTP " + str(e.code), "detail": detail}))
    sys.exit(1)
except Exception as e:
    print(json.dumps({"error": str(e)}))
    sys.exit(1)