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 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 = {}, [], False
    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 = a.get("id") or a.get("rrn")
            if rrn:
                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,
            "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)