stack/src/mail/droplet.py

"""DigitalOcean droplet lifecycle for the Maddy mail server.

All commands are touchless / stateless / idempotent:

* **Touchless** — droplet name = FQDN so DO auto-sets the PTR record;
  DNS through CF API; SSH uses account-registered keys; no panel.
* **Stateless** — DO is the source of truth for droplet metadata
  (queried by tag ``stack-mail`` on every command). Resend is the
  source of truth for domain verification + DKIM. Cloudflare is the
  source of truth for DNS. ``.state/mail/`` is a cache only.
* **Idempotent** — ``up`` adopts an existing droplet, never duplicates.
  DNS records are upserted (no error on existing). The smarthost block
  is fenced and replaced in place. Mailbox passwords use ``maddy creds
  password`` first, then fall through to ``create`` + ``imap-acct create``
  for first-time provision.
"""

from __future__ import annotations

import json
import secrets
import string
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path

from mail._helpers import (
    DOMAIN,
    DROPLET_SIZE,
    HOSTNAME,
    IMAGE,
    ROOT,
    STATE_DIR,
    TAG,
    env,
    ok,
    step,
    warn,
)

DROPLET_JSON = STATE_DIR / "droplet.json"
CREDS_JSON = STATE_DIR / "credentials.json"
GIT_MAILER_ENV = ROOT / ".state" / "git" / "mailer.env"


# ── Small utilities ─────────────────────────────────────────────


def _save_json(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(data, indent=2) + "\n")
    tmp.rename(path)


def _load_json(path: Path) -> dict:
    return json.loads(path.read_text()) if path.exists() else {}


def _gen_password(n: int = 24) -> str:
    alphabet = string.ascii_letters + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(n))


def _do_client():
    import logging

    import pydo

    logging.getLogger("azure").setLevel(logging.WARNING)
    token = env("DIGITAL_OCEAN_PAT")
    if not token:
        raise RuntimeError("DIGITAL_OCEAN_PAT not set in stack/.env")
    return pydo.Client(token=token)


def _get_ssh_key_ids(client) -> list[int]:
    return [k["id"] for k in client.ssh_keys.list().get("ssh_keys", [])]


def discover_droplet(client) -> dict | None:
    """Find the mail droplet by tag. Authoritative source of truth."""
    for d in client.droplets.list(tag_name=TAG).get("droplets") or []:
        public_ip = next(
            (
                n["ip_address"]
                for n in d.get("networks", {}).get("v4", [])
                if n["type"] == "public"
            ),
            None,
        )
        return {
            "id": d["id"],
            "name": d["name"],
            "region": (
                d["region"]["slug"]
                if isinstance(d.get("region"), dict)
                else d.get("region")
            ),
            "hostname": d["name"],
            "public_ip": public_ip,
            "status": d.get("status"),
            "created": d.get("created_at"),
        }
    return None


def _wait_for_active(client, droplet_id: int, *, timeout: int = 300) -> str:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        d = client.droplets.get(droplet_id)["droplet"]
        if d["status"] == "active":
            for n in d["networks"]["v4"]:
                if n["type"] == "public":
                    return n["ip_address"]
        time.sleep(5)
    raise RuntimeError(f"droplet {droplet_id} never went active")


def _ssh_args() -> list[str]:
    """Use the user's default SSH identity. ``STACK_SSH_KEY`` overrides."""
    override = env("STACK_SSH_KEY")
    return ["-i", override] if override else []


def _ssh(
    public_ip: str, *cmd: str, check: bool = True, input_: str | None = None
) -> subprocess.CompletedProcess:
    args = [
        "ssh",
        *_ssh_args(),
        "-o",
        "StrictHostKeyChecking=no",
        "-o",
        "UserKnownHostsFile=/dev/null",
        "-o",
        "ConnectTimeout=15",
        f"root@{public_ip}",
        *cmd,
    ]
    return subprocess.run(
        args, check=check, capture_output=True, text=True, input=input_
    )


# ── Cloud-init builder ──────────────────────────────────────────


