#!/usr/bin/env python3
import os
from pathlib import Path
from datetime import datetime, timedelta, timezone

import requests
import urllib3
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

BASE_DIR = Path(__file__).resolve().parent
load_dotenv(BASE_DIR / ".env")

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

WAZUH_URL = os.getenv("WAZUH_URL", "https://192.168.1.52:55000").rstrip("/")
WAZUH_USER = os.getenv("WAZUH_USER", "wazuh-wui")
WAZUH_PASSWORD = os.getenv("WAZUH_PASSWORD")

INDEXER_URL = os.getenv("INDEXER_URL", "https://192.168.1.52:9200").rstrip("/")
INDEXER_USER = os.getenv("INDEXER_USER")
INDEXER_PASSWORD = os.getenv("INDEXER_PASSWORD")
INDEXER_ALERT_INDEX = os.getenv("INDEXER_ALERT_INDEX", "wazuh-alerts-*")
INDEXER_VULN_INDEX = os.getenv("INDEXER_VULN_INDEX", "wazuh-states-vulnerabilities-*")

mcp = FastMCP("wazuh")

_token = None


def utc_since(hours: int) -> str:
    return (datetime.now(timezone.utc) - timedelta(hours=hours)).isoformat()


def get_token() -> str:
    global _token

    if _token:
        return _token

    if not WAZUH_PASSWORD:
        raise RuntimeError("WAZUH_PASSWORD manquant dans .env")

    r = requests.post(
        f"{WAZUH_URL}/security/user/authenticate?raw=true",
        auth=(WAZUH_USER, WAZUH_PASSWORD),
        verify=False,
        timeout=15,
    )
    r.raise_for_status()
    _token = r.text.strip()
    return _token


def wazuh_get(endpoint: str) -> dict:
    token = get_token()
    r = requests.get(
        f"{WAZUH_URL}{endpoint}",
        headers={"Authorization": f"Bearer {token}"},
        verify=False,
        timeout=30,
    )
    r.raise_for_status()
    return r.json()


def indexer_search(index: str, body: dict) -> dict:
    if not INDEXER_USER or not INDEXER_PASSWORD:
        raise RuntimeError("INDEXER_USER ou INDEXER_PASSWORD manquant dans .env")

    r = requests.post(
        f"{INDEXER_URL}/{index}/_search",
        auth=(INDEXER_USER, INDEXER_PASSWORD),
        json=body,
        verify=False,
        timeout=30,
    )
    r.raise_for_status()
    return r.json()


def simple_time_query(hours: int, size: int = 20, extra_filters=None) -> dict:
    filters = [
        {"range": {"@timestamp": {"gte": utc_since(hours)}}}
    ]

    if extra_filters:
        filters.extend(extra_filters)

    return {
        "size": size,
        "sort": [{"@timestamp": {"order": "desc"}}],
        "query": {
            "bool": {
                "filter": filters
            }
        }
    }


# -------------------------
# Wazuh Manager API
# -------------------------

@mcp.tool()
def manager_status() -> dict:
    """Retourne l'état du Wazuh Manager."""
    return wazuh_get("/manager/status")


@mcp.tool()
def manager_info() -> dict:
    """Retourne les informations du Wazuh Manager."""
    return wazuh_get("/manager/info")


@mcp.tool()
def list_agents(limit: int = 100) -> dict:
    """Liste les agents Wazuh."""
    return wazuh_get(f"/agents?limit={limit}")


@mcp.tool()
def list_active_agents(limit: int = 100) -> dict:
    """Liste les agents actifs."""
    return wazuh_get(f"/agents?status=active&limit={limit}")


@mcp.tool()
def list_disconnected_agents(limit: int = 100) -> dict:
    """Liste les agents déconnectés."""
    return wazuh_get(f"/agents?status=disconnected&limit={limit}")


@mcp.tool()
def list_never_connected_agents(limit: int = 100) -> dict:
    """Liste les agents jamais connectés."""
    return wazuh_get(f"/agents?status=never_connected&limit={limit}")


@mcp.tool()
def get_agent(agent_id: str) -> dict:
    """Retourne le détail d'un agent Wazuh par son ID."""
    return wazuh_get(f"/agents?agents_list={agent_id}")


@mcp.tool()
def get_agent_config(agent_id: str) -> dict:
    """Retourne la configuration effective d'un agent."""
    return wazuh_get(f"/agents/{agent_id}/config")


