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

# One call for an agent investigating a SentinelOne alert: what the alert is,
# the host behind it, and what else the console holds around it. An operator
# model spent ten calls rebuilding this and still lost the host — it took the
# UAM asset id (a 26-character token) for an agent id, and the raw agent list
# it fell back to ran past what a tool result can carry. Every step here is
# best effort and summarised: a step that fails is named in "errors" and the
# rest of the answer still comes back.

UUID_RE = re.compile(r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")
SHA1_RE = re.compile(r"^[0-9a-fA-F]{40}$")

# List-shape fields (the ones get_unified_alerts reads, known to every console)
# for the alert itself, fetched by an id filter.
ALERT_FIELDS = """
  id name description severity status result analystVerdict classification
  confidenceLevel detectedAt firstSeenAt lastSeenAt storylineId fileName fileHash
  analytics { name category }
  detectionSource { product vendor engine }
  asset { id name agentUuid agentVersion osType osVersion lastLoggedInUser connectivityToConsole }
  detectionTime {
    asset { domain ipV4 lastLoggedInUser osName }
    attacker { host ip }
    scope { accountName groupName siteName }
    targetUser { name domain }
  }
  process { cmdLine parentName file { certSubject md5 name path sha1 sha256 } }
"""

# What only the per-alert view carries.
DETAIL_FIELDS = """
  indicators { message attacks { tactic { name } technique { name uid } } }
  observables { name value typeName }
  relatedAlerts { id name }
  process { username }
"""

SIBLING_FIELDS = "id name severity status detectedAt classification"


def request(method, url, headers, body=None):
    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=60) as resp:
        raw = resp.read()
        return json.loads(raw) if raw else {}


def trim(v, n):
    s = "" if v is None else str(v)
    return s if len(s) <= n else s[:n] + "…"


def compact(d):
    """Drop the empty values, so the answer says only what the console said."""
    if isinstance(d, dict):
        out = {k: compact(v) for k, v in d.items()}
        return {k: v for k, v in out.items() if v not in (None, "", [], {})}
    if isinstance(d, list):
        return [compact(x) for x in d if x not in (None, "", [], {})]
    return d


def parse_time(s):
    try:
        dt = datetime.datetime.fromisoformat(str(s).replace("Z", "+00:00"))
    except (TypeError, ValueError):
        return None
    return dt if dt.tzinfo else dt.replace(tzinfo=datetime.timezone.utc)


