#!/usr/bin/env python3
"""Minimal Wazuh 5.0 API client: manager API on 55000, indexer API on 9200.

Endpoints, parameters, request bodies and response envelopes are taken from
the Wazuh 5.0 beta OpenAPI specifications:
  manager: https://documentation.wazuh.com/5.0-beta/user-manual/api/reference.html
  indexer: https://documentation.wazuh.com/5.0-beta/user-manual/indexer-api/reference.html

No response body is printed here that was not documented. The script prints
whatever the live API returns.

Requires: python3, requests
    pip install requests
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any, Iterator

import requests

# ------------------------------------------------------------------ TODO
# Supply credentials by environment variable, not by editing this file.
#
#   export WAZUH_API_URL="https://MANAGER_HOST:55000"
#   export WAZUH_API_USER="wazuh"            # default manager API admin user
#   export WAZUH_API_PASSWORD="..."          # TODO
#
#   export WAZUH_INDEXER_URL="https://INDEXER_HOST:9200"
#   export WAZUH_INDEXER_USER="admin"        # default indexer internal user
#   export WAZUH_INDEXER_PASSWORD="..."      # TODO
#
#   # Path to the CA that signed the Wazuh certificates. The installation
#   # assistant issues its own CA; point at it rather than disabling
#   # verification. Only set WAZUH_INSECURE=1 in a lab.
#   export WAZUH_CA_BUNDLE="/etc/wazuh-indexer/certs/root-ca.pem"
# ------------------------------------------------------------------------

DEFAULT_TIMEOUT = 30


def _verify_setting() -> Any:
    if os.environ.get("WAZUH_INSECURE") == "1":
        requests.packages.urllib3.disable_warnings()  # type: ignore[attr-defined]
        return False
    return os.environ.get("WAZUH_CA_BUNDLE", True)


class WazuhManagerAPI:
    """Manager API client. Basic auth once, then a Bearer JWT.

    The specification declares basicAuth as the "Intended method to get a
    token" and jwt as the global security requirement for everything else.
    The token expires after auth_token_exp_timeout seconds (default 900),
    changeable with PUT /security/config, so this client refreshes on 401.
    """

    def __init__(self, base_url: str, user: str, password: str,
                 verify: Any = True, timeout: int = DEFAULT_TIMEOUT) -> None:
        self.base_url = base_url.rstrip("/")
        self.user = user
        self.password = password
        self.verify = verify
        self.timeout = timeout
        self.session = requests.Session()
        self._token: str | None = None

    # -- authentication ---------------------------------------------------

    def login(self) -> None:
        resp = self.session.post(
            f"{self.base_url}/security/user/authenticate",
            auth=(self.user, self.password),
            verify=self.verify,
            timeout=self.timeout,
        )
        resp.raise_for_status()
        self._token = resp.json()["data"]["token"]
        self.session.headers["Authorization"] = f"Bearer {self._token}"

    def request(self, method: str, path: str, **kwargs: Any) -> dict:
        if self._token is None:
            self.login()
        kwargs.setdefault("verify", self.verify)
        kwargs.setdefault("timeout", self.timeout)
        resp = self.session.request(method, f"{self.base_url}{path}", **kwargs)
        if resp.status_code == 401:
            # Token expired mid-run. Refresh once and retry.
            self.login()
            resp = self.session.request(method, f"{self.base_url}{path}", **kwargs)
        if resp.status_code == 429:
            raise RuntimeError(
                "429 Too Many Requests. Use the endpoints' own list parameters "
                "instead of looping single-item calls."
            )
        resp.raise_for_status()
        return resp.json()

    # -- response envelope ------------------------------------------------

    @staticmethod
    def check(payload: dict) -> dict:
        """Raise if the response reports failures.

        A 200 does not mean success. Bulk operations return per-item errors
        inside data.failed_items with a 200 status, and the envelope carries
        its own top-level "error" field.
        """
        if payload.get("error"):
            raise RuntimeError(f"API error {payload['error']}: {payload.get('message')}")
        data = payload.get("data") or {}
        if data.get("total_failed_items"):
            raise RuntimeError(
                f"{data['total_failed_items']} item(s) failed: "
                f"{json.dumps(data.get('failed_items'))}"
            )
        return payload

    # -- agents -----------------------------------------------------------

    def list_agents(self, limit: int = 500, **filters: Any) -> Iterator[dict]:
        """GET /agents, paged.

        Useful filters from the specification: agents_list, select, sort,
        search, status, q, older_than, os.type, os.platform, os.version,
        os.major, os.minor, os.arch, os.name, version, group, name, ip,
        registerIP, distinct.
        """
        offset = 0
        while True:
            params = dict(filters, limit=limit, offset=offset)
            payload = self.check(self.request("GET", "/agents", params=params))
            items = payload["data"]["affected_items"]
            for item in items:
                yield item
            if len(items) < limit:
                return
            offset += limit

    def add_agent(self, name: str, ip: str | None = None) -> dict:
        """POST /agents. Returns data with id and key (base64).

        The key is what lets you provision an endpoint without it talking to
        authd, and is how the 4.x to 5.0 migration imports client.keys.
        """
        body: dict[str, Any] = {"name": name}
        if ip:
            body["ip"] = ip
        payload = self.check(self.request("POST", "/agents", json=body))
        return payload["data"]

    def get_agent_key(self, agent_id: str) -> dict:
        """GET /agents/{agent_id}/key."""
        payload = self.check(
            self.request("GET", f"/agents/{agent_id}/key")
        )
        return payload["data"]

    def assign_group(self, agent_id: str, group_id: str) -> dict:
        """PUT /agents/{agent_id}/group/{group_id}."""
        return self.check(
            self.request("PUT", f"/agents/{agent_id}/group/{group_id}")
        )

    def list_outdated(self) -> list[dict]:
        """GET /agents/outdated."""
        payload = self.check(self.request("GET", "/agents/outdated"))
        return payload["data"]["affected_items"]

    def upgrade_agents(self, agents_list: list[str] | None = None,
                       group: str | None = None,
                       wait_for_complete: bool = False,
                       **extra: Any) -> dict:
        """PUT /agents/upgrade.

        The specification recommends wait_for_complete=true when upgrading
        more than 3000 agents at once, to avoid an API timeout. Filters such
        as group, version and the os.* family also apply, so you can upgrade
        by group rather than enumerating IDs.
        """
        params: dict[str, Any] = dict(extra)
        if agents_list:
            params["agents_list"] = ",".join(agents_list)
        if group:
            params["group"] = group
        if wait_for_complete:
            params["wait_for_complete"] = "true"
        return self.check(self.request("PUT", "/agents/upgrade", params=params))

    # -- groups -----------------------------------------------------------

    def list_groups(self, limit: int = 500) -> list[dict]:
        """GET /groups."""
        payload = self.check(self.request("GET", "/groups", params={"limit": limit}))
        return payload["data"]["affected_items"]

    def create_group(self, group_id: str) -> dict:
        """POST /groups.

        group_id may contain a-z, A-Z, 0-9, '_', '-' and '.', max 128 chars.
        The names '.' and '..' are restricted.
        """
        return self.check(
            self.request("POST", "/groups", json={"group_id": group_id})
        )

    def ensure_group(self, group_id: str) -> bool:
        """Idempotent group creation. Returns True if it created one."""
        existing = {g.get("name") for g in self.list_groups()}
        if group_id in existing:
            return False
        self.create_group(group_id)
        return True

    # -- who am I ---------------------------------------------------------

    def my_policies(self) -> dict:
        """GET /security/users/me/policies.

        Run this first when an endpoint returns 403. Each operation in the
        specification declares the RBAC action it needs, for example
        group:create on POST /groups.
        """
        return self.check(self.request("GET", "/security/users/me/policies"))


class WazuhIndexerAPI:
    """Indexer API client. HTTP Basic or a Bearer JWT, per the specification."""

    def __init__(self, base_url: str, user: str, password: str,
                 verify: Any = True, timeout: int = DEFAULT_TIMEOUT) -> None:
        self.base_url = base_url.rstrip("/")
        self.verify = verify
        self.timeout = timeout
        self.session = requests.Session()
        self.session.auth = (user, password)
        self.session.headers["Content-Type"] = "application/json"

    def search(self, index_pattern: str, body: dict) -> dict:
        resp = self.session.post(
            f"{self.base_url}/{index_pattern}/_search",
            data=json.dumps(body),
            verify=self.verify,
            timeout=self.timeout,
        )
        resp.raise_for_status()
        return resp.json()

    def findings_by_rule(self, days: int = 7, size: int = 50) -> dict:
        """Count findings per rule title over the last N days.

        Fields are Wazuh Common Schema: wazuh.rule.title, wazuh.rule.level
        (a STRING: informational, low, medium, high, critical), @timestamp.
        Findings live in wazuh-findings-v5-<category> data streams.
        """
        body = {
            "size": 0,
            "query": {"range": {"@timestamp": {"gte": f"now-{days}d"}}},
            "aggs": {
                "by_rule": {
                    "terms": {"field": "wazuh.rule.title", "size": size},
                    "aggs": {
                        "by_level": {"terms": {"field": "wazuh.rule.level"}},
                        "by_agent": {"cardinality": {"field": "wazuh.agent.id"}},
                    },
                }
            },
        }
        return self.search("wazuh-findings-v5-*", body)

    def agent_state(self, index: str, agent_id: str, size: int = 100) -> dict:
        """Read a stateful index for one agent.

        These indices replace the manager API endpoints removed in 5.0:
          wazuh-states-sca                  was GET /sca/{agent_id}
          wazuh-states-fim-files            was GET /syscheck/{agent_id}
          wazuh-states-inventory-packages   was GET /syscollector/.../packages
          wazuh-states-vulnerabilities      was GET /vulnerability/{agent_id}
        Full list: wazuh-states-sca, wazuh-states-fim-files,
        wazuh-states-fim-registry-keys, wazuh-states-fim-registry-values,
        wazuh-states-vulnerabilities, and wazuh-states-inventory-{users,
        groups, packages, processes, services, ports, networks, interfaces,
        protocols, hardware, system, hotfixes, browser-extensions}.
        """
        body = {
            "size": size,
            "query": {"term": {"wazuh.agent.id": agent_id}},
        }
        return self.search(index, body)


def _manager_from_env() -> WazuhManagerAPI:
    url = os.environ.get("WAZUH_API_URL")
    user = os.environ.get("WAZUH_API_USER")
    password = os.environ.get("WAZUH_API_PASSWORD")
    if not (url and user and password):
        sys.exit("Set WAZUH_API_URL, WAZUH_API_USER and WAZUH_API_PASSWORD. See the TODO block.")
    return WazuhManagerAPI(url, user, password, verify=_verify_setting())


def _indexer_from_env() -> WazuhIndexerAPI:
    url = os.environ.get("WAZUH_INDEXER_URL")
    user = os.environ.get("WAZUH_INDEXER_USER")
    password = os.environ.get("WAZUH_INDEXER_PASSWORD")
    if not (url and user and password):
        sys.exit("Set WAZUH_INDEXER_URL, WAZUH_INDEXER_USER and WAZUH_INDEXER_PASSWORD.")
    return WazuhIndexerAPI(url, user, password, verify=_verify_setting())


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="cmd", required=True)

    sub.add_parser("agents", help="list agents (paged)")
    sub.add_parser("groups", help="list groups")
    sub.add_parser("outdated", help="list outdated agents")
    sub.add_parser("policies", help="show what this token is allowed to do")

    p = sub.add_parser("onboard", help="idempotent: ensure group, add agent, assign")
    p.add_argument("--group", required=True)
    p.add_argument("--name", required=True)
    p.add_argument("--ip")

    p = sub.add_parser("findings", help="aggregate findings by rule (indexer API)")
    p.add_argument("--days", type=int, default=7)

    p = sub.add_parser("state", help="read a wazuh-states-* index for one agent")
    p.add_argument("--index", required=True)
    p.add_argument("--agent-id", required=True)

    args = parser.parse_args()

    if args.cmd in {"agents", "groups", "outdated", "policies", "onboard"}:
        api = _manager_from_env()

    if args.cmd == "agents":
        for agent in api.list_agents(select="id,name,status,group,version,lastKeepAlive"):
            print(json.dumps(agent))
    elif args.cmd == "groups":
        print(json.dumps(api.list_groups(), indent=2))
    elif args.cmd == "outdated":
        print(json.dumps(api.list_outdated(), indent=2))
    elif args.cmd == "policies":
        print(json.dumps(api.my_policies(), indent=2))
    elif args.cmd == "onboard":
        # Groups MUST exist before an agent enrols into one in Wazuh 5.0.
        created = api.ensure_group(args.group)
        print(f"group {args.group}: {'created' if created else 'already present'}")
        data = api.add_agent(args.name, args.ip)
        print(f"agent {data['id']} added")
        api.assign_group(data["id"], args.group)
        print(f"agent {data['id']} assigned to {args.group}")
        # data["key"] is the base64 enrolment key. Do not log it.
        print("enrolment key retrieved; handle it as a secret")
    elif args.cmd == "findings":
        print(json.dumps(_indexer_from_env().findings_by_rule(days=args.days), indent=2))
    elif args.cmd == "state":
        print(json.dumps(
            _indexer_from_env().agent_state(args.index, args.agent_id), indent=2))


if __name__ == "__main__":
    main()