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"
# /alerts/ops/rrns takes a list; batching keeps one URL and one body reasonable.
RRN_BATCH = 100
# The server gives a command 90 seconds. Stop walking investigations well before
# that and report the truncation, rather than have the whole run killed and every
# alert in the window lost — a partial fetch that says so beats a silent zero.
DEFAULT_BUDGET_S = 60


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)

    # 1) The investigations in the window. This is the unit that decides what is
    #    worth ingesting: an alert Rapid7 never folded into one is not returned by
    #    this command at all, which is the whole point of it.
    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]

    # 2) Their alert RRNs. One call per investigation — there is no bulk form —
    #    so this is the step the time budget guards.
    order, by_rrn, truncated = [], {}, False
    for inv in investigations:
        if time.monotonic() - started > budget or len(order) >= max_alerts:
            truncated = True
            break
        inv_rrn = inv.get("rrn") or ""
        if not inv_rrn:
            continue
        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:
            # One unreadable investigation must not cost the whole window. The
            # alert simply does not appear; the counts below say how many did.
            continue
        for a in (page.get("data") or []):
            rrn = a.get("id") or a.get("rrn")
            if rrn and rrn not in by_rrn:
                order.append(rrn)
                by_rrn[rrn] = inv
            if len(order) >= max_alerts:
                truncated = True
                break

    # 3) The alerts themselves, in batches. This is where the detection content
    #    lives — rule_keys_of_interest, mitre_tcodes, disposition — none of which
    #    the investigation or its alert list carries.
    field_ids = _list(inp.get("field_ids"))
    alerts = []
    for i in range(0, len(order), RRN_BATCH):
        batch = order[i:i + RRN_BATCH]
        body = {"rrns": batch}
        if field_ids:
            body["field_ids"] = field_ids
        out = request("POST", AT + "alerts/ops/rrns", body=body)
        alerts.extend(out.get("alerts") or [])

    # The investigation each alert came from, carried alongside so a mapper can
    # read the container's own title, priority and disposition without a second
    # feed. Alert fields are left untouched: `investigation_rrn` already exists
    # on the alert and nothing here may shadow it.
    if str(inp.get("include_investigation") or "true").strip().lower() in ("1", "true", "yes"):
        for a in alerts:
            inv = by_rrn.get(a.get("rrn"))
            if inv:
                a["investigation"] = inv

    print(json.dumps({
        "alerts": alerts,
        "metadata": {
            "investigations_read": len(investigations),
            "alert_rrns_collected": len(order),
            "alerts_returned": len(alerts),
            "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)