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

_UNITS = {"second": 1, "minute": 60, "hour": 3600, "day": 86400, "week": 604800, "month": 2592000, "year": 31536000}

# Beyond this a single value can only be milliseconds: 1e11 seconds is the year 5138,
# and 1e11 milliseconds was 1973. Nothing a feed can legitimately ask for sits between.
_MS_FLOOR = 100000000000


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"}
    return base, headers


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


def _rel_ms_range(text):
    if not text:
        return None, None
    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")
    if num is None or not unit:
        return None, None
    now_ms = int(time.time() * 1000)
    return now_ms - num * _UNITS[unit] * 1000, now_ms


def _to_ms(v):
    """Epoch milliseconds, whatever the caller wrote.

    As an ingestion feed this command receives start_time from the scheduler's
    high-water mark, and the mark's format is an operator setting on the rule
    (iso8601, unix, unix_ms). The API only takes milliseconds. Reading all three
    here is what keeps the feed from silently fetching the wrong window — or
    nothing at all — because of a dropdown three screens away.
    """
    if v in (None, ""):
        return None
    s = str(v).strip()
    if s.isdigit():
        n = int(s)
        return n if n >= _MS_FLOOR else n * 1000
    try:
        return int(datetime.datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp() * 1000)
    except Exception:
        return None


def _num(v, default):
    try:
        return float(v) if v not in (None, "") else default
    except (TypeError, ValueError):
        return default


