Anthropic-Cybersecurity-Skills/skills/detecting-insider-threat-with-ueba/scripts/agent.py

#!/usr/bin/env python3
"""UEBA Insider Threat Agent - builds behavioral baselines and scores anomalies using Elasticsearch."""

import json
import argparse
import logging
import math
import os
from collections import defaultdict
from datetime import datetime, timedelta
from elasticsearch import Elasticsearch

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


def connect_es(hosts, api_key=None):
    """Connect to Elasticsearch cluster."""
    kwargs = {"hosts": hosts, "verify_certs": False, "request_timeout": 30}
    if api_key:
        kwargs["api_key"] = api_key
    return Elasticsearch(**kwargs)


def build_user_baseline(es, index, user_field, hours=720):
    """Build 30-day behavioral baseline per user using ES aggregations."""
    since = (datetime.utcnow() - timedelta(hours=hours)).isoformat()
    query = {
        "size": 0,
        "query": {"range": {"@timestamp": {"gte": since}}},
        "aggs": {
            "users": {
                "terms": {"field": user_field, "size": 5000},
                "aggs": {
                    "login_hours": {"histogram": {"field": "hour_of_day", "interval": 1}},
                    "daily_events": {"date_histogram": {"field": "@timestamp", "calendar_interval": "day"}},
                    "unique_hosts": {"cardinality": {"field": "host.name"}},
                    "data_volume": {"sum": {"field": "bytes_transferred"}},
                    "unique_apps": {"cardinality": {"field": "application.name"}},
                }
            }
        }
    }
    result = es.search(index=index, body=query)
    baselines = {}
    for bucket in result["aggregations"]["users"]["buckets"]:
        user = bucket["key"]
        daily_counts = [d["doc_count"] for d in bucket["daily_events"]["buckets"]]
        avg_daily = sum(daily_counts) / max(len(daily_counts), 1)
        std_daily = math.sqrt(sum((x - avg_daily) ** 2 for x in daily_counts) / max(len(daily_counts), 1))
        baselines[user] = {
            "avg_daily_events": round(avg_daily, 1),
            "std_daily_events": round(std_daily, 1),
            "unique_hosts": bucket["unique_hosts"]["value"],
            "total_data_volume": bucket["data_volume"]["value"],
            "total_events": bucket["doc_count"],
        }
    return baselines


def score_current_activity(es, index, user_field, baselines, hours=24):
    """Score current activity against baselines to find anomalies."""
    since = (datetime.utcnow() - timedelta(hours=hours)).isoformat()
    query = {
        "size": 0,
        "query": {"range": {"@timestamp": {"gte": since}}},
        "aggs": {
            "users": {
                "terms": {"field": user_field, "size": 5000},
                "aggs": {
                    "unique_hosts": {"cardinality": {"field": "host.name"}},
                    "data_volume": {"sum": {"field": "bytes_transferred"}},
                    "unique_apps": {"cardinality": {"field": "application.name"}},
                }
            }
        }
    }
    result = es.search(index=index, body=query)
    anomalies = []
    for bucket in result["aggregations"]["users"]["buckets"]:
        user = bucket["key"]
        baseline = baselines.get(user)
        if not baseline:
            anomalies.append({
                "user": user, "indicator": "new_user",
                "severity": "medium", "detail": "No baseline exists for this user",
                "risk_score": 50,
            })
            continue
        current_events = bucket["doc_count"]
        avg = baseline["avg_daily_events"]
        std = baseline["std_daily_events"]
        z_score = (current_events - avg) / max(std, 1)
        if z_score > 3:
            anomalies.append({
                "user": user, "indicator": "activity_spike",
                "severity": "high", "z_score": round(z_score, 2),
                "current": current_events, "baseline_avg": avg,
                "risk_score": min(int(z_score * 15), 100),
                "detail": f"Event count {current_events} is {z_score:.1f} std devs above baseline",
            })
        current_hosts = bucket["unique_hosts"]["value"]
        if current_hosts > baseline["unique_hosts"] * 2:
            anomalies.append({
                "user": user, "indicator": "new_host_access",
                "severity": "high",
                "current_hosts": current_hosts,
                "baseline_hosts": baseline["unique_hosts"],
                "risk_score": 70,
                "detail": f"Accessed {current_hosts} hosts vs baseline {baseline['unique_hosts']}",
            })
        current_volume = bucket["data_volume"]["value"]
        daily_avg_volume = baseline["total_data_volume"] / 30
        if current_volume > daily_avg_volume * 5 and current_volume > 100_000_000:
            anomalies.append({
                "user": user, "indicator": "data_exfiltration",
                "severity": "critical",
                "current_bytes": current_volume,
                "baseline_daily_avg": round(daily_avg_volume),
                "risk_score": 90,
                "detail": f"Transferred {current_volume / 1e6:.0f}MB vs daily avg {daily_avg_volume / 1e6:.1f}MB",
            })
    return sorted(anomalies, key=lambda x: x.get("risk_score", 0), reverse=True)


