#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-only
# Copyright (C) 2026 Sudhir Chauhan
"""Read-only Azure Log Analytics collector for Checkmk 2.5."""

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import hashlib
import json
import os
from pathlib import Path
import random
import re
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode
from urllib.request import Request, urlopen

GUID = re.compile(r"^[0-9a-fA-F]{8}-(?:[0-9a-fA-F]{4}-){3}[0-9a-fA-F]{12}$")
PROFILE = re.compile(r"^[a-z][a-z0-9_]{0,49}$")
WORKSPACE_LABEL = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$")
MAX_BYTES = 5_000_000
MAX_ROWS = 10_000
MAX_CACHE_AGE = 24 * 60 * 60
OPERATIONS = """_LogOperation
| where TimeGenerated > ago(15m)
| summarize Errors=countif(Level == 'Error'), Warnings=countif(Level == 'Warning')
"""


class CollectionError(Exception):
    pass


def validate_config(config):
    if not GUID.fullmatch(config["tenant"]) or not GUID.fullmatch(config["client"]):
        raise CollectionError("Tenant and client IDs must be GUIDs")
    raw_workspaces = config["workspaces"]
    if isinstance(raw_workspaces, str):
        # Accept direct/manual invocations and rules saved before the List form.
        raw_workspaces = raw_workspaces.splitlines()
    if not isinstance(raw_workspaces, list):
        raise CollectionError("Workspace IDs must be unique GUIDs, one per entry")
    if any(not isinstance(entry, (str, dict)) for entry in raw_workspaces):
        raise CollectionError("Workspace IDs must be unique GUIDs, one per entry")
    try:
        workspaces = [(entry if isinstance(entry, str) else entry["id"]).strip().lower()
                      for entry in raw_workspaces]
    except (KeyError, AttributeError):
        raise CollectionError("Workspace IDs must be unique GUIDs, one per entry") from None
    if not workspaces or len(set(workspaces)) != len(workspaces) or any(not GUID.fullmatch(x) for x in workspaces):
        raise CollectionError("Workspace IDs must be unique GUIDs, one per entry")
    labels = {}
    for workspace, entry in zip(workspaces, raw_workspaces):
        label = entry.get("label", workspace) if isinstance(entry, dict) else workspace
        if not isinstance(label, str) or not WORKSPACE_LABEL.fullmatch(label):
            raise CollectionError("Workspace labels must use 1-64 letters, digits, dots, underscores or hyphens")
        labels[workspace] = label
    if len({name.lower() for name in labels.values()}) != len(labels):
        raise CollectionError("Workspace labels must be unique across this rule")
    workspace_ids = set(workspaces)
    if any(label.lower() in workspace_ids - {workspace} for workspace, label in labels.items()):
        raise CollectionError("Workspace label cannot match another workspace ID")
    shared = [{"id": "workspace_operations", "query": OPERATIONS}] if config.get("workspace_operations", True) else []
    shared.extend(_validate_profiles(config.get("profiles", [])))
    per_workspace = {}
    for workspace, entry in zip(workspaces, raw_workspaces):
        per_workspace[workspace] = _validate_profiles(entry.get("profiles", [])) if isinstance(entry, dict) else []
    if len(workspaces) > 1000 or any(len({p["id"] for p in shared + per_workspace[w]}) > 50 for w in workspaces):
        raise CollectionError("Too many profiles or workspaces")
    if not shared and not any(per_workspace.values()):
        raise CollectionError("At least one profile or workspace operations query is required")
    return workspaces, shared, per_workspace, labels


def _validate_profiles(entries):
    if not isinstance(entries, list):
        raise CollectionError("Profiles must be a list")
    profiles, ids = [], {"workspace_operations"}
    for entry in entries:
        if not isinstance(entry, dict):
            raise CollectionError("Profiles need unique lowercase IDs and nonempty KQL")
        key, query_text = entry.get("id"), entry.get("query")
        if (not isinstance(key, str) or not PROFILE.fullmatch(key) or key in ids
                or not isinstance(query_text, str) or not query_text.strip()):
            raise CollectionError("Profiles need unique lowercase IDs and nonempty KQL")
        expected = entry.get("expected_services", [])
        if (not isinstance(expected, list) or len(expected) > MAX_ROWS
                or any(not isinstance(name, str) or not name or len(name) > 150 for name in expected)
                or len(expected) != len(set(expected))):
            raise CollectionError("Expected services must be unique, nonempty names (up to 10,000)")
        ids.add(key)
        profiles.append(entry)
    return profiles


