#! /usr/bin/python3
"""Link an appliance to TurnKey Hub services.

Launch registration and readiness use the installed integration package.
Mail and DNS configuration share this helper across all appliance apps.
"""
import sys

# The public command opens setup; subcommands remain available to services.
if __name__ == "__main__" and len(sys.argv) == 1:
    import os
    os.execv("/usr/lib/inithooks/bin/hubservices.py", ["hubservices.py", "--setup"])

# A partial package must remain probeable without opening saved credentials,
# repairing mail configuration, creating locks or writing Python bytecode.
if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] in ("capabilities", "register"):
    sys.dont_write_bytecode = True
    try:
        if sys.argv[1] == "capabilities":
            from libinithooks.hub_capabilities import main
        else:
            from libinithooks.hub_launch import main
        sys.exit(main(sys.argv[2:]))
    except Exception:
        if sys.argv[1] == "capabilities":
            import json
            print(json.dumps({"protocol": "hub-services-v1", "ready": False,
                              "missing": ["libinithooks.hub_capabilities"]}))
        else:
            print("Hub registration could not finish. Check the installed Hub packages and retry.", file=sys.stderr)
        sys.exit(1)

import argparse
import contextlib
import email.message
import fcntl
import hashlib
import ipaddress
import json
import os
from pathlib import Path
import re
import secrets
import smtplib
import socket
import ssl
import subprocess
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request

try:
    from libinithooks import hub_runtime
    RUNTIME = hub_runtime if hub_runtime.enabled() else None
except ImportError:
    RUNTIME = None

TURNKEY_VERSION = Path("/etc/turnkey_version")
BUILD_TAGS = Path("/etc/apt/apt.conf.d/01turnkey")
ONBOARDING_STATE = Path("/var/lib/turnkey-onboarding/state.json")
HUB_EMAIL_CACHE = Path("/var/lib/inithooks/cache/HUB_ACCOUNT_EMAIL")
HUB = "https://hub.turnkeylinux.org"
STATE_DIR = Path("/var/lib/turnkey-mail")
POSTFIX_DIR = Path("/etc/postfix")
SYSTEMD_DIR = Path("/etc/systemd/system")
BACKUP_DIR = Path("/var/lib/tklbam")
HUBDNS_DIR = Path("/var/lib/hubdns")
HELPER = "/usr/lib/inithooks/bin/turnkey-hub"
MAX_JSON = 65536
DEVICE_CODE = re.compile(r"^[A-Za-z0-9_-]{32,256}$")
ENROLLMENT_TOKEN = re.compile(r"^[a-f0-9]{32}\.[a-f0-9]{64}$")
HOST = re.compile(r"^(?=.{1,253}$)[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?$", re.I)
ADDRESS = re.compile(r"^[A-Za-z0-9.!#$%&'*+/=?^_`{|}~-]+@[A-Za-z0-9.-]+$")
ERRORS = {
    "stale_configuration": "The Hub settings changed. Synchronize again.",
    "invalid_credentials": "This Hub connection is no longer authorized. Link it again.",
    "expired": "This connection request expired. Start a new connection.",
    "denied": "The connection was not approved. Start again when ready.",
    "invalid_token": "This connection request is invalid or expired. Start a new connection.",
    "revoked": "This appliance's connection was revoked. Connect it again in the Hub.",
    "not_authorized": "This appliance is no longer authorized. Check its connection in the Hub.",
    "rate_limited": "The Hub is busy. Please wait and try again.",
}


def appliance_description():
    """The app, TurnKey release and build type of this appliance; empty if unknown.

    Sent when linking so the Hub can show which machine is being approved. The
    app and release come from /etc/turnkey_version (turnkey-wordpress-18.0-
    bookworm-amd64). The build type is the tag the build task left in the apt
    configuration (proxmox, vm, ec2, docker...); an untagged build is an ISO.
    """
    try:
        match = re.fullmatch(r"turnkey-(.+)-(\d[^-]*)-([^-]+)-([^-]+)", TURNKEY_VERSION.read_text().strip())
    except OSError:
        return {}
    if not match:
        return {}
    description = {"app": match[1][:40], "release": match[2][:16]}
    try:
        tags = re.search(r"\(([^)]+)\)", BUILD_TAGS.read_text())
    except OSError:
        tags = None
    if tags:
        description["build_target"] = " ".join(tags[1].split()[1:])[:40] or "iso"
    return description


class Error(Exception):
    def __init__(self, message, code="failed", retryable=False):
        super().__init__(message)
        self.code = code
        self.retryable = retryable