def peer_group_analysis(baselines, peer_groups):
    """Compare user activity against peer group averages."""
    findings = []
    group_stats = defaultdict(list)
    for user, baseline in baselines.items():
        group = peer_groups.get(user, "default")
        group_stats[group].append(baseline["avg_daily_events"])
    group_avgs = {g: sum(v) / len(v) for g, v in group_stats.items()}
    for user, baseline in baselines.items():
        group = peer_groups.get(user, "default")
        group_avg = group_avgs.get(group, 0)
        if group_avg > 0 and baseline["avg_daily_events"] > group_avg * 3:
            findings.append({
                "user": user, "peer_group": group,
                "user_avg": baseline["avg_daily_events"],
                "group_avg": round(group_avg, 1),
                "deviation_factor": round(baseline["avg_daily_events"] / group_avg, 1),
                "severity": "medium",
            })
    return findings


def generate_report(anomalies, peer_findings, baselines):
    critical = sum(1 for a in anomalies if a.get("severity") == "critical")
    return {
        "timestamp": datetime.utcnow().isoformat(),
        "users_baselined": len(baselines),
        "anomalies_detected": len(anomalies),
        "critical_anomalies": critical,
        "top_risk_users": anomalies[:15],
        "peer_group_outliers": peer_findings[:10],
        "risk_level": "critical" if critical > 0 else "high" if anomalies else "low",
    }


def main():
    parser = argparse.ArgumentParser(description="UEBA Insider Threat Detection Agent")
    parser.add_argument("--es-hosts", default=os.environ.get("ES_HOSTS", "https://localhost:9200"), help="Elasticsearch hosts")
    parser.add_argument("--api-key", help="Elasticsearch API key")
    parser.add_argument("--index", default="logs-*", help="Log index pattern")
    parser.add_argument("--user-field", default="user.name", help="User identity field")
    parser.add_argument("--peer-groups", help="JSON file mapping users to peer groups")
    parser.add_argument("--lookback", type=int, default=24, help="Anomaly lookback hours")
    parser.add_argument("--output", default="ueba_insider_threat_report.json")
    args = parser.parse_args()

    es = connect_es(args.es_hosts.split(","), args.api_key)
    baselines = build_user_baseline(es, args.index, args.user_field)
    anomalies = score_current_activity(es, args.index, args.user_field, baselines, args.lookback)
    peer_groups = {}
    if args.peer_groups:
        with open(args.peer_groups) as f:
            peer_groups = json.load(f)
    peer_findings = peer_group_analysis(baselines, peer_groups)
    report = generate_report(anomalies, peer_findings, baselines)
    with open(args.output, "w") as f:
        json.dump(report, f, indent=2, default=str)
    logger.info("UEBA: %d users baselined, %d anomalies (%d critical)",
                len(baselines), len(anomalies), report["critical_anomalies"])
    print(json.dumps(report, indent=2, default=str))


if __name__ == "__main__":
    main()