def request_json(url, *, data=None, headers=None, attempts=4, timeout=70):
    headers = headers or {}
    payload = None if data is None else urlencode(data).encode() if isinstance(data, dict) and headers.get("Content-Type") != "application/json" else json.dumps(data).encode()
    for attempt in range(attempts):
        try:
            req = Request(url, data=payload, headers=headers, method="POST")
            with urlopen(req, timeout=timeout) as response:
                raw = response.read(MAX_BYTES + 1)
            if len(raw) > MAX_BYTES:
                raise CollectionError("Azure response exceeds 5 MB; aggregate more in KQL")
            return json.loads(raw)
        except HTTPError as exc:
            if exc.code not in (429, 500, 502, 503, 504) or attempt == attempts - 1:
                # Error codes help diagnose KQL and RBAC failures without printing
                # request bodies, access tokens, or query contents into check output.
                try:
                    detail = json.loads(exc.read(4096)).get("error", {})
                    code = detail.get("code", "") if isinstance(detail, dict) else ""
                except (ValueError, OSError, UnicodeError):
                    code = ""
                if isinstance(code, str) and re.fullmatch(r"[A-Za-z0-9_.-]{1,80}", code):
                    raise CollectionError(f"Azure HTTP {exc.code} ({code})") from None
                raise CollectionError(f"Azure HTTP {exc.code}") from None
            retry = exc.headers.get("Retry-After", "")
            pause = min(float(retry), 20.0) if retry.isdigit() else min(2**attempt + random.random(), 20.0)
            time.sleep(pause)
        except (URLError, TimeoutError) as exc:
            if attempt == attempts - 1:
                raise CollectionError("Azure connection/timeout failure") from None
            time.sleep(min(2**attempt + random.random(), 20.0))
        except (ValueError, UnicodeError):
            raise CollectionError("Invalid Azure JSON response") from None
    raise CollectionError("Azure request failed")


def token(tenant, client, secret):
    response = request_json(
        f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token",
        data={"client_id": client, "client_secret": secret, "grant_type": "client_credentials", "scope": "https://api.loganalytics.io/.default"},
        headers={"Content-Type": "application/x-www-form-urlencoded"},
    )
    if not response.get("access_token"):
        raise CollectionError("Token response lacks access token")
    return response["access_token"]


def query(workspace, profile, access_token, label=None):
    label = label or workspace
    response = request_json(
        f"https://api.loganalytics.azure.com/v1/workspaces/{workspace}/query",
        data={"query": profile["query"], "timespan": "PT15M"} if profile["id"] == "workspace_operations" else {"query": profile["query"]},
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {access_token}"},
    )
    if response.get("error"):
        raise CollectionError("Azure returned partial or failed query result")
    tables = response.get("tables")
    if not isinstance(tables, list) or len(tables) != 1:
        raise CollectionError("Expected one Azure result table")
    table = tables[0]
    names = [c["name"] for c in table["columns"]]
    if len(names) != len(set(names)) or len(table["rows"]) > MAX_ROWS:
        raise CollectionError("Duplicate columns or over 10,000 result rows")
    rows = []
    for cells in table["rows"]:
        if len(names) != len(cells):
            raise CollectionError("Malformed Azure result row")
        record = dict(zip(names, cells))
        if profile["id"] == "workspace_operations":
            errors, warnings = int(record["Errors"]), int(record["Warnings"])
            if min(errors, warnings) < 0:
                raise ValueError("Negative operation count")
            rows.append({"kind": "item", "profile": "workspace_operations", "item": f"{label} operations", "errors": errors, "warnings": warnings, "message": f"{errors} errors, {warnings} warnings in 15 minutes"})
        else:
            service, value = str(record["Service"]), float(record["Value"])
            if not service or len(service) > 150 or not (-1e300 < value < 1e300):
                raise ValueError("Invalid service or value")
            row = {"kind": "item", "profile": profile["id"], "item": f"{label} {profile['id']} {service}", "value": value, "message": str(record.get("Message") or f"Value: {value}")[:900]}
            row.update({k: profile[k] for k in ("warn", "crit") if k in profile})
            rows.append(row)
    if profile["id"] == "workspace_operations" and len(rows) != 1:
        raise CollectionError("Operations query must return exactly one row")
    if profile["id"] != "workspace_operations" and not rows and not profile.get("expected_services"):
        raise CollectionError("Custom query returned no rows; emit explicit zero rows")
    if len({row["item"] for row in rows}) != len(rows):
        raise CollectionError("Duplicate service items in query response")
    if profile["id"] != "workspace_operations":
        returned = {row["item"] for row in rows}
        for service in profile.get("expected_services", []):
            item = f"{label} {profile['id']} {service}"
            if item not in returned:
                rows.append({"kind": "item", "profile": profile["id"], "item": item,
                             "missing": True, "message": f"Expected service {service} absent from successful query result"})
    if len(rows) > MAX_ROWS:
        raise CollectionError("Over 10,000 services after adding expected services")
    return rows


def cache_path(config, workspace, profile, cache_dir, label=None):
    identity_parts = [config["tenant"], config["client"], workspace, profile["id"],
                      profile["query"], profile.get("expected_services", [])]
    if label and label != workspace:
        identity_parts.append(label)
    identity = json.dumps(identity_parts, sort_keys=True)
    return cache_dir / (hashlib.sha256(identity.encode()).hexdigest() + ".json")