def atomic(path, content, mode=0o600):
    path = Path(path)
    fd, temporary = tempfile.mkstemp(prefix="." + path.name + "-", dir=str(path.parent))
    try:
        os.fchmod(fd, mode)
        with os.fdopen(fd, "wb") as stream:
            stream.write(content.encode() if isinstance(content, str) else content)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def run(command, data=None, timeout=30):
    try:
        result = subprocess.run(command, input=data, stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE, timeout=timeout)
    except (OSError, subprocess.TimeoutExpired):
        raise Error("A required system command could not finish. Check the service and try again.", retryable=True)
    if result.returncode:
        # Command stderr may include SMTP or Postfix credential data.
        raise Error("The mail service could not apply these settings. Existing settings were restored.", retryable=True)
    return result.stdout


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def request(route, payload):
    """Only post credentials to the fixed, certificate-verified Hub origin."""
    if not re.fullmatch(r"[a-z/]+", route):
        raise Error("Invalid Hub operation.")
    body = json.dumps(payload).encode()
    req = urllib.request.Request(HUB + "/api/appmail/" + route, data=body,
                                 headers={"Content-Type": "application/json", "Accept": "application/json"})
    opener = urllib.request.build_opener(NoRedirect(), urllib.request.HTTPSHandler(context=ssl.create_default_context()))
    try:
        with opener.open(req, timeout=15) as response:
            raw = response.read(MAX_JSON + 1)
    except urllib.error.HTTPError as exc:
        code = "unavailable"
        try:
            detail = json.loads(exc.read(MAX_JSON + 1))
            code = detail.get("error", code)
        except (ValueError, AttributeError):
            pass
        retryable = exc.code in (408, 429) or exc.code >= 500
        raise Error(ERRORS.get(code, "The Hub could not complete this request. Check the connection in the Hub and try again."),
                    code=code if code in ERRORS else "unavailable", retryable=retryable)
    except (OSError, urllib.error.URLError):
        raise Error("The Hub is temporarily unavailable. Your connection will be retried.", code="unavailable", retryable=True)
    try:
        value = json.loads(raw)
        if len(raw) > MAX_JSON or not isinstance(value, dict):
            raise ValueError()
        return value
    except (ValueError, UnicodeError):
        raise Error("The Hub returned an invalid response. Please try again.", retryable=True)


def checked_config(value, password, custom=False):
    try:
        smtp = value["smtp"]
        host = smtp["host"].lower()
        port = int(smtp["port"])
        username = smtp["username"]
        if not HOST.fullmatch(host) or ".." in host or not 1 <= port <= 65535:
            raise ValueError()
        if not isinstance(username, str) or not username or any(c.isspace() or ord(c) < 32 for c in username) or ":" in username:
            raise ValueError()
        if not isinstance(password, str) or not password or any(ord(c) < 32 or ord(c) == 127 for c in password):
            raise ValueError()
        result = {"smtp": {"host": host, "port": port, "username": username}, "password": password,
                  "mode": "custom" if custom else "hub", "status": value.get("status", "ready")}
        if not custom:
            domain = value["domain"].lower()
            sender = value["sender"]
            dns = value.get("dyndns")
            if value.get("version") != 1 or result["status"] not in ("ready", "pending"):
                raise ValueError()
            if not HOST.fullmatch(domain) or not ADDRESS.fullmatch(sender) or sender.rsplit("@", 1)[1] != domain:
                raise ValueError()
            if dns and (not HOST.fullmatch(dns["hostname"]) or not dns["hostname"].endswith(".tklapp.com")
                        or dns["update_uri"] != HUB + "/api/appmail/dns/update"):
                raise ValueError()
            result.update(domain=domain, sender=sender, device_id=value["device_id"], dyndns=dns)
        return result
    except (KeyError, ValueError, TypeError, AttributeError):
        raise Error("The mail settings are incomplete or invalid. Existing settings were kept.")


def smtp_check(config):
    smtp = config["smtp"]
    try:
        context = ssl.create_default_context()
        if smtp["port"] == 465:
            client = smtplib.SMTP_SSL(smtp["host"], smtp["port"], timeout=10, context=context)
        else:
            client = smtplib.SMTP(smtp["host"], smtp["port"], timeout=10)
            client.ehlo()
            client.starttls(context=context)
            client.ehlo()
        with client:
            client.login(smtp["username"], config["password"])
    except (OSError, smtplib.SMTPException):
        raise Error("The SMTP server could not be authenticated over a verified encrypted connection. Check the settings and try again.", retryable=True)


def cf_entries(text):
    entries = {}
    for line in text.splitlines():
        if line[:1].isspace() and entries and not line.lstrip().startswith("#"):
            entries[key] += " " + line.strip()
        else:
            match = re.match(r"^([a-zA-Z0-9_]+)\s*=\s*(.*)$", line)
            if match:
                key = match[1]
                entries[key] = match[2]
    return entries


