#!/usr/bin/env python3
"""Point a fleet of Gude Expert Power Control units at Neowit's MQTT broker.

Gude's own configuration file *is* a list of console commands, and a unit accepts
one over its HTTP interface. So provisioning a rack is: generate the handful of
"mqtt 1 ..." lines each unit needs, upload, restart. That covers the whole
configuration -- broker, TLS, credentials and client ID included -- which is the
part that otherwise has to be typed into every unit's web interface by hand.

Everything here goes over the LAN to each unit. Nothing is sent through the
broker, so a unit that has never connected is provisioned exactly like one that
has, and "Permit CLI commands" does not need switching on first.

Usage
-----
    # Credentials come from the integration's Connection details.
    export GUDE_MQTT_USERNAME='<the integration id>'
    export GUDE_MQTT_PASSWORD='key:<id>:<secret>'

    # And the units' own web logins.
    export GUDE_HTTP_PASSWORD='<the unit admin password>'

    # Show what would be sent. Changes nothing.
    python3 gude_provision.py --unit 10.0.0.11=pdu1 --unit 10.0.0.12=pdu2

    # Do it.
    python3 gude_provision.py --unit 10.0.0.11=pdu1 --apply

Each --unit is ADDRESS=CLIENT_ID. Client IDs must be unique per unit: two units
sharing one will disconnect each other.

Every run finishes by reading the unit's MQTT settings back from its own
configuration JSON, so a file that was uploaded but never applied is reported as
a failure rather than a success.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import ssl
import sys
import time
import urllib.error
import urllib.request
from collections import Counter
from dataclasses import dataclass

# The settings written to every unit, in the order they are applied. These are
# Gude console commands: the same syntax the unit's own "Config File Export"
# produces, so anything valid here is valid in a hand-edited config file too.
#
# Deliberately absent: "system fabsettings". A generated Gude config file starts
# with it and resets the unit to factory state first, which would take its
# network settings with it. Without that line the commands apply on top of the
# unit's current configuration, which is what provisioning wants.
def config_lines(broker: str, port: int, username: str, password: str,
                 client_id: str, topic: str, data_timer: int,
                 credentials: bool = True) -> list[str]:
    lines = [
        f'mqtt 1 server set "{broker}"',
        f"mqtt 1 port set {port}",
        "mqtt 1 tls enabled set 1",
    ]
    # Leaving the credentials alone is for re-provisioning units already talking
    # to Neowit: theirs are right already, and the run then needs no secrets.
    if credentials:
        lines += [
            f'mqtt 1 user set "{username}"',
            f'mqtt 1 passwd set "{password}"',
        ]
    lines += [
        # Unique per unit. The broker drops one of two clients sharing an ID.
        f'mqtt 1 client set "{client_id}"',
        # Neowit identifies a unit from the mac_addr in its telemetry, not from
        # its topic, so any prefix works. This is the factory default, and
        # converging a fleet on it means one less thing that varies per unit.
        # "[mac]" is literal: the unit substitutes its own MAC address.
        f'mqtt 1 topic set "{topic}"',
        # At most once means the broker may drop a summary or an outlet-changed
        # event with no retry. At least once costs nothing here.
        "mqtt 1 qos set 1",
        # The other half of QoS 1. Without a clean session the broker queues
        # commands for a unit while it is offline and delivers them in a burst
        # when it returns -- a power cycle nobody wants any more, minutes late.
        "mqtt 1 clean set 1",
        "mqtt 1 keepalive set 60",
        # How often the unit sends its readings. Neowit reads this value from the
        # unit and sizes its offline check against it, so anything from 60 to 900
        # is fine. Faster means more stored readings for no gain in how quickly an
        # outlet change is noticed -- those arrive immediately, whatever this is.
        f"mqtt 1 device data timer set {data_timer}",
        # Without this the unit reports readings but silently ignores every switch
        # and power-cycle command.
        "mqtt 1 console enabled set 1",
        # Last, so the unit only dials out once it has everything it needs.
        "mqtt 1 enabled set 1",
    ]
    return lines


@dataclass
class Unit:
    address: str
    client_id: str


class Device:
    """One unit's HTTP interface."""

    def __init__(self, address: str, scheme: str, auth: str, timeout: float,
                 context: ssl.SSLContext | None):
        self.base = f"{scheme}://{address}"
        self.auth = auth
        self.timeout = timeout
        self.context = context

    def _open(self, req: urllib.request.Request) -> bytes:
        # Sent up front rather than waiting for a 401: the unit answers some
        # paths anonymously and only challenges on others.
        req.add_header("Authorization", f"Basic {self.auth}")
        with urllib.request.urlopen(req, timeout=self.timeout,
                                    context=self.context) as resp:
            return resp.read()

    def get(self, path: str) -> bytes:
        return self._open(urllib.request.Request(self.base + path))

    def post_file(self, path: str, filename: str, payload: bytes) -> bytes:
        # Endpoint, field name and type numbers match Gude's own deployment tool
        # (github.com/gudesystems/gude-device-manager), which is the closest thing
        # to a specification for this: fwupdate.txt, a "fwupload" part, and
        # type=2 for a configuration file.
        boundary = "----gude-provision-boundary"
        body = b"".join([
            f'--{boundary}\r\nContent-Disposition: form-data; name="fwupload"; '
            f'filename="{filename}"\r\n'.encode(),
            b"Content-Type: application/octet-stream\r\n\r\n",
            payload,
            f"\r\n--{boundary}--\r\n".encode(),
        ])
        req = urllib.request.Request(self.base + path, data=body, method="POST")
        req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
        return self._open(req)

    # Component bitmasks on statusjsn.js / cfgjsn.js. 16 is the "misc" block
    # carrying product name and firmware; 32 on cfgjsn.js is the MQTT section.
    SYSINFO = 16
    MQTT_CONFIG = 32

    def identify(self) -> str:
        """Reachability and login check, reported as whatever the unit says it is."""
        raw = self.get(f"/statusjsn.js?components={self.SYSINFO}")
        try:
            misc = json.loads(raw).get("misc") or {}
        except ValueError:
            return "reachable"
        name = misc.get("product_name") or "reachable"
        firmware = misc.get("firm_v")
        return f"{name} (firmware {firmware})" if firmware else name

    def mqtt_config(self) -> dict:
        """The unit's current MQTT settings, as it reports them."""
        raw = self.get(f"/cfgjsn.js?components={self.MQTT_CONFIG}")
        brokers = (json.loads(raw).get("mqtt") or {}).get("broker") or [{}]
        return brokers[0]

    def upload_config(self, text: str) -> None:
        self.post_file("/fwupdate.txt?type=2", "config.txt", text.encode())

    def export_config(self) -> bytes:
        """The unit's whole configuration, as the console commands that recreate it."""
        return self.get("/config.txt")

    def restart(self) -> None:
        # An uploaded configuration is only applied on restart. A unit that goes
        # down before answering is the expected case, not a failure -- what
        # matters is whether it comes back, which wait_until_back decides.
        try:
            self.get("/?cmd=39")
        except (urllib.error.URLError, OSError):
            pass

    def wait_until_back(self, attempts: int = 20, pause: float = 3.0) -> bool:
        for _ in range(attempts):
            time.sleep(pause)
            try:
                self.get(f"/statusjsn.js?components={self.SYSINFO}")
                return True
            except (urllib.error.URLError, OSError):
                continue
        return False