def _iso(ms):
    """The window as it was queried, in the form the agent reports it."""
    return datetime.datetime.fromtimestamp(ms / 1000, tz=datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def request(method, url, params=None):
    base, headers = _cfg()
    full = url if url.startswith("http") else base + url.lstrip("/")
    if params:
        clean = {k: v for k, v in params.items() if v not in (None, "")}
        if clean:
            full += ("&" if "?" in full else "?") + urllib.parse.urlencode(clean, doseq=True)
    req = urllib.request.Request(full, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=90) as r:
        raw = r.read()
        return json.loads(raw) if raw else {}


def _collect(first, limit):
    """Walk the pagination links, stopping at `limit`.

    The walk is bounded because this runs unattended every few minutes: a query
    that matches more than it was meant to would otherwise hold the run open for
    as many pages as the window contains, and the next tick would find the rule
    still in flight and skip itself.
    """
    events = []
    queue = [first]
    while queue and len(events) < limit:
        page = queue.pop(0)
        events.extend(page.get("events", []) or [])
        for link in page.get("links", []) or []:
            href = link.get("href")
            if href:
                queue.append(request("GET", href))
    return events[:limit]


def _identify(events, fallback_log_id, detection_name):
    """Give every event a stable id, and the name of the rule it stands for.

    dedup_key needs a field that is identical across runs for the same log line.
    The scheduler deliberately re-reads the last minute of the previous window on
    every run (so an event written while a poll was in flight is not lost), which
    means the same line WILL come back — anything derived from the fetch time, or
    from the position in the page, would make it a second incident every time.
    log_id + sequence_number is the line's own address in Log Search; the hash is
    the fallback for a source that does not number its entries.

    detection_name is stamped here because the log line has no idea which basic
    detection rule it is being replayed for. Without it an incident is titled by
    its own first characters, and an analyst reads a log line where the source
    console shows an alert name.
    """
    out = []
    for e in events:
        if not isinstance(e, dict):
            continue
        log_id = e.get("log_id") or fallback_log_id or ""
        seq = e.get("sequence_number")
        if seq in (None, ""):
            digest = hashlib.sha1(
                ("%s|%s|%s" % (log_id, e.get("timestamp"), e.get("message"))).encode("utf-8", "replace")
            ).hexdigest()
            e["event_id"] = "%s:%s" % (log_id, digest)
        else:
            e["event_id"] = "%s:%s" % (log_id, seq)
        if not e.get("log_id"):
            e["log_id"] = log_id
        if detection_name:
            e["detection_name"] = detection_name
        out.append(e)
    return out


# What a summary keeps: the values an analyst reads first, per field.
_TOP_VALUES = 8
_MAX_FIELDS = 30
_SAMPLES = 5
# Rapid7's whole vendor payload, kilobytes per event: it stays out of a summary.
_SKIP_FIELDS = {"source_json", "source_data"}


def _fields(e):
    """One Log Search event as flat fields. `message` is a JSON-encoded string
    on parsed sources; nested objects are flattened one level (a.b)."""
    msg = e.get("message")
    data = None
    if isinstance(msg, str):
        try:
            data = json.loads(msg)
        except ValueError:
            data = None
    elif isinstance(msg, dict):
        data = msg
    if not isinstance(data, dict):
        return {"message": msg} if msg not in (None, "") else {}
    flat = {}
    for k, v in data.items():
        if k in _SKIP_FIELDS:
            continue
        if isinstance(v, dict):
            for k2, v2 in v.items():
                if not isinstance(v2, (dict, list)) and v2 not in (None, ""):
                    flat[k + "." + k2] = v2
        elif not isinstance(v, list) and v not in (None, ""):
            flat[k] = v
    return flat


def _parse_only(v):
    """"field=value, field2=value2" -> {field: value}. Applied to the events'
    fields after the search, so narrowing a result needs no query syntax."""
    out = {}
    for part in str(v or "").split(","):
        if "=" in part:
            k, val = part.split("=", 1)
            if k.strip():
                out[k.strip()] = val.strip().strip("\"'")
    return out


def _matches(fields, only):
    """Every condition holds, comparing as text and ignoring case."""
    for k, want in only.items():
        if str(fields.get(k, "")).lower() != want.lower():
            return False
    return True


def _summarize(events, limit):
    """Counts instead of lines: per field, its most frequent values.

    An agent asked who a user logged on to does not need a thousand raw lines —
    on a real run it got 1000 events, could read about ten, and reported the
    accounts and hosts of those ten. Top values per field answer the question
    over the whole window and fit in a model's context.
    """
    counts = {}
    first_ts, last_ts = None, None
    for e in events:
        ts = e.get("timestamp")
        if isinstance(ts, (int, float)):
            first_ts = ts if first_ts is None or ts < first_ts else first_ts
            last_ts = ts if last_ts is None or ts > last_ts else last_ts
        for k, v in _fields(e).items():
            counts.setdefault(k, {})
            key = str(v)[:200]
            counts[k][key] = counts[k].get(key, 0) + 1
    # Fields that vary least first: they are the ones that describe the activity
    # (result, service, destination) rather than identify a line (timestamps, ids).
    ranked = sorted(counts.items(), key=lambda kv: (len(kv[1]), kv[0]))
    top = {}
    for k, vals in ranked[:_MAX_FIELDS]:
        top[k] = {"distinct": len(vals),
                  "top": [[v, n] for v, n in sorted(vals.items(), key=lambda x: -x[1])[:_TOP_VALUES]]}
    return {
        "count": len(events),
        "truncated": len(events) >= limit,
        "first_event": _iso(first_ts) if first_ts is not None else None,
        "last_event": _iso(last_ts) if last_ts is not None else None,
        "fields": top,
        "fields_omitted": max(0, len(counts) - _MAX_FIELDS),
        "samples": [_fields(e) for e in events[:_SAMPLES]],
    }


def _log_sets():
    """Every log set of the account, as (id, name)."""
    out = request("GET", "log_search/management/logsets")
    sets = out.get("logsets", []) if isinstance(out, dict) else []
    return [(x.get("id", ""), x.get("name", "")) for x in sets if isinstance(x, dict)]


def _resolve_name(name):
    """The id of the log set called name: exact match first, ignoring case, then
    the one log set whose name contains it. Anything else is refused with the
    names that do exist, so the next call can pick one rather than guess."""
    sets = _log_sets()
    want = name.lower()
    exact = [i for i, n in sets if n.lower() == want]
    if len(exact) == 1:
        return exact[0]
    partial = [(i, n) for i, n in sets if want in n.lower()]
    if len(partial) == 1:
        return partial[0][0]
    names = sorted(n for _, n in sets)
    if partial:
        raise ValueError("log_set_name %r matches several log sets: %s" % (name, ", ".join(sorted(n for _, n in partial))))
    raise ValueError("no log set is called %r. Log sets: %s" % (name, ", ".join(names)))


NEW_MARKERS = ("new_authentication",)
_TRUE = ("true", "1", "yes")


def _norm_account(v):
    """jdoe, DOMAIN\\jdoe, jdoe@domain.tld -> jdoe."""
    v = str(v or "").strip().lower()
    if "\\" in v:
        v = v.rsplit("\\", 1)[1]
    if "@" in v:
        v = v.split("@", 1)[0]
    return v


def _is_users(fields, user):
    """The event is this user's: an account or user field holds the account.

    The keyword search matches the text anywhere in the line, so it also brings
    other people's logons — on a real run the only "new" event it found was
    another account's. Only the account fields decide."""
    want = _norm_account(user)
    for k, v in fields.items():
        kl = k.lower()
        if ("account" in kl or "user" in kl) and "sid" not in kl and _norm_account(v) == want:
            return True
    return False


def _is_new(fields):
    """Rapid7 marks a first-time authentication either with a boolean field
    (new_authentication: true) or with that word as a value."""
    for k, v in fields.items():
        if "new_authentication" in k.lower() and str(v).strip().lower() in _TRUE:
            return True
        if str(v).strip().lower() in NEW_MARKERS:
            return True
    return False


def run():
    """One call for the question a behaviour (UBA) authentication alert asks:
    which assets did this user authenticate to around the alert, and which of
    them were new. The agent gives a user and the alert's time; the log set,
    the query, the window and the split are done here — on real runs the model
    mangled log set ids, computed the window a day off, put a filter inside the
    LEQL and got a 400, and read ten lines out of a thousand."""
    inp = _inputs()
    user = str(inp.get("user") or "").strip().strip("\"'")
    if not user or '"' in user:
        raise ValueError("user is required: the account name as the alert gives it, e.g. jdoe")
    anchor = _to_ms(inp.get("anchor_time"))
    if anchor is None:
        raise ValueError("anchor_time is required: the alert's first seen time, exactly as the data gives it")
    before = _num(inp.get("hours_before"), 24.0)
    after = _num(inp.get("minutes_after"), 30.0)
    start_time = int(anchor - before * 3600 * 1000)
    end_time = int(anchor + after * 60 * 1000)
    name = str(inp.get("log_set_name") or "Asset Authentication").strip()
    log_set_id = _resolve_name(name)
    try:
        limit = max(int(inp.get("limit") or 10000), 10000)
    except (TypeError, ValueError):
        limit = 10000
    query = 'where("%s")' % user
    first = request("GET", "log_search/query/logsets/" + urllib.parse.quote(log_set_id, safe=""),
                    params={"query": query, "from": start_time, "to": end_time})
    matched = _identify(_collect(first, limit), "", None)
    events = [e for e in matched if _is_users(_fields(e), user)]
    new = [e for e in events if _is_new(_fields(e))]
    print(json.dumps({
        "user": user,
        "log_set": name,
        "query": query,
        "keyword_matches": len(matched),
        "user_matches": len(events),
        "truncated": len(matched) >= limit,
        "window": "%s to %s" % (_iso(start_time), _iso(end_time)),
        "all": _summarize(events, limit),
        "new_authentications": _summarize(new, limit),
        "note": ("all and new_authentications hold only this user's events (an account field equals "
                 "the user): keyword_matches counts every line the text appeared in, other accounts "
                 "included. new_authentications holds the events Rapid7 marked as a first "
                 "authentication: their destination assets are the new assets a behaviour alert counts."),
    }))


try:
    run()
except urllib.error.HTTPError as e:
    out = {"error": "HTTP " + str(e.code), "detail": e.read().decode("utf-8", "replace")}
    if e.code == 404:
        out["hint"] = "No log set with this log_set_id in the configured region. Check the ID against list_log_sets, and the instance region."
    print(json.dumps(out))
    sys.exit(1)
except Exception as e:
    print(json.dumps({"error": str(e)}))
    sys.exit(1)