def change_cf(text, changes):
    """Replace exact parameter assignments, preserving unrelated configuration."""
    result = []
    skip = False
    for line in text.splitlines(keepends=True):
        if line[:1].isspace() and not line.lstrip().startswith("#"):
            if not skip:
                result.append(line)
            continue
        match = re.match(r"^([a-zA-Z0-9_]+)\s*=", line)
        skip = bool(match and match[1] in changes)
        if not skip:
            result.append(line)
    suffix = "".join("%s = %s\n" % (key, value) for key, value in changes.items() if value is not None)
    return "".join(result).rstrip("\n") + "\n" + suffix


def sender_maps(domain, sender):
    # Envelope-only canonical mapping: Reply-To is a header_sender in Postfix.
    escaped = re.escape(domain)
    envelope = "/^([^@]+@%s)$/ $1\n/./ %s\n" % (escaped, sender)
    # SMTP header checks touch only From/Sender. Existing Reply-To, recipients,
    # body and attachments are deliberately left to the ordinary mail pipeline.
    header = (r"/^(From|Sender):[[:space:]]*([^<>]*<)?[^@<>[:space:]]+@%s>?[[:space:]]*$/ DUNNO" % escaped
              + "\n" + r"/^(From|Sender):[[:space:]]*(.*<)[^<>]*>[[:space:]]*$/ REPLACE $1: ${2}%s>" % sender
              + "\n" + r"/^(From|Sender):/ REPLACE $1: %s" % sender + "\n")
    return envelope, header


def journey():
    correlation = {}
    try:
        state = json.loads(ONBOARDING_STATE.read_text())
        if (all(re.fullmatch(r"[a-f0-9]{32}", state.get(key, "")) for key in ("installation_id", "flow_id"))
                and state.get("flow_kind") in ("firstboot", "confconsole")):
            correlation = {key: state[key] for key in ("installation_id", "flow_id", "flow_kind")}
    except (OSError, ValueError, TypeError, AttributeError):
        pass
    return correlation