def load_cache(path, now):
    try:
        document = json.loads(path.read_text())
        stamp = int(document["timestamp"])
        if stamp > now or stamp < 0 or now - stamp > MAX_CACHE_AGE:
            return [], None
        rows = document["rows"]
        if not isinstance(rows, list) or len(rows) > MAX_ROWS:
            return [], None
        return [dict(row, cached=True, age_seconds=now - stamp) for row in rows], now - stamp
    except (OSError, ValueError, TypeError, KeyError):
        return [], None


def save_cache(path, rows, now):
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    with tmp.open("w") as stream:
        os.chmod(tmp, 0o600)
        json.dump({"timestamp": now, "rows": rows}, stream)
    os.replace(tmp, path)


def collect(config, secret, cache_dir):
    workspaces, shared, per_workspace, labels = validate_config(config)
    now = int(time.time())
    jobs = []
    for workspace in workspaces:
        selected = {profile["id"]: profile for profile in shared}
        selected.update({profile["id"]: profile for profile in per_workspace[workspace]})
        jobs.extend((workspace, profile, labels[workspace]) for profile in selected.values())
    try:
        access_token = token(config["tenant"], config["client"], secret)
    except CollectionError as exc:
        failures = [str(exc)] * len(jobs)
        rows = []
        ages = []
        results = []
        for w, p, label in jobs:
            cached, age = load_cache(cache_path(config, w, p, cache_dir, label), now)
            rows.extend(cached)
            if age is not None:
                ages.append(age)
            results.append({"workspace": w, "label": label, "profile": p["id"], "status": "not_run",
                            "error": f"Token acquisition failed: {exc}", "cached_age_seconds": age})
        return collector(len(jobs), failures, ages, results), rows

    failures, ages, rows = [], [], []
    results = {}
    with ThreadPoolExecutor(max_workers=3) as pool:
        futures = {pool.submit(query, w, p, access_token, label): (w, p, label)
                   for w, p, label in jobs}
        for future in as_completed(futures):
            w, p, label = futures[future]
            path = cache_path(config, w, p, cache_dir, label)
            try:
                fresh = future.result()
                save_cache(path, fresh, now)
                rows.extend(fresh)
                results[(w, p["id"])] = {"workspace": w, "label": label, "profile": p["id"], "status": "ok"}
            except (CollectionError, ValueError, TypeError, KeyError, OSError) as exc:
                error = type(exc).__name__ if not isinstance(exc, CollectionError) else str(exc)
                failures.append(f"{w}/{p['id']}: {error}")
                cached, age = load_cache(path, now)
                rows.extend(cached)
                if age is not None:
                    ages.append(age)
                results[(w, p["id"])] = {"workspace": w, "label": label, "profile": p["id"], "status": "failed",
                                           "error": error, "cached_age_seconds": age}
    return collector(len(jobs), failures, ages,
                     [results[(w, p["id"])] for w, p, _ in jobs]), rows


def collector(total, failures, ages, results=None):
    message = f"{total - len(failures)}/{total} queries OK" if total else "Collection did not start"
    if failures:
        message += f"; {len(failures)} failed"
        if results:
            first = next((row for row in results if row["status"] != "ok"), None)
            if first:
                brief = first["error"].removeprefix("Azure ")[:90]
                if first["status"] == "not_run":
                    message += f"; token failed ({brief.removeprefix('Token acquisition failed: ')})"
                else:
                    name = first.get("label", first["workspace"])
                    short = name if name != first["workspace"] else f"{name[:8]}..."
                    message += f" in workspace {short} ({brief})"
        else:
            message += f"; first: {failures[0]}"
    if ages:
        message += f"; oldest cached data {max(ages)}s"
    return {"kind": "collector", "total": total, "failures": len(failures),
            "message": message[:500], "queries": (results or [])[:500],
            "omitted_queries": max(0, len(results or []) - 500)}


def main():
    from cmk.password_store.v1_unstable import parser_add_secret_option, resolve_secret_option

    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    parser_add_secret_option(parser, long="--secret", required=True, help="Entra application secret")
    args = parser.parse_args()
    try:
        config = json.loads(args.config)
        secret = resolve_secret_option(args, "secret").reveal()
        cache_dir = Path(os.environ.get("OMD_ROOT", str(Path.home()))) / "var/check_mk/azure_loganalytics_cache"
        status, rows = collect(config, secret, cache_dir)
    except CollectionError as exc:
        status, rows = collector(0, [str(exc)], []), []
    except (ValueError, TypeError, KeyError, OSError):
        status, rows = collector(0, ["Configuration or password-store error"], []), []
    print("<<<azure_loganalytics:sep(0)>>>")
    print(json.dumps(status, ensure_ascii=True))
    for row in rows:
        print(json.dumps(row, ensure_ascii=True))


if __name__ == "__main__":
    main()