def _cloud_init() -> str:
    """Build the user_data script. Mailbox passwords are NOT included —
    cloud-init only stands up Maddy + cert; mailboxes get seeded later
    over SSH via `seed_mailboxes`. This keeps secrets off the DO
    metadata service, which is queryable from inside the droplet by
    any process with HTTP access."""
    script = (ROOT / "infra" / "droplets" / "mail-setup.sh").read_text()
    tpl = (ROOT / "infra" / "mail" / "maddy.conf.tpl").read_text()
    script = script.replace("__MADDY_CONF_TPL_PLACEHOLDER__", tpl)

    cf_token = env("CF_API_TOKEN") or env("CLOUDFLARE_API_TOKEN")
    if not cf_token:
        raise RuntimeError(
            "CF_API_TOKEN not set — required for DNS-01 cert "
            "issuance inside cloud-init."
        )
    # certbot-dns-cloudflare credentials.ini expects the var named
    # CLOUDFLARE_API_TOKEN regardless of which side our env had it.
    return f"""#!/bin/bash
set -euo pipefail
export HOSTNAME={HOSTNAME!r}
export PRIMARY_DOMAIN={DOMAIN!r}
export CLOUDFLARE_API_TOKEN={cf_token!r}

{script}
"""


# ── Lifecycle commands ─────────────────────────────────────────


def up(region: str = "nyc3") -> dict:
    """Adopt-or-create the mail droplet. Idempotent."""
    step("Mail droplet up")
    client = _do_client()

    existing = discover_droplet(client)
    if existing:
        ok(f"adopting droplet {existing['id']} at {existing['public_ip']}")
        _save_json(DROPLET_JSON, existing)
        return existing

    ssh_key_ids = _get_ssh_key_ids(client)
    if not ssh_key_ids:
        raise RuntimeError(
            "No SSH keys on the DO account — register one at "
            "https://cloud.digitalocean.com/account/security so this and "
            "future machines can manage the droplet."
        )

    step(f"Creating droplet {HOSTNAME} in {region}")
    resp = client.droplets.create(
        body={
            "name": HOSTNAME,  # name=FQDN → DO auto-PTR
            "region": region,
            "size": DROPLET_SIZE,
            "image": IMAGE,
            "user_data": _cloud_init(),
            "tags": [TAG],
            "ssh_keys": ssh_key_ids,
        }
    )
    droplet_id = resp["droplet"]["id"]
    public_ip = _wait_for_active(client, droplet_id)

    info = {
        "id": droplet_id,
        "name": HOSTNAME,
        "region": region,
        "hostname": HOSTNAME,
        "public_ip": public_ip,
        "created": datetime.now(timezone.utc).isoformat(),
    }
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    _save_json(DROPLET_JSON, info)
    ok(f"droplet live at {public_ip}")
    return info


def down(*, confirm: bool = False) -> None:
    """Destroy the droplet. Idempotent no-op if none exists."""
    client = _do_client()
    existing = discover_droplet(client)
    if not existing:
        ok("no mail droplet exists — nothing to do")
        for f in (DROPLET_JSON, CREDS_JSON):
            f.unlink(missing_ok=True)
        return
    if not confirm:
        raise RuntimeError(f"refusing to destroy {existing['name']} without --yes")
    try:
        client.droplets.destroy(existing["id"])
        ok(f"destroyed droplet {existing['id']}")
    except Exception as e:  # noqa: BLE001
        warn(f"destroy failed (maybe already gone): {e}")
    for f in (DROPLET_JSON, CREDS_JSON):
        f.unlink(missing_ok=True)


def status() -> dict | None:
    """Return droplet state from DO + PTR check."""
    import socket

    client = _do_client()
    meta = discover_droplet(client)
    if not meta:
        return None
    _save_json(DROPLET_JSON, meta)
    try:
        meta["ptr"] = socket.gethostbyaddr(meta["public_ip"])[0]
        meta["ptr_match"] = meta["ptr"] == meta["hostname"]
    except (socket.herror, socket.gaierror):
        meta["ptr"] = None
        meta["ptr_match"] = False
    return meta


# ── DNS + smarthost orchestration ──────────────────────────────


def apply_dns(claim_apex: bool = False) -> None:
    """Upsert mail DNS in Cloudflare. Idempotent.

    Apex MX/SPF/DMARC require ``claim_apex=True`` — default skips
    them so we don't trample an existing provider (Proton, Fastmail,
    Google Workspace) at the bare domain.
    """
    import httpx

    from mail.cloudflare import ensure_mail_dns

    meta = discover_droplet(_do_client())
    if not meta:
        raise RuntimeError("no mail droplet exists")
    with httpx.Client(timeout=30) as client:
        ensure_mail_dns(
            client,
            hostname=meta["hostname"],
            public_ip=meta["public_ip"],
            domain=DOMAIN,
            claim_apex=claim_apex,
        )


