#!/usr/bin/env python3
"""Checkmk special agent for Proxmox Mail Gateway (PMG) via its REST API.

Unlike Proxmox VE, PMG has no API-token mechanism (/access/users/{id}/token
does not exist). Authentication is ticket-based:

    1. POST /api2/json/access/ticket with username[@realm] + password
       -> {"ticket": ..., "CSRFPreventionToken": ..., ...}
    2. Send the ticket back as Cookie: PMGAuthCookie=<ticket> on every
       subsequent request. CSRFPreventionToken is only required for
       state-changing (POST/PUT/DELETE) calls; this agent is read-only.

The web/API port is hard-coded to 8006 in PMG itself (no config option).

Emits Checkmk agent sections:

    <<<pmg_statistics:sep(0)>>>   {"mail": {...}, "rejectcount": [...]}
    <<<pmg_queue:sep(0)>>>        [{"queue": "deferred", "count": N}, ...]
    <<<pmg_quarantine:sep(0)>>>   {"spam": {...}, "virus": {...}}
    <<<pmg_quarantine_queue:sep(0)>>>
                                  {"spam": N, "virus": N, "attachment": N}
    <<<pmg_clamav:sep(0)>>>       [{"type": ..., "version": ..., ...}, ...]
    <<<pmg_spamassassin:sep(0)>>> [{"channel": ..., "update_avail": ...}, ...]
    <<<pmg_node:sep(0)>>>         {"status": {...}, "subscription": {...},
                                  "updates": [...], "certificates": [...]}

Note on `statistics/mail`, `statistics/rejectcount` and
`quarantine/spamstatus`/`quarantine/virusstatus`: the PMG API applies an
undocumented, non-configurable rolling 24h window to these when no
starttime/endtime is given (see PMG::API2::Statistics: "$start = time -
86400"). `quarantine/{spam,virus,attachment}` (used for the *_queue sections
below) are different again -- they return unfiltered per-mail lists unless
starttime/endtime is passed, which is why --quarantine-lookback-days exists
for those.

A read-only PMG user (role "Audit") is sufficient for all of these
endpoints.
"""
import argparse
import json
import ssl
import sys
import time
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError


QUEUES = ("deferred", "active", "incoming", "hold")
QUARANTINE_QUEUES = ("spam", "virus", "attachment")


def parse_args(argv):
    p = argparse.ArgumentParser(description="Checkmk PMG special agent")
    p.add_argument("--username", required=True,
                   help="PMG username, e.g. 'checkmk' (realm via --realm)")
    p.add_argument("--password", required=True, help="PMG password")
    p.add_argument("--realm", default="pmg",
                   help="PMG authentication realm (default: pmg)")
    p.add_argument("--port", type=int, default=8006,
                   help="HTTPS port (default 8006, hard-coded in PMG)")
    p.add_argument("--no-cert-check", action="store_true",
                   help="Disable TLS certificate verification (self-signed)")
    p.add_argument("--timeout", type=int, default=20,
                   help="Per-request timeout (s)")
    p.add_argument("--quarantine-lookback-days", type=int, default=30,
                   help="Fetch quarantine/{spam,virus,attachment} mails from "
                        "the last N days (default 30) for the "
                        "*Queue services")
    p.add_argument("hostname", help="PMG host / address")
    return p.parse_args(argv)


def make_context(no_cert_check):
    ctx = ssl.create_default_context()
    if no_cert_check:
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE
    return ctx


class PMGClient:
    def __init__(self, host, port, username, password, realm, ctx, timeout):
        if ":" in host and not host.startswith("["):
            host = "[%s]" % host
        self.base = "https://%s:%d/api2/json" % (host, port)
        self.ctx = ctx
        self.timeout = timeout
        self.ticket = None
        self._login(username, password, realm)

    def _login(self, username, password, realm):
        url = "%s/access/ticket" % self.base
        payload = {
            "username": username,
            "password": password,
            "realm": realm,
        }
        data = json.dumps(payload).encode("utf-8")
        req = Request(url, data=data, method="POST",
                      headers={"Content-Type": "application/json"})
        with urlopen(req, context=self.ctx, timeout=self.timeout) as resp:
            body = json.loads(resp.read().decode("utf-8"))
        result = body.get("data", {})
        self.ticket = result.get("ticket")
        if not self.ticket:
            raise ValueError("PMG login did not return a ticket")

    def call(self, path, params=None):
        url = "%s/%s" % (self.base, path)
        if params:
            from urllib.parse import urlencode
            url += "?" + urlencode(params)
        req = Request(url, method="GET")
        req.add_header("Cookie", "PMGAuthCookie=%s" % self.ticket)
        with urlopen(req, context=self.ctx, timeout=self.timeout) as resp:
            body = json.loads(resp.read().decode("utf-8"))
        return body.get("data")

    def node_name(self):
        nodes = self.call("nodes")
        if isinstance(nodes, list) and nodes:
            return nodes[0].get("node", "localhost")
        return "localhost"