class Mail:
    def __init__(self):
        self.state = self.load()
        self.recover()

    def load(self):
        path = STATE_DIR / "state.json"
        try:
            value = json.loads(path.read_text())
            if not isinstance(value, dict):
                raise ValueError()
            return value
        except FileNotFoundError:
            return {}
        except (ValueError, OSError):
            raise Error("Saved Hub settings could not be read. Restore the state file before changing the connection.")

    def save(self):
        atomic(STATE_DIR / "state.json", json.dumps(self.state, sort_keys=True) + "\n")

    def recover(self):
        # A root-only journal covers interruption between config and state writes.
        transaction = self.state.get("applying")
        if not transaction:
            return
        atomic(POSTFIX_DIR / "main.cf", transaction["original"], 0o644)
        self.remove_files(transaction["files"])
        run(["postfix", "reload"])
        self.state.pop("applying")
        self.save()

    def guard(self):
        """A fresh automatic enrollment queues SMTP while Hub is unavailable."""
        text = (POSTFIX_DIR / "main.cf").read_text()
        values = cf_entries(text)
        if self.state.get("active") or values.get("relayhost") or self.state.get("guard"):
            return
        old = values.get("defer_transports")
        held = " ".join(filter(None, (old, "smtp")))
        self.state["guard"] = {"before": old, "value": held}
        self.save()
        atomic(POSTFIX_DIR / "main.cf", change_cf(text, {"defer_transports": held}), 0o644)
        try:
            run(["postfix", "reload"])
        except Error:
            atomic(POSTFIX_DIR / "main.cf", text, 0o644)
            self.state.pop("guard", None)
            self.save()
            raise

    def linked(self):
        # Read older combined state until the next successful configuration change.
        config = self.state.get("link", self.state.get("active", {}))
        return config if config.get("mode") == "hub" else {}

    def apply(self, config, verify=True, restore_sender_policy=None):
        if verify:
            smtp_check(config)
        linked = self.linked()
        original = (POSTFIX_DIR / "main.cf").read_text()
        values = cf_entries(original)
        suffix = secrets.token_hex(8)
        sasl = POSTFIX_DIR / ("turnkey-mail-" + suffix + ".sasl")
        created = [sasl, Path(str(sasl) + ".db")]
        target = "[%s]:%s" % (config["smtp"]["host"], config["smtp"]["port"])
        options = {
            "relayhost": target, "smtp_sasl_auth_enable": "yes",
            "smtp_sasl_password_maps": "hash:" + str(sasl),
            "smtp_sasl_security_options": "noanonymous", "smtp_sasl_tls_security_options": "noanonymous",
            "smtp_tls_security_level": "secure", "smtp_tls_CAfile": "/etc/ssl/certs/ca-certificates.crt",
            "smtp_tls_wrappermode": "yes" if config["smtp"]["port"] == 465 else "no",
            "smtp_fallback_relay": "", "smtp_sasl_auth_soft_bounce": "yes",
        }
        if config["mode"] == "hub":
            envelope, header = sender_maps(config["domain"], config["sender"])
            for extension, content in (("sender", envelope), ("headers", header)):
                path = POSTFIX_DIR / ("turnkey-mail-" + suffix + "." + extension)
                atomic(path, content)
                created.append(path)
            options.update(sender_canonical_classes="envelope_sender",
                           sender_canonical_maps="regexp:" + str(created[2]),
                           smtp_header_checks="regexp:" + str(created[3]))
        elif self.state.get("active", {}).get("mode") == "hub":
            # Switching to custom SMTP restores the prior sender policy.
            for key in ("sender_canonical_classes", "sender_canonical_maps", "smtp_header_checks"):
                options[key] = self.state.get("baseline", {}).get(key)
        options.update(restore_sender_policy or {})
        if self.state.get("guard"):
            options["defer_transports"] = self.state["guard"]["before"]
        guarded_queue = config["mode"] == "hub" and self.state.get("guard")
        initial_options = dict(options)
        if guarded_queue:
            initial_options["defer_transports"] = guarded_queue["value"]
        baseline = dict(self.state.get("baseline", {}))
        for key in options:
            if key not in baseline:
                baseline[key] = values.get(key)
        if self.state.get("guard"):
            baseline["defer_transports"] = self.state["guard"]["before"]
        old_state = json.loads(json.dumps(self.state))
        self.state["applying"] = {"original": original, "files": [str(path) for path in created]}
        self.save()
        try:
            atomic(sasl, "%s %s:%s\n" % (target, config["smtp"]["username"], config["password"]))
            run(["postmap", "hash:" + str(sasl)])
            os.chmod(str(sasl) + ".db", 0o600)
            atomic(POSTFIX_DIR / "main.cf", change_cf(original, initial_options), 0o644)
            run(["postfix", "check"])
            run(["postfix", "reload"])
            if guarded_queue:
                # Canonical maps run at queue intake. Reprocess mail accepted
                # before these maps existed, keeping SMTP held until it is
                # queued for cleanup with the new sender policy installed.
                run(["postsuper", "-r", "ALL"])
                atomic(POSTFIX_DIR / "main.cf", change_cf(original, options), 0o644)
                run(["postfix", "reload"])
            self.state.update(active=config, baseline=baseline, owned_options=options,
                              owned_files=[str(path) for path in created])
            if config["mode"] == "hub" and not linked.get("control_secret"):
                self.state["link"] = dict(config)
            elif linked:
                self.state["link"] = linked
            if config["mode"] == "custom":
                self.state["mail_preference"] = "manual"
            self.state.pop("guard", None)
            self.state.pop("applying", None)
            self.save()
        except (Error, OSError):
            atomic(POSTFIX_DIR / "main.cf", original, 0o644)
            for path in created:
                path.unlink(missing_ok=True)
            with contextlib.suppress(Error):
                run(["postfix", "reload"])
            self.state = old_state
            self.save()
            raise Error("The mail service could not apply these settings. Existing settings were restored.", retryable=True)
        self.remove_files(old_state.get("owned_files", []))

    def remove_files(self, files):
        for name in files:
            path = Path(name)
            if path.parent == POSTFIX_DIR and path.name.startswith("turnkey-mail-"):
                path.unlink(missing_ok=True)

    def enroll(self, token=None, flow=None):
        if token:
            if not ENROLLMENT_TOKEN.fullmatch(token):
                raise Error("The connection code is invalid. Start a new connection.")
            if self.state.get("enrolled_token_hash") == hashlib.sha256(token.encode()).hexdigest() and self.linked():
                return self.status(refresh=False)
            if self.state.get("pending", {}).get("token") != token:
                self.state["pending"] = {"token": token, "control_secret": secrets.token_urlsafe(32),
                                         "dns_password": secrets.token_urlsafe(32),
                                         "flow": journey() if flow is None else flow}
            self.save()
        pending = self.state.get("pending")
        if not pending:
            return self.status(refresh=False)
        if self.state.get('mail_preference', 'hub') == 'hub':
            self.guard()
        try:
            result = request("enroll", {key: pending[key] for key in ("token", "control_secret", "dns_password")})
            if result.get("version") != 2 or not re.fullmatch(r"[a-f0-9]{32}", result.get("device_id", "")):
                raise Error("The Hub returned an invalid connection.")
            if self.linked().get("device_id") != result["device_id"]:
                HUB_EMAIL_CACHE.unlink(missing_ok=True)
                for key in ("smtp_secret", "backup_secret", "services"):
                    self.state.pop(key, None)
            self.state["link"] = dict(mode="hub", device_id=result["device_id"],
                                      control_secret=pending["control_secret"], dns_password=pending["dns_password"],
                                      flow=pending.get("flow", {}))
            self.state["enrolled_token_hash"] = hashlib.sha256(pending["token"].encode()).hexdigest()
            for key in ("pending", "pairing", "error"):
                self.state.pop(key, None)
            self.state.setdefault("mail_preference", "hub")
            self.save()
            return self.sync()
        except Error as error:
            self.state["error"] = {"code": error.code, "message": str(error)}
            if not error.retryable and self.state.get("pending"):
                self.state["pending"]["stopped"] = True
            self.save()
            if self.linked().get("control_secret"):
                return self.status(refresh=False)
            raise

    def retry(self):
        pending = self.state.get("pending")
        if pending and not pending.get("stopped"):
            return self.enroll()
        if self.state.get("error") and self.linked().get("control_secret"):
            return self.sync()
        return self.status(refresh=False)

    def sync(self):
        link = self.linked()
        if not link.get("control_secret"):
            return self.status(refresh=False)
        proof = {key: link[key] for key in ("device_id", "control_secret")}
        value = request("configure", proof)
        if value.get("device_id") != link["device_id"] or value.get("connection_status") != "linked":
            raise Error("The Hub returned an invalid connection.")
        registration = {}
        mail = value["mail"]
        if mail.get("enabled") and self.state.get("mail_preference", "hub") == "hub":
            self.guard()
            grant = mail.get("grant_id", "")
            if not re.fullmatch(r"[a-f0-9]{32}", grant):
                raise Error("The Hub returned an invalid email grant.")
            if self.state.get("smtp_secret", {}).get("grant_id") != grant:
                self.state["smtp_secret"] = {"grant_id": grant, "password": secrets.token_urlsafe(32)}
                self.save()
            if mail["status"] == "secret_required":
                registration.update(smtp_grant_id=grant, smtp_password=self.state["smtp_secret"]["password"])
        if registration:
            value = request("configure", dict(proof, **registration))
        account_email = value.get("account_email")
        if isinstance(account_email, str) and len(account_email) <= 254 and ADDRESS.fullmatch(account_email):
            HUB_EMAIL_CACHE.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
            atomic(HUB_EMAIL_CACHE, account_email + "\n")
        self.state["services"] = value
        dns = value["dns"]
        link["dyndns"] = {"hostname": dns["hostname"], "update_uri": dns["update_uri"]} if dns["enabled"] else None
        self.save()
        mail = value["mail"]
        if (self.state.get("mail_preference", "hub") == "hub" and mail.get("enabled") and
                mail["status"] in ("active", "pending_verification") and mail.get("smtp")):
            envelope = dict(version=1, device_id=link["device_id"], smtp=mail["smtp"],
                            domain=mail["domain"], sender=mail["sender"],
                            status="ready" if mail["status"] == "active" else "pending")
            config = checked_config(envelope, self.state["smtp_secret"]["password"])
            if self.state.get("active") != config:
                try:
                    self.apply(config)
                except Error as error:
                    mail["error"] = str(error)
        account_email = value.get("account_email")
        if isinstance(account_email, str) and len(account_email) <= 254 and ADDRESS.fullmatch(account_email):
            HUB_EMAIL_CACHE.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
            atomic(HUB_EMAIL_CACHE, account_email + "\n")
        self.state["services"] = value
        self.state.pop("error", None)
        self.save()
        return self.status(refresh=False)

    def credential(self, service):
        """Return service-scoped authorization; consumers own initialization."""
        link = self.linked()
        if service == "dns" and link.get("dns_password") and link.get("dyndns"):
            return {"credential": "hubdns2.%s.%s" % (link["device_id"], link["dns_password"]),
                    "hostname": link["dyndns"]["hostname"]}
        if not link.get("control_secret"):
            raise Error("Connect this server by running turnkey-hub, then try again.")
        proof = {key: link[key] for key in ("device_id", "control_secret")}
        value = request("configure", proof)
        if value.get("device_id") != link["device_id"] or value.get("connection_status") != "linked":
            raise Error("The Hub returned an invalid connection.")
        if service == "dns":
            dns = value["dns"]
            if not dns.get("enabled"):
                raise Error("Enable DNS for this server in the Hub, then run hubdns-init again.")
            secret = link.get("dns_password")
            if not secret:
                raise Error("The DNS credential is missing. Reconnect this server with turnkey-hub.")
            return {"credential": "hubdns2.%s.%s" % (link["device_id"], secret),
                    "hostname": dns["hostname"]}
        backup = value["backup"]
        status = backup.get("status")
        if status == "disabled":
            raise Error("Enable backups for this server in the Hub, then run tklbam-backup again.")
        if status not in ("active", "secret_required"):
            action = urllib.parse.urljoin(HUB, backup.get("action_url", "/backup/"))
            if urllib.parse.urlsplit(action).netloc != urllib.parse.urlsplit(HUB).netloc:
                action = HUB + "/backup/"
            message = ("Finish migrating your backup account" if status == "migration_required" else
                       "Finish backup setup and connect your AWS account")
            raise Error(message + " in the Hub:\n\n  " + action + "\n\nThen retry your command.")
        if not self.state.get("backup_secret"):
            self.state["backup_secret"] = secrets.token_urlsafe(32)
            self.save()
        secret = self.state["backup_secret"]
        result = request("configure", dict(proof, backup_secret=secret))
        if result.get("backup", {}).get("status") != "active":
            raise Error("Backup setup could not finish. Check backup settings in the Hub and try again.")
        return {"credential": "tklbam2.%s.%s" % (link["device_id"], secret)}

    def use_hub(self):
        self.state["mail_preference"] = "hub"
        self.save()
        return self.sync()

    def start(self):
        correlation = journey()
        result = request("device/start", dict(correlation, **appliance_description()))
        try:
            uri = urllib.parse.urlsplit(result["verification_uri"])
            if uri.scheme != "https" or uri.netloc != urllib.parse.urlsplit(HUB).netloc:
                raise ValueError()
            if not DEVICE_CODE.fullmatch(result["device_code"]):
                raise ValueError()
            if not re.fullmatch(r"[A-Z0-9]{4}[&-][A-Z0-9]{4}", result["user_code"]):
                raise ValueError()
            ttl = min(900, max(1, int(result["expires_in"])))
            interval = min(30, max(5, int(result["interval"])))
        except (KeyError, TypeError, ValueError):
            raise Error("The Hub returned an invalid connection request.")
        # Older helpers require an ampersand on the wire; display a dash locally.
        result["user_code"] = result["user_code"].replace("&", "-")
        self.state["pairing"] = {"device_code": result["device_code"], "expires": time.time() + ttl,
                                 "interval": interval, "next_poll": 0, "flow": correlation,
                                 "verification_uri": result["verification_uri"], "user_code": result["user_code"]}
        try:
            from libinithooks.hub_push import validated_notification
            notification = validated_notification(result.get("notification"))
            if notification:
                self.state["pairing"]["notification"] = notification
        except ImportError:
            pass
        self.save()
        return {key: result[key] for key in ("verification_uri", "user_code", "expires_in", "interval")}

    def poll(self, notified=False):
        pairing = self.state.get("pairing", {})
        if pairing.get("expires", 0) <= time.time():
            return {"status": "expired"}
        if not notified and pairing.get("next_poll", 0) > time.time():
            return {"status": "pending", "interval": pairing["interval"]}
        pairing["next_poll"] = time.time() + pairing["interval"]
        self.save()
        result = request("device/poll", {"device_code": pairing["device_code"]})
        if result.get("status") == "approved":
            token = result.get("enrollment_token", "")
            if not isinstance(token, str) or not ENROLLMENT_TOKEN.fullmatch(token):
                raise Error("The Hub returned an invalid connection approval.")
            return self.enroll(token, pairing.get("flow", {}))
        if result.get("status") not in ("pending", "expired", "denied"):
            raise Error("The Hub returned an invalid connection status.")
        return {"status": result["status"], "interval": pairing["interval"]}

    def status(self, refresh=True):
        link = self.linked()
        if refresh and link.get("control_secret"):
            return self.sync()
        active = self.state.get("active", {})
        services = self.state.get("services", {})
        return {"status": active.get("status", "linked" if link.get("control_secret") else "connecting" if self.state.get("pending") else "disconnected"),
                "connection_status": "linked" if link.get("control_secret") else "disconnected",
                "mode": active.get("mode", ""), "mail_preference": self.state.get("mail_preference", "hub"),
                "domain": active.get("domain", ""), "sender": active.get("sender", ""),
                "host": active.get("smtp", {}).get("host", ""), "services": services,
                "message": self.state.get("error", {}).get("message", self.state.get("legacy_dns", ""))}

    def legacy_dns(self):
        return (Path("/var/lib/hubdns/sub_apikey").exists()
                and os.access("/etc/cron.hourly/hubdns-update", os.X_OK))

    def dns_update(self):
        raise Error("Run hubdns-update to update DNS.")

    def reset_mail(self):
        return self.reset(mail_only=True)

    def reset(self, mail_only=False):
        original = (POSTFIX_DIR / "main.cf").read_text()
        values = cf_entries(original)
        restore = {key: self.state.get("baseline", {}).get(key)
                   for key, value in self.state.get("owned_options", {}).items() if values.get(key) == value}
        guard = self.state.get("guard")
        if guard and values.get("defer_transports") == guard["value"]:
            restore["defer_transports"] = guard["before"]
        try:
            if restore:
                atomic(POSTFIX_DIR / "main.cf", change_cf(original, restore), 0o644)
                run(["postfix", "check"])
                run(["postfix", "reload"])
        except Error:
            atomic(POSTFIX_DIR / "main.cf", original, 0o644)
            with contextlib.suppress(Error):
                run(["postfix", "reload"])
            raise
        self.remove_files(self.state.get("owned_files", []))
        names = () if mail_only else ("turnkey-mail-retry", "turnkey-mail-dns")
        for name in names:
            if RUNTIME:
                RUNTIME.schedule("dns-update" if name.endswith("dns") else "retry", False)
                continue
            timer = SYSTEMD_DIR / (name + ".timer")
            if timer.exists():
                run(["systemctl", "disable", "--now", name + ".timer"])
            for extension in ("timer", "service"):
                (SYSTEMD_DIR / (name + "." + extension)).unlink(missing_ok=True)
        if not RUNTIME:
            run(["systemctl", "daemon-reload"])
        credential = BACKUP_DIR / "sub_apikey"
        if not mail_only and credential.exists() and self.linked().get("device_id"):
            if credential.read_text().startswith("tklbam2." + self.linked()["device_id"] + "."):
                credential.unlink()
                (BACKUP_DIR / "credentials").unlink(missing_ok=True)
        retained = {key: value for key, value in self.state.items() if key not in ("active", "baseline", "owned_options", "owned_files", "guard", "applying")} if mail_only else {}
        if mail_only:
            retained["mail_preference"] = "manual"
        if retained and "enrolled_token_hash" in self.state:
            retained["enrolled_token_hash"] = self.state["enrolled_token_hash"]
        if not mail_only:
            HUB_EMAIL_CACHE.unlink(missing_ok=True)
        self.state = retained
        self.save()
        return {"status": "disconnected"}

    def test(self, recipient):
        if not ADDRESS.fullmatch(recipient) or len(recipient) > 254:
            raise Error("Enter a valid recipient email address.")
        config = self.state.get("active")
        if not config:
            raise Error("Connect a mail relay before sending a test.")
        sender = config.get("sender", "root@localhost")
        message = email.message.EmailMessage()
        message["From"] = sender
        message["To"] = recipient
        message["Subject"] = "TurnKey app email test"
        message.set_content("This message was submitted through your appliance's system mail service.\n")
        run(["/usr/sbin/sendmail", "-i", "-f", sender, "--", recipient], message.as_bytes())
        return {"status": "queued", "message": "Test message accepted by the local mail queue. Check the recipient inbox for delivery."}