def export_dkim() -> None:
    """Read maddy's generated DKIM key and publish to Cloudflare. Idempotent."""
    import httpx

    from mail.cloudflare import ensure_dkim_record

    meta = discover_droplet(_do_client())
    if not meta:
        raise RuntimeError("no mail droplet exists")
    r = _ssh(
        meta["public_ip"],
        f"cat /srv/mail/data/dkim_keys/{DOMAIN}_default.dns",
        check=False,
    )
    if r.returncode != 0 or not r.stdout.strip():
        raise RuntimeError(
            f"could not read DKIM key: {r.stderr.strip()}\n"
            "Cloud-init likely still running; wait a minute and retry."
        )
    with httpx.Client(timeout=30) as client:
        ensure_dkim_record(client, domain=DOMAIN, rdata=r.stdout.strip())


def attach_smarthost() -> None:
    """Register at Postmark, push DNS, install smarthost block on droplet.

    Postmark (not Resend) because Resend's free tier caps at 1 domain
    and corwins.media already used it. Postmark allows multiple
    domains free, transactional-only, has an account+server token
    model that mints us a per-Server SMTP credential.
    """
    import httpx

    from mail.postmark import (
        ensure_postmark_domain,
        ensure_postmark_server,
        publish_postmark_dns,
        verify_postmark_domain,
    )

    meta = discover_droplet(_do_client())
    if not meta:
        raise RuntimeError("no mail droplet exists — `stack mail provision` first")
    ip = meta["public_ip"]

    step("Provisioning Postmark Server + Domain")
    server_token = ensure_postmark_server("fhirworx")
    domain_info = ensure_postmark_domain(DOMAIN)
    with httpx.Client(timeout=30) as client:
        publish_postmark_dns(domain_info, client)
    # Trigger verification — Postmark won't auto-verify after Domain creation.
    verify_postmark_domain(domain_info["ID"])

    # Postmark SMTP: 25 / 587 (STARTTLS) / 2525 (also STARTTLS, alt
    # port for hosts that block 587). DO blocks 25/465/587 silently;
    # 2525 is the universally-open one so we use it. Maddy's `tcp://`
    # scheme means connect plaintext + negotiate STARTTLS — exactly
    # what Postmark's 2525 speaks.
    block = (
        "# SMARTHOST-BEGIN managed by `stack mail` — do not edit by hand.\n"
        "target.smtp outbound_smarthost {\n"
        "    targets tcp://smtp.postmarkapp.com:2525\n"
        f'    auth plain "{server_token}" "{server_token}"\n'
        "}\n"
        "# SMARTHOST-END\n"
    )

    step("Installing smarthost block on droplet")
    # Maddy errors out if `outbound_smarthost` is defined twice, so the
    # patch always strips:
    #   1. The fenced SMARTHOST-BEGIN/END region (prior install)
    #   2. Any unfenced `target.smtp outbound_smarthost { … }` (the
    #      stub from the base template, written by cloud-init)
    # Then prepends the fresh fenced block. Idempotent + state-agnostic.
    py_script = (
        "import re, pathlib\n"
        "BLOCK = " + repr(block) + "\n"
        "CONF = pathlib.Path('/srv/mail/data/maddy.conf')\n"
        "txt = CONF.read_text()\n"
        "txt = re.sub(r'# SMARTHOST-BEGIN.*?# SMARTHOST-END\\n', '', txt, flags=re.DOTALL)\n"
        "txt = re.sub(r'target\\.smtp outbound_smarthost \\{[^{}]*\\}\\n*', '', txt)\n"
        "txt = BLOCK + txt\n"
        "CONF.write_text(txt)\n"
    )
    # Pipe via stdin — ssh's argv-quoting otherwise mangles multi-line Python.
    subprocess.run(
        [
            "ssh",
            *_ssh_args(),
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            f"root@{ip}",
            "python3 -",
        ],
        input=py_script,
        text=True,
        check=True,
        capture_output=True,
    )
    _ssh(ip, "docker", "restart", "mail")
    ok("smarthost attached — outbound via smtp.postmarkapp.com:2525")


# ── Mailbox + Gitea wiring ─────────────────────────────────────


