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

PAGE = 1000           # the relay's largest page
DEFAULT_LIMIT = 10000  # the relay's default max_pending per source


def _cfg():
    s = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
    base = str(s.get("base_url") or "").strip().rstrip("/")
    if not base.startswith("https://"):
        raise ValueError("base_url must be the relay's https:// address")
    return base, {"Authorization": "Bearer " + str(s.get("read_token") or "").strip(), "Accept": "application/json"}


def _get(path, params):
    base, headers = _cfg()
    url = base + path + "?" + urllib.parse.urlencode({k: v for k, v in params.items() if v not in (None, "")})
    req = urllib.request.Request(url, headers=headers, method="GET")
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read() or b"{}")


# InsightIDR basic detection rules post a single field, `payload`, holding the
# alert as a JSON-encoded string — as a JSON body {"payload": "..."} or as a form
# body payload=... depending on the delivery. Both are opened here, in Python:
# event.s is a 64-bit sequence that a float-based JSON reader would round.
OPTED_OUT = "The log line that triggered this alert is not included"


def _payload(m):
    j = m.get("json")
    if isinstance(j, dict) and "payload" in j:
        p = j["payload"]
        return json.loads(p) if isinstance(p, str) else p
    body = m.get("body") or ""
    if isinstance(j, dict):
        return j
    form = urllib.parse.parse_qs(body)
    if "payload" in form:
        return json.loads(form["payload"][0])
    raise ValueError("no Rapid7 payload in this message")


def _iso_ms(t):
    try:
        return datetime.datetime.fromtimestamp(int(t) / 1000, datetime.timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
    except (TypeError, ValueError, OverflowError, OSError):
        return None


def rapid7(m):
    out = {k: m.get(k) for k in ("id", "source", "received_at", "remote_ip", "content_type", "verified", "body")}
    # A message the relay could not tie to a Rapid7 alert still becomes an
    # alert: dropping it would hide a delivery nobody would then look for.
    out["alert_id"] = "relay:" + str(m.get("id"))
    try:
        r = _payload(m)
    except Exception as e:
        out["decode_error"] = str(e)
        return out
    if not isinstance(r, dict):
        out["decode_error"] = "payload is not an object"
        return out
    alert, event, log, host = (r.get(k) if isinstance(r.get(k), dict) else {} for k in ("alert", "event", "log", "host"))
    if r.get("correlation_id"):
        # Rapid7's id for this firing: a redelivery carries it again, a new
        # relay message id does not, so it is the dedup key.
        out["alert_id"] = str(r["correlation_id"])
    out["correlation_id"] = r.get("correlation_id")
    out["rule"] = alert.get("name")
    out["region"] = r.get("region")
    out["event_time"] = _iso_ms(event.get("t"))
    out["event_sequence"] = str(event["s"]) if event.get("s") is not None else None
    line = event.get("m")
    out["log_line_withheld"] = isinstance(line, str) and line.startswith(OPTED_OUT)
    out["log_line"] = None if out["log_line_withheld"] else line
    if isinstance(line, str) and not out["log_line_withheld"]:
        try:
            out["log_event"] = json.loads(line)
        except ValueError:
            pass
    out["log_name"] = log.get("name")
    out["log_key"] = log.get("key")
    # In a basic rule's payload `host.name` is the log set (Asset Authentication),
    # not a machine; host.hostname is the machine, when Rapid7 knows one.
    out["log_set"] = host.get("name")
    out["hostname"] = host.get("hostname")
    out["rapid7"] = r
    return out


def run():
    inp = json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))
    source = str(inp.get("source") or "").strip()
    if not source:
        raise ValueError("source is required: the relay source id, e.g. r7-acme")
    try:
        limit = int(inp.get("limit") or DEFAULT_LIMIT)
    except (TypeError, ValueError):
        limit = DEFAULT_LIMIT
    since = str(inp.get("since") or "").strip()

    # Every page is read within this one run. The next run's watermark is the
    # instant of this one, so a page left for later would never be asked for.
    messages, after = [], ""
    while True:
        page = _get("/relay/v1/messages", {"source": source, "since": since, "after": after, "limit": PAGE})
        messages.extend(page.get("messages") or [])
        if not page.get("more"):
            break
        after = page.get("next") or ""
        if not after:
            break
        if len(messages) >= limit:
            # Failing keeps the watermark where it is; returning a truncated
            # list would move it past what was never read.
            raise RuntimeError(
                "more than %d messages wait on source %s since %s: raise this command's limit "
                "(nothing was ingested; the next run reads the same window)" % (limit, source, since or "the beginning"))
    alerts = [rapid7(m) for m in messages]
    print(json.dumps({"alerts": alerts, "count": len(alerts), "source": source, "since": since}))


try:
    run()
except urllib.error.HTTPError as e:
    detail = e.read().decode("utf-8", "replace")
    out = {"error": "HTTP " + str(e.code), "detail": detail}
    if e.code == 401:
        out["hint"] = "The relay refused the read token, or this Riposte's address is not in RELAY_CONSUMER_ALLOW_IPS."
    elif e.code == 404:
        out["hint"] = ("Unknown source on this relay, or this Riposte's address is not allowed by the proxy in front "
                       "of it (it answers 404 to readers it does not know). Run test_connection.")
    print(json.dumps(out))
    sys.exit(1)
except Exception as e:
    print(json.dumps({"error": str(e)}))
    sys.exit(1)