@mcp.tool()
def search_agents(query: str, limit: int = 100) -> dict:
    """Recherche des agents par nom, IP ou information associée."""
    return wazuh_get(f"/agents?q={query}&limit={limit}")


# -------------------------
# Wazuh Indexer / Alerts
# -------------------------

@mcp.tool()
def list_alerts(hours: int = 24, limit: int = 20) -> dict:
    """Liste les dernières alertes Wazuh sur une période donnée."""
    body = simple_time_query(hours=hours, size=limit)
    return indexer_search(INDEXER_ALERT_INDEX, body)


@mcp.tool()
def list_critical_alerts(hours: int = 24, limit: int = 20) -> dict:
    """Liste les alertes critiques, niveau 12 ou supérieur."""
    body = simple_time_query(
        hours=hours,
        size=limit,
        extra_filters=[
            {"range": {"rule.level": {"gte": 12}}}
        ],
    )
    return indexer_search(INDEXER_ALERT_INDEX, body)


@mcp.tool()
def top_alerting_agents(hours: int = 24, limit: int = 10) -> dict:
    """Retourne les agents générant le plus d'alertes."""
    body = {
        "size": 0,
        "query": {
            "range": {
                "@timestamp": {
                    "gte": utc_since(hours)
                }
            }
        },
        "aggs": {
            "agents": {
                "terms": {
                    "field": "agent.name",
                    "size": limit
                }
            }
        }
    }
    return indexer_search(INDEXER_ALERT_INDEX, body)


@mcp.tool()
def search_alerts(query: str, hours: int = 24, limit: int = 20) -> dict:
    """Recherche libre dans les alertes Wazuh."""
    body = {
        "size": limit,
        "sort": [{"@timestamp": {"order": "desc"}}],
        "query": {
            "bool": {
                "filter": [
                    {"range": {"@timestamp": {"gte": utc_since(hours)}}}
                ],
                "must": [
                    {
                        "query_string": {
                            "query": query
                        }
                    }
                ]
            }
        }
    }
    return indexer_search(INDEXER_ALERT_INDEX, body)


# -------------------------
# Auth / SSH / sudo
# -------------------------

@mcp.tool()
def get_auth_failures(hours: int = 24, limit: int = 30) -> dict:
    """Liste les échecs d'authentification récents."""
    body = simple_time_query(
        hours=hours,
        size=limit,
        extra_filters=[
            {
                "bool": {
                    "should": [
                        {"match_phrase": {"rule.groups": "authentication_failed"}},
                        {"match_phrase": {"rule.groups": "sshd"}},
                        {"match_phrase": {"full_log": "Failed password"}},
                        {"match_phrase": {"full_log": "authentication failure"}}
                    ],
                    "minimum_should_match": 1
                }
            }
        ],
    )
    return indexer_search(INDEXER_ALERT_INDEX, body)


@mcp.tool()
def get_sudo_events(hours: int = 24, limit: int = 30) -> dict:
    """Liste les événements sudo récents."""
    body = simple_time_query(
        hours=hours,
        size=limit,
        extra_filters=[
            {
                "bool": {
                    "should": [
                        {"match_phrase": {"rule.groups": "sudo"}},
                        {"match_phrase": {"full_log": "sudo"}}
                    ],
                    "minimum_should_match": 1
                }
            }
        ],
    )
    return indexer_search(INDEXER_ALERT_INDEX, body)


# -------------------------
# FIM
# -------------------------

@mcp.tool()
def get_fim_events(hours: int = 24, limit: int = 30) -> dict:
    """Liste les événements d'intégrité fichiers FIM/syscheck."""
    body = simple_time_query(
        hours=hours,
        size=limit,
        extra_filters=[
            {
                "bool": {
                    "should": [
                        {"match_phrase": {"rule.groups": "syscheck"}},
                        {"exists": {"field": "syscheck.path"}}
                    ],
                    "minimum_should_match": 1
                }
            }
        ],
    )
    return indexer_search(INDEXER_ALERT_INDEX, body)


@mcp.tool()
def get_fim_events_for_path(path: str, hours: int = 24, limit: int = 30) -> dict:
    """Liste les modifications FIM pour un chemin donné."""
    body = simple_time_query(
        hours=hours,
        size=limit,
        extra_filters=[
            {"wildcard": {"syscheck.path": f"*{path}*"}}
        ],
    )
    return indexer_search(INDEXER_ALERT_INDEX, body)