# What the unit should report back once the configuration has taken, keyed by the
# field name it uses in cfgjsn.js. Anything not listed here isn't checked.
def expected_config(broker: str, port: int, username: str, client_id: str,
                    topic: str, data_timer: int) -> dict:
    return {
        "enabled": 1,
        "tls": 1,
        "allow_cli": 1,
        "hostname": broker,
        "port": port,
        "user": username,
        "client_id": client_id,
        "tprefix": topic,
        "qos": 1,
        "cln_ses": 1,
        "keep_alive": 60,
        "telemetry_interv": data_timer,
    }


def verify(device: Device, expected: dict) -> list[str]:
    """Return the settings the unit does not report as asked for."""
    actual = device.mqtt_config()
    return [f"{key}: asked for {value!r}, unit reports {actual.get(key)!r}"
            for key, value in expected.items() if actual.get(key) != value]


def provision(unit: Unit, device: Device, config: str, expected: dict,
              apply_changes: bool) -> bool:
    print(unit.address)
    reachable = True
    try:
        print(f"  {'unit':14} {device.identify()}")
    except urllib.error.HTTPError as err:
        detail = "wrong web login" if err.code in (401, 403) else f"HTTP {err.code}"
        print(f"  {'unit':14} unreachable: {detail}")
        reachable = False
    except (urllib.error.URLError, OSError) as err:
        print(f"  {'unit':14} unreachable: {err}")
        reachable = False

    # A dry run still shows what would be sent even when the unit can't be
    # reached, so a config can be reviewed away from the rack.
    if not apply_changes:
        print(f"  {'client id':14} {unit.client_id}")
        for line in config.splitlines():
            if line and not line.startswith("#"):
                print(f"    {line}")
        if reachable:
            differences = verify(device, expected)
            print(f"  {'to change':14} "
                  + (f"{len(differences)} setting(s)" if differences else "nothing"))
            for difference in differences:
                print(f"    {difference}")
        return reachable
    if not reachable:
        return False
    print(f"  {'client id':14} {unit.client_id}")

    try:
        device.upload_config(config)
        print(f"  {'config':14} uploaded")
        device.restart()
        print(f"  {'restart':14} requested")
    except (urllib.error.URLError, OSError) as err:
        print(f"  FAILED: {err}")
        return False

    if not device.wait_until_back():
        print(f"  {'status':14} did not come back -- check it by hand")
        return False

    # Read the settings back rather than trusting the upload: a staged file that
    # was never applied, or a command the firmware ignored, both look like
    # success up to this point.
    differences = verify(device, expected)
    if differences:
        print(f"  {'status':14} back up, but not configured as asked:")
        for difference in differences:
            print(f"    {difference}")
        return False
    print(f"  {'status':14} back up, settings confirmed")
    return True