# Mailboxes the orchestrator always seeds. Full addresses so each
# can sit on whichever domain makes sense:
#   apex (fhirworx.io) for outbound branding (Gitea From=)
#   subdomain (mail.fhirworx.io) for inbound when apex MX is owned
#     by another provider (e.g. Proton at the apex)
DEFAULT_MAILBOXES = (
    f"postmaster@{DOMAIN}",
    f"postmaster@{HOSTNAME}",  # DMARC rua/ruf destination; apex MX is Proton
    f"git@{DOMAIN}",  # Gitea's outbound SMTP submission AS this user
    f"cmsupdates@{HOSTNAME}",  # CMS list signups — MX must reach our Maddy
)


def _addr_to_key(addr: str) -> str:
    """Cache-key for an address. ``user@apex`` → ``user`` (back-compat for
    the older single-domain credentials.json layout); subdomain
    addresses get the full address as the key so they're unambiguous."""
    if addr.endswith(f"@{DOMAIN}"):
        return addr.split("@", 1)[0]
    return addr


def _maddy_imap_accts(ip: str) -> set[str]:
    """Set of existing IMAP account addresses on the droplet."""
    r = _ssh(ip, "docker", "exec", "mail", "maddy", "imap-acct", "list")
    return {line.strip() for line in r.stdout.splitlines() if line.strip()}


def _ensure_imap_acct(ip: str, addr: str, existing: set[str]) -> None:
    """Create the IMAP mailbox for ``addr`` if it doesn't already exist.

    Maddy keeps SMTP credentials (creds) and IMAP storage (imap-acct)
    in separate stores. Having one without the other means auth
    succeeds but delivered mail has nowhere to land — which silently
    discards e.g. DMARC aggregate reports. Always pair them.

    ``existing`` is mutated so callers can list once and reuse the set.
    """
    if addr in existing:
        return
    _ssh(ip, "docker", "exec", "mail", "maddy", "imap-acct", "create", addr)
    existing.add(addr)


def rotate_creds(user_or_addr: str) -> str:
    """Mint new password and push to droplet. ``user_or_addr`` may be
    ``user`` (assumed at apex) or full ``user@host``.

    maddy subcommands:
        creds password   → change password on existing account
        creds create     → create new account (also needs imap-acct create)
    Earlier versions of this function called ``creds set``, which doesn't
    exist in maddy — it silently failed and fell through to ``creds
    create`` which then errored on the existing account. Net result:
    local credentials.json got the new password but the droplet kept
    the old one. The poller then loops forever on Invalid credentials.
    """
    addr = user_or_addr if "@" in user_or_addr else f"{user_or_addr}@{DOMAIN}"
    meta = discover_droplet(_do_client())
    if not meta:
        raise RuntimeError("no mail droplet exists")
    ip = meta["public_ip"]
    new_pw = _gen_password()

    # Try password change first (account exists). If the account
    # doesn't exist yet, create it.
    r = _ssh(
        ip,
        "docker",
        "exec",
        "mail",
        "maddy",
        "creds",
        "password",
        "--password",
        new_pw,
        addr,
        check=False,
    )
    if r.returncode != 0:
        _ssh(
            ip,
            "docker",
            "exec",
            "mail",
            "maddy",
            "creds",
            "create",
            "--password",
            new_pw,
            addr,
        )
    # IMAP mailbox may be missing even when SMTP creds exist (saw this
    # in production with postmaster@mail.fhirworx.io after a partial
    # earlier provision). Ensure it independently.
    _ensure_imap_acct(ip, addr, _maddy_imap_accts(ip))
    creds = _load_json(CREDS_JSON)
    # Save under both forms so lookup works whether the caller passes
    # the bare key or the full address.
    creds[_addr_to_key(addr)] = new_pw
    creds[addr] = new_pw
    _save_json(CREDS_JSON, creds)
    CREDS_JSON.chmod(0o600)
    ok(f"rotated {addr}")
    return new_pw


