#!/usr/bin/env python3
"""Verify the broker is configured the way the firmware expects.

Reads include/secrets.h by path and uses the values in place. It never prints
the password, and never echoes it to argv or the environment -- run it freely,
the output is safe to paste into a chat or an issue.

  python3 tools/check_broker.py

Exit status is 0 only if every check that matters passes.
"""

import argparse
import os
import re
import sys
import time

try:
    import paho.mqtt.client as mqtt
except ImportError:
    sys.exit("paho-mqtt not installed:  pip install paho-mqtt")

# paho 2.x normalises MQTT 3.1.1 CONNACK codes onto the MQTT 5 reason-code
# numbering, so both ranges can show up depending on callback API version.
CONNACK = {
    # MQTT 3.1.1
    0: "accepted",
    1: "refused: unacceptable protocol version",
    2: "refused: client id rejected",
    3: "refused: server unavailable",
    4: "refused: bad username or password",
    5: "refused: not authorized",
    # MQTT 5
    128: "refused: unspecified error",
    132: "refused: unsupported protocol version",
    133: "refused: client identifier not valid",
    134: "refused: bad username or password",
    135: "refused: not authorized",
    136: "refused: server unavailable",
    137: "refused: server busy",
    138: "refused: banned",
    151: "refused: quota exceeded",
    159: "refused: connection rate exceeded",
}

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_SECRETS = os.path.join(HERE, os.pardir, "include", "secrets.h")


def parse_secrets(path):
    """Pull the MQTT settings out of secrets.h without echoing them."""
    if not os.path.exists(path):
        sys.exit("no secrets.h at %s -- copy include/secrets.h.example first" % path)
    text = open(path).read()

    def macro(name):
        m = re.search(r'^\s*#define\s+%s\s+"([^"]*)"' % name, text, re.M)
        return m.group(1) if m else None

    def macro_int(name):
        m = re.search(r"^\s*#define\s+%s\s+(\d+)" % name, text, re.M)
        return int(m.group(1)) if m else None

    cfg = {
        "host": macro("MQTT_HOST"),
        "port": macro_int("MQTT_PORT") or 1883,
        "user": macro("MQTT_USER"),
        "password": macro("MQTT_PASSWORD"),
    }
    if not cfg["host"]:
        sys.exit("MQTT_HOST not found in %s" % path)
    return cfg


def rc_to_int(rc):
    """paho 2.x hands the v2 callback a ReasonCode; paho 1.x hands it an int."""
    return int(getattr(rc, "value", rc))


def try_connect(host, port, user=None, password=None, client_id="broker-check"):
    """Return the CONNACK code, or a string describing a transport failure."""
    result = {}

    def on_connect(client, userdata, flags, rc, properties=None):
        result["rc"] = rc_to_int(rc)

    c = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, client_id=client_id)
    if user:
        c.username_pw_set(user, password)
    c.on_connect = on_connect
    try:
        c.connect(host, port, 10)
        c.loop_start()
        for _ in range(80):
            if "rc" in result:
                break
            time.sleep(0.1)
        c.loop_stop()
        c.disconnect()
    except Exception as exc:  # noqa: BLE001 - surfacing the transport error is the point
        return "transport error: %s" % exc
    return result.get("rc", "no CONNACK within 8s")


def check_publish(cfg, topic, retain=False):
    """True if a publish to `topic` is accepted and echoed back to us."""
    got = []
    c = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, client_id="broker-check-pub")
    if cfg["user"]:
        c.username_pw_set(cfg["user"], cfg["password"])
    c.on_message = lambda cl, u, m: got.append(m.payload)
    try:
        c.connect(cfg["host"], cfg["port"], 10)
        c.loop_start()
        c.subscribe(topic, qos=0)
        time.sleep(0.4)
        c.publish(topic, "check", qos=0, retain=retain)
        time.sleep(1.2)
        if retain:
            # Leave nothing behind: an empty retained payload clears it.
            c.publish(topic, "", qos=0, retain=True)
            time.sleep(0.6)
        c.loop_stop()
        c.disconnect()
    except Exception:  # noqa: BLE001
        return False
    return bool(got)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--secrets", default=DEFAULT_SECRETS)
    args = ap.parse_args()

    cfg = parse_secrets(args.secrets)
    print("broker   : %s:%d" % (cfg["host"], cfg["port"]))
    print("username : %s" % (cfg["user"] or "<none>"))
    print("password : %s\n" % ("<set, %d chars>" % len(cfg["password"]) if cfg["password"] else "<empty>"))

    failures = []
    warnings = []

    # 1. Anonymous must be refused once an authenticator exists.
    rc = try_connect(cfg["host"], cfg["port"], client_id="broker-check-anon")
    if not isinstance(rc, int):
        # Never treat "couldn't tell" as "secure" -- that is how an open broker
        # gets mistaken for a locked-down one.
        print("[?] anonymous connect     : INCONCLUSIVE (%s)" % rc)
        failures.append("Could not determine whether anonymous access is allowed: %s" % rc)
    elif rc == 0:
        print("[!] anonymous connect     : ACCEPTED  <-- broker is still open")
        warnings.append(
            "Anonymous connections are still accepted. No authenticator is "
            "configured, so the credentials below are not actually being checked."
        )
    else:
        print("[ok] anonymous connect    : %s" % CONNACK.get(rc, rc))

    # 2. Made-up credentials must be refused.
    rc = try_connect(cfg["host"], cfg["port"], "not-a-real-user", "not-a-real-password",
                     client_id="broker-check-bogus")
    if not isinstance(rc, int):
        print("[?] bogus credentials     : INCONCLUSIVE (%s)" % rc)
    elif rc == 0:
        print("[!] bogus credentials     : ACCEPTED  <-- not being verified")
    else:
        print("[ok] bogus credentials    : %s" % CONNACK.get(rc, rc))

    # 3. The real credentials must be accepted.
    if not cfg["user"] or "CHANGEME" in (cfg["user"] or ""):
        print("[!] configured credentials: MQTT_USER still unset/CHANGEME")
        failures.append("Fill in MQTT_USER and MQTT_PASSWORD in secrets.h.")
    else:
        rc = try_connect(cfg["host"], cfg["port"], cfg["user"], cfg["password"],
                         client_id="broker-check-real")
        label = CONNACK.get(rc, rc) if isinstance(rc, int) else rc
        if rc == 0:
            print("[ok] configured credentials: accepted")
        elif not isinstance(rc, int):
            print("[?] configured credentials: INCONCLUSIVE (%s)" % rc)
            failures.append("Could not test the configured account: %s" % rc)
        else:
            print("[!] configured credentials: %s" % label)
            failures.append(
                "The account in secrets.h was rejected (%s). It must match the "
                "user you created in EMQX exactly." % label
            )

        # 4. ACLs must permit the two topic trees the firmware publishes to.
        for topic, why in (
            ("soil/_check/probe", "state payloads"),
            ("homeassistant/sensor/_check/config", "HA discovery configs"),
        ):
            ok = check_publish(cfg, topic, retain=True)
            print("[%s] publish %-38s (%s)" % ("ok" if ok else "!", topic, why))
            if not ok:
                failures.append(
                    "Publishing to %s was blocked. Check the authorization rules "
                    "for user %r -- without this, %s never reach the broker."
                    % (topic, cfg["user"], why)
                )

    print()
    for w in warnings:
        print("WARNING: %s" % w)
    for f in failures:
        print("FAIL: %s" % f)
    if not failures and not warnings:
        print("All checks passed. The broker is ready for the node.")
    return 1 if failures else 0


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