def parse_unit(value: str) -> Unit:
    address, _, client_id = value.partition("=")
    if not address or not client_id:
        raise argparse.ArgumentTypeError(
            f"--unit wants ADDRESS=CLIENT_ID, got {value!r}")
    return Unit(address, client_id)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--unit", action="append", default=[], type=parse_unit,
                        metavar="ADDRESS=CLIENT_ID",
                        help="a unit's address on your network and the client ID "
                             "to give it; repeatable")
    parser.add_argument("--broker", default="mqtt.neowit.io")
    parser.add_argument("--broker-port", type=int, default=8883)
    parser.add_argument("--username", default=os.environ.get("GUDE_MQTT_USERNAME", ""),
                        help="the Neowit integration id (GUDE_MQTT_USERNAME)")
    parser.add_argument("--password", default=os.environ.get("GUDE_MQTT_PASSWORD", ""),
                        help="key:<id>:<secret> from the connection details "
                             "(GUDE_MQTT_PASSWORD)")
    parser.add_argument("--http-user", default=os.environ.get("GUDE_HTTP_USER", "admin"),
                        help="the units' web login (GUDE_HTTP_USER)")
    parser.add_argument("--http-password", default=os.environ.get("GUDE_HTTP_PASSWORD", ""),
                        help="the units' web password (GUDE_HTTP_PASSWORD). Leave "
                             "unset for a unit with its HTTP password turned off.")
    parser.add_argument("--topic", default="de/gudesystems/epc/[mac]",
                        help="topic prefix to set. [mac] is substituted by the unit.")
    parser.add_argument("--data-timer", type=int, default=300,
                        help="telemetry interval in seconds")
    parser.add_argument("--scheme", choices=("https", "http"), default="https",
                        help="the units carry the broker password, so https unless "
                             "a unit has no TLS enabled")
    parser.add_argument("--insecure", action="store_true",
                        help="accept the unit's own certificate without verifying "
                             "it -- normal, factory certificates are self-signed")
    parser.add_argument("--keep-credentials", action="store_true",
                        help="don't write the username and password. For units "
                             "already connected to Neowit, whose credentials are "
                             "right; no MQTT secrets are then needed at all.")
    parser.add_argument("--timeout", type=float, default=10.0)
    parser.add_argument("--apply", action="store_true",
                        help="upload and restart. Without this, nothing is changed.")
    parser.add_argument("--dry-run", action="store_true",
                        help="report what would change and stop. This is the "
                             "default; the flag just says so out loud.")
    args = parser.parse_args()

    if args.apply and args.dry_run:
        parser.error("--apply and --dry-run ask for opposite things")
    if not args.unit:
        parser.error("give at least one --unit ADDRESS=CLIENT_ID")
    if not args.keep_credentials and (not args.username or not args.password):
        parser.error("--username and --password are required "
                     "(or set GUDE_MQTT_USERNAME / GUDE_MQTT_PASSWORD), "
                     "unless you pass --keep-credentials")
    # No check on the web password: a unit with its HTTP password switched off
    # serves and accepts everything unauthenticated, and many racks are set up
    # that way.
    counts = Counter(u.client_id for u in args.unit)
    duplicates = [name for name, n in counts.items() if n > 1]
    if duplicates:
        parser.error(f"client IDs must be unique per unit; repeated: "
                     f"{', '.join(sorted(duplicates))}")

    context = None
    if args.scheme == "https":
        context = ssl.create_default_context()
        if args.insecure:
            context.check_hostname = False
            context.verify_mode = ssl.CERT_NONE

    auth = base64.b64encode(
        f"{args.http_user}:{args.http_password}".encode()).decode()

    if not args.apply:
        print("dry run -- nothing will be changed. Re-run with --apply.\n")

    failures = 0
    for unit in args.unit:
        lines = config_lines(args.broker, args.broker_port, args.username,
                             args.password, unit.client_id, args.topic,
                             args.data_timer, not args.keep_credentials)
        config = "# Written by gude_provision.py for Neowit.\n" + "\n".join(lines) + "\n"
        shown = config.replace(args.password, "********") if args.password else config
        expected = expected_config(args.broker, args.broker_port, args.username,
                                   unit.client_id, args.topic, args.data_timer)
        if args.keep_credentials:
            del expected["user"]
        device = Device(unit.address, args.scheme, auth, args.timeout, context)
        if not provision(unit, device, shown if not args.apply else config,
                         expected, args.apply):
            failures += 1
        print()

    done = len(args.unit) - failures
    verb = "provisioned" if args.apply else "ready to provision"
    print(f"{done} {verb}, {failures} failed")
    if args.apply and done:
        print("Settings above were read back from each unit. They appear under "
              "Devices in Neowit within one telemetry interval.")
    return 1 if failures else 0


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