def seed_mailboxes(addrs: tuple[str, ...] = DEFAULT_MAILBOXES) -> None:
    """Ensure each address exists on droplet with cache-matching password
    AND has a matching IMAP mailbox.

    Idempotent: when an address has a cached password, push it to the
    droplet (no-op if matches, replace if differs, create if missing).
    When no cached password exists, mint and push a fresh one. Either
    way, the IMAP mailbox is verified to exist (independently of the
    SMTP creds store — they can drift).
    """
    meta = discover_droplet(_do_client())
    if not meta:
        raise RuntimeError("no mail droplet exists")
    ip = meta["public_ip"]
    creds = _load_json(CREDS_JSON)
    imap_accts = _maddy_imap_accts(ip)

    for addr in addrs:
        key = _addr_to_key(addr)
        pw = creds.get(key)
        if pw:
            r = _ssh(
                ip,
                "docker",
                "exec",
                "mail",
                "maddy",
                "creds",
                "password",
                "--password",
                pw,
                addr,
                check=False,
            )
            if r.returncode == 0:
                _ensure_imap_acct(ip, addr, imap_accts)
                ok(f"mailbox {addr}: cached password in sync with droplet")
                continue
            _ssh(
                ip,
                "docker",
                "exec",
                "mail",
                "maddy",
                "creds",
                "create",
                "--password",
                pw,
                addr,
            )
            _ensure_imap_acct(ip, addr, imap_accts)
            ok(f"mailbox {addr}: created on droplet from cached password")
        else:
            rotate_creds(addr)
            imap_accts.add(addr)


def write_git_mailer_env() -> bool:
    """Write the Gitea SMTP env file. Returns True if content changed.

    Loaded by the ``git`` service in compose.yml as an optional
    ``env_file``. Once changed, recreate the container with
    ``docker compose up -d git`` to apply (a plain restart doesn't
    re-read compose).
    """
    if not CREDS_JSON.exists() or not DROPLET_JSON.exists():
        warn("git mailer: mail droplet not provisioned — skipping")
        return False
    creds = _load_json(CREDS_JSON)
    droplet = _load_json(DROPLET_JSON)
    if "git" not in creds:
        warn("git mailer: no `git` mailbox creds — `stack mail rotate-creds git`")
        return False

    hostname = droplet.get("hostname") or HOSTNAME
    desired = (
        "GITEA__mailer__ENABLED=true\n"
        "GITEA__mailer__PROTOCOL=smtps\n"
        f"GITEA__mailer__SMTP_ADDR={hostname}\n"
        "GITEA__mailer__SMTP_PORT=465\n"
        f"GITEA__mailer__USER=git@{DOMAIN}\n"
        f"GITEA__mailer__PASSWD={creds['git']}\n"
        f"GITEA__mailer__FROM=git@{DOMAIN}\n"
    )
    GIT_MAILER_ENV.parent.mkdir(parents=True, exist_ok=True)
    if GIT_MAILER_ENV.exists() and GIT_MAILER_ENV.read_text() == desired:
        return False
    tmp = GIT_MAILER_ENV.with_suffix(".tmp")
    tmp.write_text(desired)
    tmp.rename(GIT_MAILER_ENV)
    return True


# ── End-to-end orchestrator ────────────────────────────────────


def provision(region: str = "nyc3", wait_seconds: int = 360) -> None:
    """End-to-end idempotent: up → wait → DNS → DKIM → seed → smarthost → git env."""
    up(region=region)
    meta = _load_json(DROPLET_JSON)
    public_ip = meta["public_ip"]

    step("Waiting for Maddy + DKIM keypair (cloud-init)")
    deadline = time.monotonic() + wait_seconds
    dkim_path = f"/srv/mail/data/dkim_keys/{DOMAIN}_default.dns"
    while time.monotonic() < deadline:
        r = _ssh(public_ip, f"test -f {dkim_path} && cat {dkim_path}", check=False)
        if r.returncode == 0 and r.stdout.strip():
            ok("DKIM key present on droplet")
            break
        time.sleep(15)  # pragma: no cover — timing-dependent wait loop
    else:
        warn("DKIM key never appeared in time — `stack mail dkim-export` later")

    apply_dns()
    try:
        export_dkim()
    except Exception as e:  # noqa: BLE001
        warn(f"DKIM step skipped: {e}")

    # Seed mailboxes BEFORE smarthost: the smarthost block reads
    # `git`'s creds when it writes the env file later, and an early
    # `seed_mailboxes` run also surfaces auth errors before we paint
    # over the maddy config with the smarthost block.
    seed_mailboxes()
    attach_smarthost()

    if write_git_mailer_env():
        ok("git mailer env written → recreate gitea: `docker compose up -d git`")
    else:
        ok("git mailer env unchanged or skipped")
    ok(f"mail droplet provisioned at {HOSTNAME} ({public_ip})")