# -------------------------
# SCA / compliance
# -------------------------

@mcp.tool()
def get_sca_failures(hours: int = 24, limit: int = 30) -> dict:
    """Liste les contrôles SCA en échec."""
    body = simple_time_query(
        hours=hours,
        size=limit,
        extra_filters=[
            {
                "bool": {
                    "should": [
                        {"match_phrase": {"rule.groups": "sca"}},
                        {"exists": {"field": "data.sca"}}
                    ],
                    "minimum_should_match": 1
                }
            },
            {
                "bool": {
                    "should": [
                        {"match_phrase": {"data.sca.check.result": "failed"}},
                        {"match_phrase": {"full_log": "failed"}}
                    ],
                    "minimum_should_match": 1
                }
            }
        ],
    )
    return indexer_search(INDEXER_ALERT_INDEX, body)


# -------------------------
# Vulnerabilities
# -------------------------

@mcp.tool()
def list_vulnerabilities(limit: int = 50) -> dict:
    """Liste les vulnérabilités présentes dans l'index Wazuh."""
    body = {
        "size": limit,
        "sort": [{"@timestamp": {"order": "desc"}}],
        "query": {
            "match_all": {}
        }
    }
    return indexer_search(INDEXER_VULN_INDEX, body)


@mcp.tool()
def list_critical_vulnerabilities(limit: int = 50) -> dict:
    """Liste les vulnérabilités critiques."""
    body = {
        "size": limit,
        "sort": [{"@timestamp": {"order": "desc"}}],
        "query": {
            "bool": {
                "should": [
                    {"match_phrase": {"vulnerability.severity": "Critical"}},
                    {"match_phrase": {"vulnerability.severity": "CRITICAL"}},
                    {"range": {"vulnerability.score.base": {"gte": 9}}}
                ],
                "minimum_should_match": 1
            }
        }
    }
    return indexer_search(INDEXER_VULN_INDEX, body)


@mcp.tool()
def get_agent_vulnerabilities(agent_name: str, limit: int = 50) -> dict:
    """Liste les vulnérabilités d'un agent par nom."""
    body = {
        "size": limit,
        "sort": [{"@timestamp": {"order": "desc"}}],
        "query": {
            "bool": {
                "filter": [
                    {"match_phrase": {"agent.name": agent_name}}
                ]
            }
        }
    }
    return indexer_search(INDEXER_VULN_INDEX, body)


@mcp.tool()
def top_vulnerable_agents(limit: int = 10) -> dict:
    """Retourne les agents avec le plus de vulnérabilités."""
    body = {
        "size": 0,
        "query": {"match_all": {}},
        "aggs": {
            "agents": {
                "terms": {
                    "field": "agent.name",
                    "size": limit
                }
            }
        }
    }
    return indexer_search(INDEXER_VULN_INDEX, body)


# -------------------------
# Résumé sécurité
# -------------------------

@mcp.tool()
def security_summary(hours: int = 24) -> dict:
    """Produit un résumé sécurité global sur les dernières heures."""
    summary = {}

    try:
        summary["manager_status"] = manager_status()
    except Exception as e:
        summary["manager_status_error"] = str(e)

    try:
        summary["disconnected_agents"] = list_disconnected_agents(limit=100)
    except Exception as e:
        summary["disconnected_agents_error"] = str(e)

    try:
        summary["critical_alerts"] = list_critical_alerts(hours=hours, limit=10)
    except Exception as e:
        summary["critical_alerts_error"] = str(e)

    try:
        summary["top_alerting_agents"] = top_alerting_agents(hours=hours, limit=10)
    except Exception as e:
        summary["top_alerting_agents_error"] = str(e)

    try:
        summary["auth_failures"] = get_auth_failures(hours=hours, limit=10)
    except Exception as e:
        summary["auth_failures_error"] = str(e)

    try:
        summary["critical_vulnerabilities"] = list_critical_vulnerabilities(limit=10)
    except Exception as e:
        summary["critical_vulnerabilities_error"] = str(e)

    try:
        summary["fim_events"] = get_fim_events(hours=hours, limit=10)
    except Exception as e:
        summary["fim_events_error"] = str(e)

    return summary