def iso(dt):
    return dt.astimezone(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.000Z")


def group_count(items, key, extra=None, top=10):
    """Count items per key value, most frequent first, with the latest time seen."""
    groups = {}
    for it in items:
        k = key(it) or "?"
        g = groups.setdefault(k, {"value": k, "count": 0, "last": ""})
        g["count"] += 1
        t = extra(it) if extra else ""
        if t and t > g["last"]:
            g["last"] = t
    out = sorted(groups.values(), key=lambda g: -g["count"])[:top]
    if not extra:
        for g in out:
            g.pop("last")
    return out


class Console:
    def __init__(self, base, headers):
        self.base, self.headers, self.errors = base, headers, []

    def get(self, step, path, qs=None, missing=None):
        """missing, when given, is what a 404 means for this step: an answer,
        not a failed check."""
        url = self.base + path + ("?" + urllib.parse.urlencode({k: v for k, v in (qs or {}).items() if v not in (None, "")}) if qs else "")
        try:
            return request("GET", url, self.headers)
        except urllib.error.HTTPError as e:
            if e.code == 404 and missing is not None:
                return missing
            self.errors.append({"step": step, "error": "HTTP %d" % e.code,
                                "detail": trim(e.read().decode("utf-8", "replace"), 300)})
        except Exception as e:
            self.errors.append({"step": step, "error": trim(e, 300)})
        return None

    def gql(self, step, query):
        try:
            resp = request("POST", self.base + "/unifiedalerts/graphql", self.headers, {"query": query})
        except urllib.error.HTTPError as e:
            self.errors.append({"step": step, "error": "HTTP %d" % e.code,
                                "detail": trim(e.read().decode("utf-8", "replace"), 300)})
            return None
        except Exception as e:
            self.errors.append({"step": step, "error": trim(e, 300)})
            return None
        if resp.get("errors"):
            self.errors.append({"step": step, "error": "GraphQL error",
                                "detail": trim(json.dumps(resp["errors"]), 300)})
        return resp.get("data") or {}


def alerts_query(filters, fields, first):
    # Both spellings of the sort argument exist across console versions; none
    # is needed here, so the query carries no sort at all.
    return """query Alerts { alerts(first: %d viewType: ALL filters: [ %s ]) {
  totalCount edges { node { %s } } } }""" % (first, " ".join(filters), fields)


def summarize_agent(a):
    ips = []
    for nic in a.get("networkInterfaces") or []:
        for ip in nic.get("inet") or []:
            if ip and ip not in ips:
                ips.append(ip)
    return compact({
        "computer_name": a.get("computerName"),
        "agent_id": a.get("id"),
        "os": a.get("osName"),
        "os_revision": a.get("osRevision"),
        "machine_type": a.get("machineType"),
        "domain": a.get("domain"),
        "ips": ips[:6],
        "external_ip": a.get("externalIp"),
        "last_logged_in_user": a.get("lastLoggedInUserName"),
        "agent_version": a.get("agentVersion"),
        "is_active": a.get("isActive"),
        "network_status": a.get("networkStatus"),
        "infected": a.get("infected"),
        "active_threats": a.get("activeThreats"),
        "last_active": a.get("lastActiveDate"),
        "site": a.get("siteName"),
        "group": a.get("groupName"),
        "account": a.get("accountName"),
    })


def find_agent(con, agent_uuid=None, agent_id=None, name=None):
    """The host's agent, by the most precise key the alert gave."""
    for step, qs in (("agent_by_id", {"ids": agent_id}),
                     ("agent_by_uuid", {"uuids": agent_uuid}),
                     ("agent_by_name", {"computerName": name})):
        if not list(qs.values())[0]:
            continue
        resp = con.get(step, "/agents", dict(qs, limit=5))
        data = (resp or {}).get("data") or []
        if data:
            return data[0], step
    return None, None


def threats_summary(con, step, qs):
    resp = con.get(step, "/threats", dict(qs, limit=100, sortBy="createdAt", sortOrder="desc"))
    if resp is None:
        return None
    data = resp.get("data") or []
    total = ((resp.get("pagination") or {}).get("totalItems"))
    return compact({
        "total": total if total is not None else len(data),
        "by_name": group_count(data, lambda t: (t.get("threatInfo") or {}).get("threatName"),
                               lambda t: (t.get("threatInfo") or {}).get("createdAt") or ""),
        "by_host": group_count(data, lambda t: (t.get("agentRealtimeInfo") or {}).get("agentComputerName")),
    })


def hash_checks(con, sha1, site_id):
    out = {"sha1": sha1}
    verdict = con.get("hash_verdict", "/hashes/%s/verdict" % sha1,
                      missing={"data": {"verdict": "unknown to SentinelOne (404)"}})
    if verdict is not None:
        out["reputation"] = ((verdict.get("data") or {}).get("verdict")) or "unknown"
    blocked = []
    for scope_qs in ({"tenant": "true"}, {"siteIds": site_id} if site_id else None):
        if not scope_qs:
            continue
        resp = con.get("blocklist", "/restrictions", dict(scope_qs, type="black_hash", value__contains=sha1, limit=10))
        for r in (resp or {}).get("data") or []:
            blocked.append(compact({"scope": r.get("scopeName") or r.get("scope"), "description": trim(r.get("description"), 120)}))
    out["blocklisted"] = bool(blocked)
    if blocked:
        out["blocklist_entries"] = blocked[:5]
    out["threats_fleet_wide"] = threats_summary(con, "threats_by_hash", {"contentHashes": sha1})
    return compact(out)


def investigate_uam(con, alert_id, days):
    data = con.gql("alert", alerts_query(['{ fieldId: "id" stringEqual: { value: %s } }' % json.dumps(alert_id)],
                                         ALERT_FIELDS, 1))
    if data is None:
        # The console could not be asked: that is a failed call, not an alert
        # that does not exist.
        raise RuntimeError("the alert could not be read: %s" % json.dumps(con.errors[-1]))
    edges = (data.get("alerts") or {}).get("edges") or []
    if not edges:
        return {"found": False, "alert_id": alert_id,
                "hint": "No unified alert has this id. Copy the id whole from get_incident_context (raw id), or pass a threat_id for a threat."}
    a = edges[0]["node"]
    det = (con.gql("alert_detail", "query D { alert(id: %s) { %s } }" % (json.dumps(alert_id), DETAIL_FIELDS)) or {}).get("alert") or {}

    proc = a.get("process") or {}
    f = proc.get("file") or {}
    dt_asset = (a.get("detectionTime") or {}).get("asset") or {}
    scope = (a.get("detectionTime") or {}).get("scope") or {}
    indicators = []
    for ind in (det.get("indicators") or [])[:10]:
        attacks = ["%s / %s (%s)" % ((x.get("tactic") or {}).get("name"), (x.get("technique") or {}).get("name"),
                                     (x.get("technique") or {}).get("uid")) for x in ind.get("attacks") or []]
        indicators.append(compact({"message": trim(ind.get("message"), 220), "mitre": attacks[:3]}))
    observables, seen = [], set()
    for ob in (det.get("observables") or []):
        key = (ob.get("typeName") or ob.get("name"), ob.get("value"))
        if key in seen or not ob.get("value"):
            continue
        seen.add(key)
        observables.append("%s: %s" % (key[0], trim(key[1], 160)))
    out = {
        "kind": "unified alert",
        "alert": compact({
            "id": a.get("id"), "name": a.get("name"), "description": trim(a.get("description"), 400),
            "severity": a.get("severity"), "classification": a.get("classification"),
            "confidence": a.get("confidenceLevel"), "status": a.get("status"),
            "analyst_verdict": a.get("analystVerdict"), "mitigation_result": a.get("result"),
            "detected_at": a.get("detectedAt"), "first_seen": a.get("firstSeenAt"), "last_seen": a.get("lastSeenAt"),
            "rule": (a.get("analytics") or {}).get("name"), "rule_category": (a.get("analytics") or {}).get("category"),
            "detected_by": " ".join(x for x in ((a.get("detectionSource") or {}).get("product"),
                                                (a.get("detectionSource") or {}).get("engine")) if x),
            "storyline_id": a.get("storylineId"),
            "attacker": (a.get("detectionTime") or {}).get("attacker"),
            "target_user": (a.get("detectionTime") or {}).get("targetUser"),
            "scope": " / ".join(x for x in (scope.get("accountName"), scope.get("siteName"), scope.get("groupName")) if x),
        }),
        "process": compact({
            "command_line": trim(proc.get("cmdLine"), 600), "parent": proc.get("parentName"),
            "user": det.get("process", {}).get("username") if det.get("process") else None,
            "file": f.get("name") or a.get("fileName"), "path": f.get("path"),
            "sha1": f.get("sha1"), "sha256": f.get("sha256"), "md5": f.get("md5"),
            "signer": f.get("certSubject"),
        }),
        "indicators": indicators,
        "observables": observables[:25],
        "related_alerts": [trim(r.get("name"), 120) for r in (det.get("relatedAlerts") or [])[:10]],
    }

    asset = a.get("asset") or {}
    agent, how = find_agent(con, agent_uuid=asset.get("agentUuid"), name=asset.get("name"))
    out["host"] = summarize_agent(agent) if agent else compact({
        "computer_name": asset.get("name"), "os": dt_asset.get("osName") or asset.get("osVersion"),
        "ip_at_detection": dt_asset.get("ipV4"), "domain": dt_asset.get("domain"),
        "last_logged_in_user": dt_asset.get("lastLoggedInUser") or asset.get("lastLoggedInUser"),
        "agent_version": asset.get("agentVersion"), "connectivity": asset.get("connectivityToConsole"),
        "note": "no SentinelOne agent found for this asset (not an endpoint, or uninstalled): fields from the alert",
    })
    if agent:
        out["host"]["found_by"] = how
        if dt_asset.get("ipV4"):
            out["host"]["ip_at_detection"] = dt_asset.get("ipV4")

    anchor = parse_time(a.get("detectedAt")) or datetime.datetime.now(datetime.timezone.utc)
    start, end = anchor - datetime.timedelta(days=days), anchor + datetime.timedelta(days=1)
    window = "%s to %s (%d days before the alert's detection to 1 day after)" % (iso(start), iso(end), days)
    elsewhere = {"window": window}
    if asset.get("id"):
        sib = con.gql("same_host_alerts", alerts_query([
            '{ fieldId: "asset.id" stringEqual: { value: %s } }' % json.dumps(asset["id"]),
            '{ fieldId: "detectedAt" dateTimeRange: { start: %d end: %d } }' % (start.timestamp() * 1000, end.timestamp() * 1000),
        ], SIBLING_FIELDS, 100))
        alerts = (sib or {}).get("alerts") or {}
        edges_all = [e["node"] for e in alerts.get("edges") or [] if e.get("node")]
        nodes = [n for n in edges_all if n.get("id") != alert_id]
        if sib is not None:
            # Counted from what was read, never from totalCount minus one: a
            # real run reported "1 other alert" with none to name, because the
            # console's total did not hold this alert. The console's own count
            # is only given when more alerts exist than were read.
            other = {"count": len(nodes),
                     "by_name": group_count(nodes, lambda n: n.get("name"), lambda n: n.get("detectedAt") or "")}
            total = alerts.get("totalCount")
            if total is not None and total > len(edges_all):
                other["console_total_including_this_alert"] = total
                other["note"] = "more alerts than the 100 read"
            elsewhere["same_host_other_alerts"] = compact(other) if nodes else "none"
    if agent:
        elsewhere["same_host_threats"] = threats_summary(con, "threats_on_host", {
            "agentIds": agent.get("id"), "createdAt__gt": iso(start), "createdAt__lt": iso(end)})
    sha1 = f.get("sha1") or (a.get("fileHash") if SHA1_RE.match(str(a.get("fileHash") or "")) else None)
    if sha1:
        out["file_checks"] = hash_checks(con, sha1, agent.get("siteId") if agent else None)
    else:
        out["file_checks"] = "no SHA1 on this alert: reputation, blocklist and fleet-wide checks skipped"
    out["elsewhere"] = compact(elsewhere)
    return out


def investigate_threat(con, threat_id, days):
    resp = con.get("threat", "/threats", {"ids": threat_id})
    if resp is None:
        raise RuntimeError("the threat could not be read: %s" % json.dumps(con.errors[-1]))
    data = (resp or {}).get("data") or []
    if not data:
        return {"found": False, "threat_id": threat_id,
                "hint": "No threat has this id. Copy the id whole from get_incident_context (threatInfo.threatId), or pass alert_id for a unified alert."}
    t = data[0]
    ti, rt, di = t.get("threatInfo") or {}, t.get("agentRealtimeInfo") or {}, t.get("agentDetectionInfo") or {}
    out = {
        "kind": "threat",
        "threat": compact({
            "id": ti.get("threatId"), "name": ti.get("threatName"), "classification": ti.get("classification"),
            "confidence": ti.get("confidenceLevel"), "analyst_verdict": ti.get("analystVerdict"),
            "incident_status": ti.get("incidentStatus"), "mitigation_status": ti.get("mitigationStatus"),
            "detection_type": ti.get("detectionType"), "engines": ti.get("engines"),
            "created_at": ti.get("createdAt"), "identified_at": ti.get("identifiedAt"),
            "storyline": ti.get("storyline"),
        }),
        "process": compact({
            "file_path": ti.get("filePath"), "sha1": ti.get("sha1"), "sha256": ti.get("sha256"), "md5": ti.get("md5"),
            "signer": ti.get("publisherName"), "verified": ti.get("fileVerificationType"),
            "originating_process": ti.get("originatorProcess"), "process_user": ti.get("processUser"),
            "command_line": trim(ti.get("maliciousProcessArguments"), 600),
            "initiated_by": ti.get("initiatedBy"),
        }),
    }
    notes = con.get("threat_notes", "/threats/%s/notes" % threat_id, {"limit": 10})
    out["notes"] = [compact({"by": n.get("creator"), "at": n.get("createdAt"), "text": trim(n.get("text"), 300)})
                    for n in (notes or {}).get("data") or []]
    acts = con.get("threat_timeline", "/activities", {"threatIds": threat_id, "limit": 30, "sortBy": "createdAt", "sortOrder": "asc"})
    out["timeline"] = ["%s %s" % (x.get("createdAt"), trim(x.get("primaryDescription"), 200))
                       for x in (acts or {}).get("data") or []][:20]

    agent, how = find_agent(con, agent_id=rt.get("agentId"), name=rt.get("agentComputerName"))
    out["host"] = summarize_agent(agent) if agent else compact({
        "computer_name": rt.get("agentComputerName"), "os": rt.get("agentOsName"),
        "agent_version": rt.get("agentVersion"), "network_status": rt.get("networkStatus"),
        "note": "agent not found in the console: fields from the threat",
    })
    if agent:
        out["host"]["found_by"] = how
    out["host"]["user_at_detection"] = di.get("agentLastLoggedInUserName")
    out["host"]["ip_at_detection"] = di.get("agentIpV4")
    out["host"] = compact(out["host"])

    anchor = parse_time(ti.get("createdAt")) or datetime.datetime.now(datetime.timezone.utc)
    start, end = anchor - datetime.timedelta(days=days), anchor + datetime.timedelta(days=1)
    elsewhere = {"window": "%s to %s (%d days before the threat to 1 day after)" % (iso(start), iso(end), days)}
    if agent:
        elsewhere["same_host_threats"] = threats_summary(con, "threats_on_host", {
            "agentIds": agent.get("id"), "createdAt__gt": iso(start), "createdAt__lt": iso(end)})
    out["elsewhere"] = compact(elsewhere)
    if ti.get("sha1"):
        out["file_checks"] = hash_checks(con, ti["sha1"], agent.get("siteId") if agent else None)
    return out


def main():
    secrets = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
    inputs = json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))
    base = secrets.get("url", "").rstrip("/") + "/web/api/v2.1"
    con = Console(base, {
        "Authorization": "ApiToken " + secrets.get("api_token", ""),
        "Accept": "application/json",
        "Content-Type": "application/json",
    })
    alert_id = str(inputs.get("alert_id") or "").strip()
    threat_id = str(inputs.get("threat_id") or "").strip()
    days = max(1, min(int(inputs.get("days_before") or 7), 30))
    # A long number in alert_id is a threat id given in the wrong field.
    if alert_id.isdigit() and not threat_id:
        threat_id, alert_id = alert_id, ""
    if threat_id:
        if not threat_id.isdigit():
            raise ValueError("threat_id is a long number (about 19 digits); for a UUID use alert_id")
        out = investigate_threat(con, threat_id, days)
    elif alert_id:
        if not UUID_RE.match(alert_id):
            raise ValueError("alert_id is the unified alert's UUID (8-4-4-4-12 hex), copied whole; got %r (asset.id and agentUuid are not alert ids)" % alert_id)
        out = investigate_uam(con, alert_id, days)
    else:
        raise ValueError("give alert_id (unified alert UUID) or threat_id (threat number)")
    if con.errors:
        out["errors"] = con.errors[:8]
    print(json.dumps(out, ensure_ascii=False, separators=(",", ":")))


try:
    main()
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)