def initialize_linked_dns(result):
    """Delegate first-link DNS setup to the optional service package."""
    if (not result.get("services", {}).get("dns", {}).get("enabled") or
            (HUBDNS_DIR / "sub_apikey").exists()):
        return
    try:
        version = subprocess.run(["dpkg-query", "-W", "-f=${Version}", "hubdns"],
                                 capture_output=True, text=True, timeout=5)
        if version.returncode:
            return
        compatible = subprocess.run(["dpkg", "--compare-versions", version.stdout.strip(), "ge", "1.5.0"],
                                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=5)
        if compatible.returncode:
            return
        initialized = subprocess.run(["/usr/bin/hubdns-init"], stdin=subprocess.DEVNULL,
                                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=30)
        if not initialized.returncode:
            return
    except (OSError, subprocess.TimeoutExpired):
        pass
    result["dns_setup_error"] = "Automatic DNS setup could not finish. Run hubdns-init to try again."


@contextlib.contextmanager
def locked():
    STATE_DIR.mkdir(mode=0o700, parents=True, exist_ok=True)
    os.chmod(STATE_DIR, 0o700)
    with (STATE_DIR / "lock").open("a") as lock:
        os.chmod(STATE_DIR / "lock", 0o600)
        fcntl.flock(lock, fcntl.LOCK_EX)
        yield


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    enroll = commands.add_parser("enroll")
    enroll.add_argument("--token-stdin", action="store_true", required=True)
    enroll.add_argument("--require-backup", action="store_true")
    credential = commands.add_parser("credential")
    credential.add_argument("service", choices=("dns", "backup"))
    credential.add_argument("--connect", action="store_true")
    connect = commands.add_parser("connect")
    mode = connect.add_mutually_exclusive_group()
    mode.add_argument("--start", action="store_true")
    mode.add_argument("--poll", action="store_true")
    connect.add_argument("--notified", action="store_true")
    custom = commands.add_parser("configure")
    custom.add_argument("--stdin", action="store_true", required=True)
    status = commands.add_parser("status")
    status.add_argument("--local", action="store_true")
    for name in ("reset", "reset-mail", "retry", "dns-update", "sync", "use-hub"):
        commands.add_parser(name)
    test = commands.add_parser("test")
    test.add_argument("--recipient", required=True)
    for command in commands.choices.values():
        command.add_argument("--json", action="store_true")
    args = parser.parse_args()
    if args.command == "connect" and args.notified and not args.poll:
        parser.error("--notified requires --poll")
    try:
        if os.geteuid() != 0:
            raise Error("Run this command as root.")
        os.umask(0o077)
        if args.command == "dns-update":
            os.execv("/usr/bin/hubdns-update", ["hubdns-update"])
        if args.command == "credential" and args.connect and not Mail().linked().get("control_secret"):
            if not sys.stdin.isatty():
                raise Error("Connect this server by running turnkey-hub, then try again.")
            # Linking owns its own lock and interactive output. Keep the exported
            # credential on stdout for the calling service to consume privately.
            result = subprocess.call([HELPER], stdout=sys.stderr)
            if result:
                raise Error("Connection was not completed. Run turnkey-hub to try again.")
        with locked():
            mail = Mail()
            was_linked = bool(mail.linked().get("control_secret"))
            if args.command == "enroll":
                result = mail.enroll(sys.stdin.read(512).strip())
            elif args.command == "credential":
                result = mail.credential(args.service)
            elif args.command == "configure":
                try:
                    raw = sys.stdin.read(MAX_JSON + 1)
                    if len(raw) > MAX_JSON:
                        raise ValueError()
                    config = json.loads(raw)
                    mail.apply(checked_config(config, config["password"], custom=True))
                    for field in ("pending", "pairing", "error"):
                        mail.state.pop(field, None)
                    mail.save()
                    result = mail.status(refresh=False)
                except (ValueError, KeyError, TypeError):
                    raise Error("The custom SMTP settings are invalid.")
            elif args.command == "connect":
                result = mail.poll(notified=args.notified) if args.poll else mail.start()
            elif args.command == "dns-update":
                result = mail.dns_update()
            elif args.command == "test":
                result = mail.test(args.recipient)
            elif args.command == "status":
                result = mail.status(refresh=not args.local)
            else:
                result = getattr(mail, args.command.replace("-", "_"))()
        if args.command == "connect" and not args.start and not args.poll:
            print("Open %s\nConnection code: %s\nSign in or create a free Hub account, then approve this appliance.\nWaiting for approval (Ctrl-C to cancel)..." %
                  (result["verification_uri"], result["user_code"]), flush=True)
            while True:
                time.sleep(result.get("interval", 5))
                with locked():
                    result = Mail().poll()
                if result.get("status") != "pending" or result.get("domain"):
                    break
            if result.get("status") in ("expired", "denied"):
                raise Error(ERRORS[result["status"]], code=result["status"])
        if (args.command in ("enroll", "connect", "retry") and not was_linked and
                result.get("connection_status") == "linked"):
            initialize_linked_dns(result)
        if args.command == "enroll" and args.require_backup:
            run(["tklbam-init"], timeout=60)
        if result.get("dns_setup_error") and not args.json:
            print(result["dns_setup_error"], file=sys.stderr)
        if args.json or args.command == "credential":
            print(json.dumps(result, sort_keys=True))
        elif result.get("message"):
            print(result["message"])
        else:
            print("%s connection: %s%s" % ("Mail" if args.command == "configure" else "Hub", result.get("status", "waiting"),
                  (" (" + result["domain"] + ")") if result.get("domain") else ""))
        return 0
    except Error as error:
        if args.json:
            print(json.dumps({"status": "error", "error": error.code, "message": str(error)}))
        else:
            print(str(error), file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("Linking cancelled. Existing service settings were kept.", file=sys.stderr)
        return 130
    except (OSError, ValueError, KeyError, TypeError):
        # Never emit tracebacks that could contain tokens or configuration.
        label = "Mail setup" if args.command in ("configure", "test", "reset-mail") else "Hub services setup"
        print(label + " could not finish. Check the system services and try again.", file=sys.stderr)
        return 1


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