@mcp.tool()
def vulnerability_report(limit_hosts: int = 100) -> dict:
    """Rapport des vulnérabilités par host et par criticité."""
    body = {
        "size": 0,
        "query": {
            "match_all": {}
        },
        "aggs": {
            "hosts": {
                "terms": {
                    "field": "agent.name",
                    "size": limit_hosts
                },
                "aggs": {
                    "severity": {
                        "terms": {
                            "field": "vulnerability.severity",
                            "size": 10
                        }
                    }
                }
            }
        }
    }

    result = indexer_search(INDEXER_VULN_INDEX, body)

    report = []
    buckets = result.get("aggregations", {}).get("hosts", {}).get("buckets", [])

    for host in buckets:
        counts = {
            "host": host.get("key"),
            "critical": 0,
            "high": 0,
            "medium": 0,
            "low": 0,
            "total": host.get("doc_count", 0),
        }

        for sev in host.get("severity", {}).get("buckets", []):
            severity = str(sev.get("key", "")).lower()
            count = sev.get("doc_count", 0)

            if severity == "critical":
                counts["critical"] = count
            elif severity == "high":
                counts["high"] = count
            elif severity == "medium":
                counts["medium"] = count
            elif severity == "low":
                counts["low"] = count

        report.append(counts)

    report.sort(
        key=lambda x: (
            x["critical"],
            x["high"],
            x["medium"],
            x["low"],
            x["total"],
        ),
        reverse=True,
    )

    return {
        "summary": {
            "hosts": len(report),
            "total_vulnerabilities": sum(h["total"] for h in report),
            "critical": sum(h["critical"] for h in report),
            "high": sum(h["high"] for h in report),
            "medium": sum(h["medium"] for h in report),
            "low": sum(h["low"] for h in report),
        },
        "hosts": report,
    }

