import json, os, sys, time, urllib.request, urllib.parse, urllib.error
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone

# Riposte stops a command after 60 s. Event searches get what is left of this
# budget once the alerts are listed; the margin covers container start-up.
STARTED = time.monotonic()
BUDGET = 45


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


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


def left():
    return BUDGET - (time.monotonic() - STARTED)


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 += "?" + 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=max(5, min(30, left()))) as r:
        raw = r.read()
        return json.loads(raw) if raw else {}


def nest(event):
    """Sekoia returns events flat, ECS names as dotted keys ("source.ip").
    Rebuild the tree so a mapper reads events[0].source.ip; a key whose prefix
    is already a plain value stays flat rather than overwrite it."""
    out = {}
    for key in sorted(event, key=lambda k: k.count(".")):
        node, parts = out, key.split(".")
        for p in parts[:-1]:
            if p not in node:
                node[p] = {}
            if not isinstance(node[p], dict):
                node = None
                break
            node = node[p]
        if node is None or parts[-1] in node:
            out[key] = event[key]
        else:
            node[parts[-1]] = event[key]
    return out


def start_search(alert):
    # The alert's events are the ones tagged with its short id. Sekoia's own
    # XSOAR integration searches from first_seen_at less a minute up to now:
    # events are indexed slightly after the time they carry.
    try:
        first = datetime.fromisoformat(str(alert.get("first_seen_at")).replace("Z", "+00:00"))
        earliest = (first - timedelta(minutes=1)).astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    except ValueError:
        earliest = "-7d"
    job = request("POST", "/v1/sic/conf/events/search/jobs", body={
        "earliest_time": earliest,
        "latest_time": "now",
        "term": "alert_short_ids:" + str(alert.get("short_id")),
    })
    return job.get("uuid")


def collect(job_uuid, max_events):
    path = "/v1/sic/conf/events/search/jobs/" + urllib.parse.quote(job_uuid, safe="")
    while left() > 2:
        if request("GET", path).get("status") == 2:
            res = request("GET", path + "/events", params={"limit": max_events})
            return [nest(e) for e in res.get("items") or []], res.get("total")
        time.sleep(1)
    return None, None


def attach_events(alerts, max_events):
    """Put each alert's events on it, as events (nested) and events_total. A
    search not done within the budget leaves events_pending and the job uuid
    (events_results_query reads it later): the alert still goes out, because
    the rule's watermark moves on with this run and would not offer it again."""
    def one(alert):
        try:
            job = start_search(alert)
            events, total = collect(job, max_events) if job else (None, None)
        except Exception as e:
            alert["events_error"] = str(e)
            return
        if events is None:
            alert["events_pending"] = True
            alert["events_job_uuid"] = job
        else:
            alert["events"] = events
            alert["events_total"] = total
    with ThreadPoolExecutor(max_workers=10) as pool:
        list(pool.map(one, alerts))


def _range(value):
    # A bare date is the lower bound ("<date>,now"); a full range is passed through.
    v = str(value)
    return v if "," in v else (v + ",now")


def run():
    inp = _inputs()
    params = {
        "offset": inp.get("offset") or 0,
        "direction": inp.get("direction") or "asc",
        "sort": inp.get("sort_by") or "created_at",
    }
    if inp.get("limit"):
        params["limit"] = inp["limit"]
    if inp.get("status"):
        params["match[status_name]"] = inp["status"]
    if inp.get("created_at"):
        params["date[created_at]"] = _range(inp["created_at"])
    if inp.get("updated_at"):
        params["date[updated_at]"] = _range(inp["updated_at"])
    if inp.get("urgency"):
        params["range[urgency]"] = inp["urgency"]
    if inp.get("alerts_type"):
        params["match[type_value]"] = inp["alerts_type"]
    out = request("GET", "/v1/sic/alerts", params=params)
    if str(inp.get("with_events") or "").lower() in ("true", "1", "yes"):
        attach_events(out.get("items") or [], int(inp.get("max_events") or 10))
    print(json.dumps(out))


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)