def emit(section, payload):
    sys.stdout.write("<<<%s:sep(0)>>>\n" % section)
    sys.stdout.write(json.dumps(payload) + "\n")


def safe_call(client, path, params=None):
    try:
        return client.call(path, params)
    except (HTTPError, URLError, ValueError) as exc:
        return {"_error": str(exc)}


def main(argv=None):
    args = parse_args(argv if argv is not None else sys.argv[1:])
    ctx = make_context(args.no_cert_check)
    try:
        client = PMGClient(args.hostname, args.port, args.username,
                           args.password, args.realm, ctx, args.timeout)
    except (HTTPError, URLError, ValueError) as exc:
        # Without a ticket nothing else can be fetched -- emit error
        # sections for everything so checks surface UNKNOWN instead of
        # going stale/missing.
        err = {"_error": "PMG login failed: %s" % exc}
        for section in ("pmg_statistics", "pmg_queue", "pmg_quarantine",
                        "pmg_quarantine_queue", "pmg_clamav",
                        "pmg_spamassassin", "pmg_node"):
            emit(section, err)
        return 1

    node = client.node_name()

    # --- Statistics (mail counters + reject counters) ---
    statistics = {
        "mail": safe_call(client, "statistics/mail"),
        "rejectcount": safe_call(client, "statistics/rejectcount"),
    }
    emit("pmg_statistics", statistics)

    # --- Postfix queue depths ---
    queue_data = []
    for queue in QUEUES:
        result = safe_call(client, "nodes/%s/postfix/queue/%s" % (node, queue))
        if isinstance(result, dict) and "_error" in result:
            queue_data.append({"queue": queue, "_error": result["_error"]})
        else:
            count = len(result) if isinstance(result, list) else 0
            queue_data.append({"queue": queue, "count": count})
    emit("pmg_queue", queue_data)

    # --- Quarantine status (spam + virus) ---
    quarantine = {
        "spam": safe_call(client, "quarantine/spamstatus"),
        "virus": safe_call(client, "quarantine/virusstatus"),
    }
    emit("pmg_quarantine", quarantine)

    # --- Quarantine queue (current backlog, needs release/delete) ---
    # Unlike spamstatus/virusstatus (unfiltered full DB history), these
    # endpoints need an explicit starttime/endtime window.
    endtime = int(time.time())
    starttime = endtime - args.quarantine_lookback_days * 86400
    quarantine_queue = {}
    for qtype in QUARANTINE_QUEUES:
        result = safe_call(client, "quarantine/%s" % qtype,
                            params={"starttime": starttime, "endtime": endtime})
        if isinstance(result, dict) and "_error" in result:
            quarantine_queue[qtype] = result
        else:
            quarantine_queue[qtype] = len(result) if isinstance(result, list) else 0
    emit("pmg_quarantine_queue", quarantine_queue)

    # --- ClamAV virus database status ---
    clamav = safe_call(client, "nodes/%s/clamav/database" % node)
    emit("pmg_clamav", clamav)

    # --- SpamAssassin rules status ---
    spamassassin = safe_call(client, "nodes/%s/spamassassin/rules" % node)
    emit("pmg_spamassassin", spamassassin)

    # --- Node status, subscription, updates, certificates ---
    node_info = {
        "status": safe_call(client, "nodes/%s/status" % node),
        "subscription": safe_call(client, "nodes/%s/subscription" % node),
        "updates": safe_call(client, "nodes/%s/apt/update" % node),
        "certificates": safe_call(client, "nodes/%s/certificates/info" % node),
    }
    emit("pmg_node", node_info)
    return 0


if __name__ == "__main__":
    sys.exit(main())