@mcp.tool()
def cve_report_by_host(limit_hosts: int = 100, limit_cves_per_host: int = 50) -> dict:
    """
    Rapport détaillé des CVE par host.

    Exemple de sortie :
    host1:
      CVE-2024-XXXX : faille sur le paquet openssl - Critical
      CVE-2025-YYYY : faille sur le paquet curl - High
    """

    body = {
        "size": 0,
        "query": {
            "match_all": {}
        },
        "aggs": {
            "hosts": {
                "terms": {
                    "field": "agent.name",
                    "size": limit_hosts
                },
                "aggs": {
                    "cves": {
                        "terms": {
                            "field": "vulnerability.id",
                            "size": limit_cves_per_host
                        },
                        "aggs": {
                            "sample": {
                                "top_hits": {
                                    "size": 1,
                                    "_source": {
                                        "includes": [
                                            "agent.name",
                                            "vulnerability.id",
                                            "vulnerability.title",
                                            "vulnerability.description",
                                            "vulnerability.severity",
                                            "vulnerability.score.base",
                                            "package.name",
                                            "package.version",
                                            "package.architecture"
                                        ]
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    result = indexer_search(INDEXER_VULN_INDEX, body)

    report = []
    hosts = result.get("aggregations", {}).get("hosts", {}).get("buckets", [])

    severity_order = {
        "critical": 4,
        "high": 3,
        "medium": 2,
        "low": 1,
        "unknown": 0,
    }

    for host_bucket in hosts:
        host_name = host_bucket.get("key")
        cve_entries = []

        cve_buckets = host_bucket.get("cves", {}).get("buckets", [])

        for cve_bucket in cve_buckets:
            sample_hits = (
                cve_bucket
                .get("sample", {})
                .get("hits", {})
                .get("hits", [])
            )

            if not sample_hits:
                continue

            src = sample_hits[0].get("_source", {})

            vulnerability = src.get("vulnerability", {})
            package = src.get("package", {})

            cve_id = vulnerability.get("id", cve_bucket.get("key", "CVE inconnue"))
            title = vulnerability.get("title") or vulnerability.get("description") or "faille détectée"
            severity = vulnerability.get("severity", "Unknown")
            score = vulnerability.get("score", {}).get("base")
            package_name = package.get("name", "paquet inconnu")
            package_version = package.get("version", "version inconnue")
            package_arch = package.get("architecture", "")

            cve_entries.append({
                "cve": cve_id,
                "package": package_name,
                "version": package_version,
                "architecture": package_arch,
                "severity": severity,
                "score": score,
                "title": title,
                "text": f"{cve_id} : faille sur le paquet {package_name} {package_version} - {severity}",
            })

        cve_entries.sort(
            key=lambda x: (
                severity_order.get(str(x["severity"]).lower(), 0),
                x["score"] or 0,
                x["cve"]
            ),
            reverse=True
        )

        report.append({
            "host": host_name,
            "total_cves": len(cve_entries),
            "cves": cve_entries,
        })

    report.sort(
        key=lambda x: x["total_cves"],
        reverse=True
    )

    return {
        "summary": {
            "hosts": len(report),
            "total_cves": sum(h["total_cves"] for h in report),
        },
        "hosts": report,
    }

@mcp.tool()
def cve_report_by_severity(severity: str = "Critical", limit_hosts: int = 100, limit_cves_per_host: int = 50) -> dict:
    """
    Rapport des CVE par host filtré par criticité.
    severity peut être : Critical, High, Medium, Low
    """

    body = {
        "size": 0,
        "query": {
            "match_phrase": {
                "vulnerability.severity": severity
            }
        },
        "aggs": {
            "hosts": {
                "terms": {
                    "field": "agent.name",
                    "size": limit_hosts
                },
                "aggs": {
                    "cves": {
                        "terms": {
                            "field": "vulnerability.id",
                            "size": limit_cves_per_host
                        },
                        "aggs": {
                            "sample": {
                                "top_hits": {
                                    "size": 1,
                                    "_source": {
                                        "includes": [
                                            "agent.name",
                                            "vulnerability.id",
                                            "vulnerability.title",
                                            "vulnerability.description",
                                            "vulnerability.severity",
                                            "vulnerability.score.base",
                                            "package.name",
                                            "package.version",
                                            "package.architecture"
                                        ]
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    result = indexer_search(INDEXER_VULN_INDEX, body)

    report = []
    hosts = result.get("aggregations", {}).get("hosts", {}).get("buckets", [])

    for host_bucket in hosts:
        host_name = host_bucket.get("key")
        cve_entries = []

        for cve_bucket in host_bucket.get("cves", {}).get("buckets", []):
            sample_hits = (
                cve_bucket
                .get("sample", {})
                .get("hits", {})
                .get("hits", [])
            )

            if not sample_hits:
                continue

            src = sample_hits[0].get("_source", {})
            vulnerability = src.get("vulnerability", {})
            package = src.get("package", {})

            cve_id = vulnerability.get("id", cve_bucket.get("key", "CVE inconnue"))
            package_name = package.get("name", "paquet inconnu")
            package_version = package.get("version", "version inconnue")
            title = vulnerability.get("title") or vulnerability.get("description") or "faille détectée"
            score = vulnerability.get("score", {}).get("base")

            cve_entries.append({
                "cve": cve_id,
                "package": package_name,
                "version": package_version,
                "severity": vulnerability.get("severity", severity),
                "score": score,
                "title": title,
                "text": f"{cve_id} : faille sur le paquet {package_name} {package_version} - {severity}",
            })

        report.append({
            "host": host_name,
            "severity": severity,
            "total_cves": len(cve_entries),
            "cves": cve_entries,
        })

    report.sort(key=lambda x: x["total_cves"], reverse=True)

    return {
        "summary": {
            "severity": severity,
            "hosts": len(report),
            "total_cves": sum(h["total_cves"] for h in report),
        },
        "hosts": report,
    }
@mcp.tool()
def sca_summary_by_host(hours: int = 168, limit_hosts: int = 100) -> dict:
    """
    Résumé de l'évaluation des configurations SCA par host.
    Donne le nombre de checks passés, échoués, non applicables, etc.
    """
    body = {
        "size": 0,
        "query": {
            "bool": {
                "filter": [
                    {"range": {"@timestamp": {"gte": utc_since(hours)}}},
                    {
                        "bool": {
                            "should": [
                                {"match_phrase": {"rule.groups": "sca"}},
                                {"exists": {"field": "data.sca.check.result"}}
                            ],
                            "minimum_should_match": 1
                        }
                    }
                ]
            }
        },
        "aggs": {
            "hosts": {
                "terms": {
                    "field": "agent.name",
                    "size": limit_hosts
                },
                "aggs": {
                    "results": {
                        "terms": {
                            "field": "data.sca.check.result",
                            "size": 10
                        }
                    },
                    "policies": {
                        "terms": {
                            "field": "data.sca.policy",
                            "size": 10
                        }
                    }
                }
            }
        }
    }

    result = indexer_search(INDEXER_ALERT_INDEX, body)
    hosts = result.get("aggregations", {}).get("hosts", {}).get("buckets", [])

    report = []

    for host in hosts:
        counts = {
            "host": host.get("key"),
            "passed": 0,
            "failed": 0,
            "not_applicable": 0,
            "unknown": 0,
            "total": host.get("doc_count", 0),
            "policies": []
        }

        for bucket in host.get("results", {}).get("buckets", []):
            key = str(bucket.get("key", "")).lower()
            count = bucket.get("doc_count", 0)

            if key in ["passed", "pass"]:
                counts["passed"] += count
            elif key in ["failed", "fail"]:
                counts["failed"] += count
            elif key in ["not applicable", "not_applicable", "n/a"]:
                counts["not_applicable"] += count
            else:
                counts["unknown"] += count

        counts["policies"] = [
            {
                "policy": p.get("key"),
                "count": p.get("doc_count", 0)
            }
            for p in host.get("policies", {}).get("buckets", [])
        ]

        effective_total = counts["passed"] + counts["failed"]
        counts["compliance_percent"] = round(
            (counts["passed"] / effective_total) * 100, 2
        ) if effective_total > 0 else None

        report.append(counts)

    report.sort(key=lambda x: (x["failed"], x["total"]), reverse=True)

    return {
        "summary": {
            "hosts": len(report),
            "total_checks": sum(h["total"] for h in report),
            "passed": sum(h["passed"] for h in report),
            "failed": sum(h["failed"] for h in report),
            "not_applicable": sum(h["not_applicable"] for h in report),
        },
        "hosts": report
    }

@mcp.tool()
def sca_failed_checks_by_host(hours: int = 168, limit_hosts: int = 50, limit_checks_per_host: int = 20) -> dict:
    """
    Liste les contrôles SCA échoués par host.
    Utile pour demander : résumé de l'évaluation des configurations.
    """
    body = {
        "size": 0,
        "query": {
            "bool": {
                "filter": [
                    {"range": {"@timestamp": {"gte": utc_since(hours)}}},
                    {"match_phrase": {"data.sca.check.result": "failed"}}
                ]
            }
        },
        "aggs": {
            "hosts": {
                "terms": {
                    "field": "agent.name",
                    "size": limit_hosts
                },
                "aggs": {
                    "failed_checks": {
                        "terms": {
                            "field": "data.sca.check.title",
                            "size": limit_checks_per_host
                        },
                        "aggs": {
                            "sample": {
                                "top_hits": {
                                    "size": 1,
                                    "_source": {
                                        "includes": [
                                            "@timestamp",
                                            "agent.name",
                                            "data.sca.policy",
                                            "data.sca.check.id",
                                            "data.sca.check.title",
                                            "data.sca.check.result",
                                            "data.sca.check.description",
                                            "data.sca.check.rationale",
                                            "data.sca.check.remediation",
                                            "rule.level",
                                            "rule.description"
                                        ]
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    result = indexer_search(INDEXER_ALERT_INDEX, body)
    hosts = result.get("aggregations", {}).get("hosts", {}).get("buckets", [])

    report = []

    for host_bucket in hosts:
        checks = []

        for check_bucket in host_bucket.get("failed_checks", {}).get("buckets", []):
            hits = (
                check_bucket
                .get("sample", {})
                .get("hits", {})
                .get("hits", [])
            )

            if not hits:
                continue

            src = hits[0].get("_source", {})
            sca = src.get("data", {}).get("sca", {})
            check = sca.get("check", {})

            checks.append({
                "id": check.get("id"),
                "title": check.get("title", check_bucket.get("key")),
                "policy": sca.get("policy"),
                "description": check.get("description"),
                "rationale": check.get("rationale"),
                "remediation": check.get("remediation"),
                "rule_level": src.get("rule", {}).get("level"),
                "count": check_bucket.get("doc_count", 0),
            })

        report.append({
            "host": host_bucket.get("key"),
            "failed_total": host_bucket.get("doc_count", 0),
            "failed_checks": checks,
        })

    report.sort(key=lambda x: x["failed_total"], reverse=True)

    return {
        "summary": {
            "hosts": len(report),
            "total_failed_checks": sum(h["failed_total"] for h in report),
        },
        "hosts": report
    }


if __name__ == "__main__":
    mcp.run()