netadmin/agents/mail_agent.py

#!/usr/bin/env python3
"""
mail_agent.py — NetAdmin Mail Agent (PostfixAdmin MySQL schema)
Runs on the Postfix/Dovecot server.
Subscribes to Redis mail.commands, writes to the existing PostfixAdmin
MariaDB database, then publishes ACK.

Schema targeted:
  domain   (domain, description, aliases, mailboxes, maxquota, quota,
             transport, backupmx, created, modified, active)
  mailbox  (username, password, name, maildir, quota, local_part,
             domain, created, modified, active)
  alias    (address, goto, domain, created, modified, active)

Install on the mail server:
    pip install redis pyyaml PyMySQL
    python3 mail_agent.py --config /etc/netadmin/config.yaml

Postfix/Dovecot already query MySQL directly — no file reload needed
for account changes. We call 'postfix reload' only when domain config
changes (transport, etc.).
"""

import argparse
import json
import logging
import ssl
import sys
import subprocess
import datetime
from pathlib import Path

import redis
import yaml

try:
    import pymysql
    import pymysql.cursors
except ImportError:
    print("ERROR: PyMySQL not installed. Run: pip install PyMySQL")
    sys.exit(1)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [mail-agent] %(levelname)s %(message)s"
)
log = logging.getLogger("mail-agent")

NOW_DEFAULT = datetime.datetime(2000, 1, 1)


# ──────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────

def load_config(path: str) -> dict:
    with open(path) as f:
        raw = yaml.safe_load(f) or {}

    cfg = {}

    # Section mail_agent
    mail_cfg = raw.get("mail_agent", {})
    # Map nested mail_agent.db → flat db_* keys
    db = mail_cfg.get("db", {})
    if db:
        cfg["db_host"]     = db.get("host",     "localhost")
        cfg["db_port"]     = int(db.get("port", 3306))
        cfg["db_name"]     = db.get("name",     "netadmin_mail")
        cfg["db_user"]     = db.get("user",     "postfix")
        cfg["db_password"] = db.get("password", "")
    postfix = mail_cfg.get("postfix", {})
    if postfix:
        cfg["postfix_reload_cmd"] = postfix.get("reload_cmd", "postfix reload")
    dkim = mail_cfg.get("dkim", {})
    if dkim:
        cfg["dkim_dir"]      = dkim.get("dir",      "/var/lib/rspamd/dkim")
        cfg["dkim_selector"] = dkim.get("selector",  "mail")

    # Section redis → mapped to redis_* keys
    redis_cfg = raw.get("redis", {})
    if redis_cfg:
        cfg["redis_host"]     = redis_cfg.get("host",     "localhost")
        cfg["redis_port"]     = int(redis_cfg.get("port", 6380))
        cfg["redis_password"] = redis_cfg.get("password", "")
        cfg["redis_db"]       = int(redis_cfg.get("db",   0))
        if redis_cfg.get("ssl"):
            cfg["tls"] = {
                "enabled":     True,
                "ca_cert":     redis_cfg.get("ssl_ca_cert",  ""),
                "client_cert": redis_cfg.get("ssl_certfile", ""),
                "client_key":  redis_cfg.get("ssl_keyfile",  ""),
                "check_hostname": redis_cfg.get("check_hostname", False),
            }

    return cfg


DEFAULT_CFG = {
    # MariaDB connection
    "db_host":     "localhost",
    "db_port":     3306,
    "db_name":     "postfix",
    "db_user":     "postfix",
    "db_password": "",
    # Postfix reload (only needed for domain-level changes)
    "postfix_reload_cmd": "postfix reload",
    # Redis
    "redis_host":    "localhost",
    "redis_port":    6380,
    "redis_password": "",
    "redis_db":      0,
    "tls": {
        "enabled":        False,
        "ca_cert":        "/etc/netadmin/certs/ca.crt",
        "client_cert":    "/etc/netadmin/certs/mail-agent.crt",
        "client_key":     "/etc/netadmin/certs/mail-agent.key",
        "check_hostname": True,
    },
}


# ──────────────────────────────────────────────
# mTLS helper
# ──────────────────────────────────────────────

def build_redis_ssl_context(tls_cfg: dict):
    if not tls_cfg.get("enabled", False):
        return None
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.verify_mode = ssl.CERT_REQUIRED
    ctx.load_verify_locations(cafile=tls_cfg["ca_cert"])
    ctx.load_cert_chain(
        certfile=tls_cfg["client_cert"],
        keyfile=tls_cfg["client_key"]
    )
    ctx.minimum_version = ssl.TLSVersion.TLSv1_2
    if not tls_cfg.get("check_hostname", True):
        ctx.check_hostname = False
    return ctx


# ──────────────────────────────────────────────
# DB helper
# ──────────────────────────────────────────────

def get_db(cfg: dict):
    return pymysql.connect(
        host=cfg["db_host"],
        port=int(cfg["db_port"]),
        user=cfg["db_user"],
        password=cfg["db_password"],
        database=cfg["db_name"],
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
        autocommit=False,
    )


def now() -> str:
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")


# ──────────────────────────────────────────────
# Command handlers
# ──────────────────────────────────────────────

def _sync_domain_row(cur, domain_name: str, config: dict, now) -> bool:
    """Insert or update the domain row. Returns True if inserted."""
    description  = config.get("description", "")
    max_accounts = int(config.get("max_accounts", 0))
    max_quota_mb = int(config.get("max_quota_mb", 0))
    mb_quota_mb  = int(config.get("mb_quota_mb", 0))
    active       = int(config.get("active", True))

    cur.execute("SELECT domain FROM domain WHERE domain=%s", (domain_name,))
    if cur.fetchone():
        cur.execute(
            """UPDATE domain SET
                   description=%s, max_accounts=%s, max_quota_mb=%s,
                   mb_quota_mb=%s, active=%s, modified=%s
               WHERE domain=%s""",
            (description, max_accounts, max_quota_mb, mb_quota_mb,
             active, now(), domain_name)
        )
        log.info(f"Updated domain: {domain_name}")
        return False
    else:
        cur.execute(
            """INSERT INTO domain
                   (domain, description, max_accounts, max_quota_mb,
                    mb_quota_mb, active, created, modified)
               VALUES (%s,%s,%s,%s,%s,%s,%s,%s)""",
            (domain_name, description, max_accounts, max_quota_mb,
             mb_quota_mb, active, now(), now())
        )
        log.info(f"Inserted domain: {domain_name}")
        return True


def _sync_mailboxes(cur, domain_name: str, accounts: list, now):
    """Sync mailbox rows for a domain — insert, update, delete."""
    cur.execute("SELECT username FROM mailbox WHERE domain=%s", (domain_name,))
    existing_users = {r["username"] for r in cur.fetchall()}
    incoming_users = {a["username"] for a in accounts}

    # Delete removed accounts
    for username in existing_users - incoming_users:
        cur.execute("DELETE FROM mailbox WHERE username=%s", (username,))
        cur.execute("DELETE FROM alias WHERE address=%s AND goto=%s", (username, username))
        log.info(f"Deleted mailbox: {username}")

    # Insert or update
    for acc in accounts:
        username   = acc["username"]
        password   = acc.get("password_hash", "")
        name       = acc.get("name", "")
        acc_active = int(acc.get("active", True))

        if username in existing_users:
            update_fields = {
                "name":     name,
                "quota_mb": int(acc.get("quota", 1024)),
                "active":   acc_active,
                "modified": now(),
            }
            if password:
                update_fields["password"] = password
            set_clause = ", ".join(f"{k}=%s" for k in update_fields)
            cur.execute(
                f"UPDATE mailbox SET {set_clause} WHERE username=%s",
                (*update_fields.values(), username)
            )
        else:
            cur.execute(
                """INSERT INTO mailbox
                    (username, password, name, domain, quota_mb,
                     active, created, modified)
                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s)""",
                (username, password, name, domain_name,
                 int(acc.get("quota", 1024)),
                 acc_active, now(), now())
            )
            # Self-alias required for virtual_alias_maps
            cur.execute("SELECT address FROM alias WHERE address=%s", (username,))
            if not cur.fetchone():
                cur.execute(
                    """INSERT INTO alias
                        (address, goto, domain, active, created, modified)
                       VALUES (%s,%s,%s,1,%s,%s)""",
                    (username, username, domain_name, now(), now())
                )
            log.info(f"Created mailbox + self-alias: {username}")

    return incoming_users


def _sync_aliases(cur, domain_name: str, aliases: list, incoming_users: set, now):
    """Sync alias rows for a domain — upsert and delete stale entries."""
    cur.execute("SELECT address FROM alias WHERE domain=%s", (domain_name,))
    existing_aliases = {r["address"] for r in cur.fetchall()}
    incoming_alias_addresses = {a["address"] for a in aliases}

    # Remove stale aliases (keep self-aliases — managed with mailbox lifecycle)
    for addr in existing_aliases - incoming_alias_addresses - incoming_users:
        cur.execute("DELETE FROM alias WHERE address=%s AND domain=%s", (addr, domain_name))
        log.info(f"Deleted alias: {addr}")

    # Upsert
    for alias in aliases:
        address      = alias["address"]
        goto         = alias["goto"]
        alias_active = int(alias.get("active", True))

        if address in existing_aliases:
            cur.execute(
                "UPDATE alias SET goto=%s, active=%s, modified=%s WHERE address=%s",
                (goto, alias_active, now(), address)
            )
        else:
            cur.execute(
                """INSERT INTO alias
                    (address, goto, domain, active, created, modified)
                   VALUES (%s,%s,%s,%s,%s,%s)""",
                (address, goto, domain_name, alias_active, now(), now())
            )
            log.info(f"Created alias: {address} → {goto}")


def handle_apply_domain(payload: dict, cfg: dict) -> dict:
    """
    Sync domain row + all mailboxes + aliases for one domain.
    Targets the netadmin_mail schema (not PostfixAdmin).
    """
    domain_name = payload.get("domain")
    config      = payload.get("config", {})
    accounts    = payload.get("accounts", [])
    aliases     = payload.get("aliases", [])

    if not domain_name:
        return {"success": False, "error": "Missing 'domain' in payload"}

    try:
        conn = get_db(cfg)
        with conn:
            with conn.cursor() as cur:
                _sync_domain_row(cur, domain_name, config, now)
                incoming_users = _sync_mailboxes(cur, domain_name, accounts, now)
                _sync_aliases(cur, domain_name, aliases, incoming_users, now)
            conn.commit()

        _reload_postfix(cfg)
        log.info(f"apply_domain '{domain_name}' complete — "
                 f"{len(accounts)} mailbox(es), {len(aliases)} alias(es)")
        return {"success": True}

    except Exception as e:
        log.error(f"apply_domain error for '{domain_name}': {e}")
        return {"success": False, "error": str(e)}


def handle_delete_domain(payload: dict, cfg: dict) -> dict:
    """Remove all mailboxes and aliases for a domain, then the domain itself."""
    domain_name = payload["domain"]
    try:
        conn = get_db(cfg)
        with conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM mailbox WHERE domain=%s", (domain_name,))
                cur.execute("DELETE FROM alias   WHERE domain=%s", (domain_name,))
                cur.execute("DELETE FROM domain  WHERE domain=%s", (domain_name,))
            conn.commit()
        log.info(f"Domain '{domain_name}' deleted from MariaDB")
        _reload_postfix(cfg)
        return {"success": True}
    except Exception as e:
        log.error(f"delete_domain error: {e}")
        return {"success": False, "error": str(e)}


def handle_get_quota_usage(payload: dict, cfg: dict) -> dict:
    """
    Return current quota usage per mailbox for a domain.
    Reads from quota2 table (Dovecot quota backend) if available.
    """
    domain_name = payload["domain"]
    try:
        conn = get_db(cfg)
        usage = {}
        with conn:
            with conn.cursor() as cur:
                # quota2 is Dovecot's quota backend table — may not exist in netadmin_mail
                # If absent, quota usage simply won't be reported (non-fatal)
                cur.execute("SHOW TABLES LIKE 'quota2'")
                has_quota2 = cur.fetchone() is not None

                if has_quota2:
                    cur.execute(
                        """SELECT username,
                                  bytes   AS used_bytes,
                                  ROUND(bytes/1048576) AS used_mb
                           FROM quota2
                           WHERE username LIKE %s""",
                        (f"%@{domain_name}",)
                    )
                    for row in cur.fetchall():
                        usage[row["username"]] = {
                            "used_bytes": row["used_bytes"],
                            "used_mb":    int(row["used_mb"] or 0),
                        }
                else:
                    # Fallback: quota table (older Dovecot setup)
                    cur.execute(
                        """SELECT username,
                                  current_bytes   AS used_bytes,
                                  ROUND(current_bytes/1048576) AS used_mb
                           FROM quota
                           WHERE username LIKE %s""",
                        (f"%@{domain_name}",)
                    )
                    for row in cur.fetchall():
                        usage[row["username"]] = {
                            "used_bytes": row["used_bytes"],
                            "used_mb":    int(row["used_mb"] or 0),
                        }

        return {"success": True, "domain": domain_name, "usage": usage}
    except Exception as e:
        log.error(f"get_quota_usage error: {e}")
        return {"success": False, "error": str(e)}


def _reload_postfix(cfg: dict):
    cmd = cfg.get("postfix_reload_cmd", "postfix reload")
    try:
        result = subprocess.run(
            cmd.split(), capture_output=True, text=True, timeout=15
        )
        if result.returncode != 0:
            log.warning(f"postfix reload: {result.stderr.strip()}")
        else:
            log.info("Postfix reloaded")
    except FileNotFoundError:
        log.warning(f"postfix command not found: '{cmd}'")
    except Exception as e:
        log.warning(f"postfix reload error: {e}")



def handle_get_state(payload: dict, cfg: dict) -> dict:
    """Return the full state of MariaDB for all domains (or a specific one).

    This is the canonical source of truth — used by the backend to sync
    its local SQLite cache at startup and on demand.

    Payload:
      {}                      → all domains
      {"domain": "example.com"} → one domain only

    Response:
    {
      "success": true,
      "domains": [
        {
          "name":           "example.com",
          "active":         true,
          "max_accounts":   0,
          "max_quota_mb":   0,
          "accounts": [
            {
              "username":   "user",         # local part
              "email":      "user@example.com",
              "quota_mb":   1024,           # converted from bytes
              "active":     true
              # password_hash intentionally NOT returned for security
            }
          ],
          "aliases": [
            {"address": "alias@example.com", "goto": "dest@example.com", "active": true}
          ]
        }
      ]
    }
    """
    filter_domain = payload.get("domain")
    try:
        conn = get_db(cfg)
        result = []
        with conn:
            with conn.cursor() as cur:
                # Fetch domains
                if filter_domain:
                    cur.execute("SELECT * FROM domain WHERE domain=%s", (filter_domain,))
                else:
                    cur.execute("SELECT * FROM domain WHERE active=1 ORDER BY domain")
                domains = cur.fetchall()

                for dom in domains:
                    domain_name = dom["domain"]

                    # Mailboxes — no local_part/maildir in netadmin_mail
                    cur.execute(
                        """SELECT username, name, quota_mb, active
                           FROM mailbox
                           WHERE domain=%s
                           ORDER BY username""",
                        (domain_name,)
                    )
                    mailboxes = cur.fetchall()

                    # Aliases
                    cur.execute(
                        """SELECT address, goto, active
                           FROM alias
                           WHERE domain=%s
                           ORDER BY address""",
                        (domain_name,)
                    )
                    aliases = cur.fetchall()

                    result.append({
                        "name":         domain_name,
                        "active":       bool(dom["active"]),
                        "max_accounts": int(dom.get("max_accounts", 0)),
                        "max_quota_mb": int(dom.get("max_quota_mb", 0)),
                        "mb_quota_mb":  int(dom.get("mb_quota_mb", 0)),
                        "accounts": [
                            {
                                "username": m["username"],
                                "name":     m["name"] or "",
                                "quota_mb": int(m["quota_mb"] or 0),
                                "active":   bool(m["active"]),
                            }
                            for m in mailboxes
                        ],
                        "aliases": [
                            {
                                "address": a["address"],
                                "goto":    a["goto"],
                                "active":  bool(a["active"]),
                            }
                            for a in aliases
                        ],
                    })

        return {"success": True, "domains": result}
    except Exception as e:
        log.error(f"get_state error: {e}")
        return {"success": False, "error": str(e)}


def handle_generate_dkim(payload: dict, cfg: dict) -> dict:
    """
    Generate a DKIM keypair for a domain using rspamadm dkim_keygen.
    Writes private key and map config to the configured DKIM directory.

    Payload: { "domain": "example.com", "selector": "mail" }
    Returns: { "success": true, "selector": "mail", "txt_record": "v=DKIM1; k=rsa; p=..." }
    """
    domain   = payload["domain"]
    selector = payload.get("selector") or cfg.get("dkim_selector", "mail")
    dkim_dir = cfg.get("dkim_dir", "/var/lib/rspamd/dkim")

    try:
        Path(dkim_dir).mkdir(parents=True, exist_ok=True)

        key_file = Path(dkim_dir) / f"{selector}.{domain}.key"
        txt_file = Path(dkim_dir) / f"{selector}.{domain}.pub.txt"

        # Generate keypair
        result = subprocess.run(
            ["rspamadm", "dkim_keygen",
             "-s", selector,
             "-d", domain,
             "-k", str(key_file)],
            capture_output=True, text=True, timeout=30
        )

        if result.returncode != 0:
            err = result.stderr.strip() or result.stdout.strip()
            log.error(f"dkim_keygen failed for {domain}: {err}")
            return {"success": False, "error": f"rspamadm error: {err}"}

        # rspamadm writes the TXT record content to stdout
        txt_output = result.stdout.strip()
        # Save for reference
        txt_file.write_text(txt_output)

        # Extract just the p= value for the TXT record
        # rspamadm outputs: selector._domainkey.domain. IN TXT "v=DKIM1; k=rsa; p=MIIB..."
        # (possibly split across multiple quoted strings)
        import re
        parts = re.findall(r'"([^"]+)"', txt_output)
        txt_value = "".join(parts)  # join multi-part TXT

        if not txt_value:
            # Fallback: try to extract from the raw output
            m = re.search(r'(v=DKIM1[^"\n]+)', txt_output)
            txt_value = m.group(1).strip() if m else txt_output

        # Write Rspamd signing config entry
        signing_conf = Path(dkim_dir) / "signing.conf"
        entry = f'\n# {domain}\n{domain}\t{{\n\tpath = "{key_file}";\n\tselector = "{selector}";\n}}\n'
        existing = signing_conf.read_text() if signing_conf.exists() else ""
        if domain not in existing:
            with open(signing_conf, "a") as f:
                f.write(entry)
            log.info(f"Added {domain} to {signing_conf}")

        log.info(f"DKIM generated for {domain} selector={selector}")
        return {
            "success":    True,
            "domain":     domain,
            "selector":   selector,
            "txt_record": txt_value,
            "key_file":   str(key_file),
        }

    except FileNotFoundError:
        return {"success": False, "error": "rspamadm not found — is Rspamd installed?"}
    except subprocess.TimeoutExpired:
        return {"success": False, "error": "rspamadm timed out"}
    except Exception as e:
        log.error(f"generate_dkim error: {e}")
        return {"success": False, "error": str(e)}

HANDLERS = {
    "apply_domain":      handle_apply_domain,
    "delete_domain":     handle_delete_domain,
    "get_quota_usage":   handle_get_quota_usage,
    "get_state":         handle_get_state,
    "generate_dkim":     handle_generate_dkim,
}


# ──────────────────────────────────────────────
# Main loop
# ──────────────────────────────────────────────

def run(cfg: dict):
    redis_kwargs = {
        "host":                  cfg["redis_host"],
        "port":                  int(cfg["redis_port"]),
        "db":                    int(cfg["redis_db"]),
        "decode_responses":      True,
        "socket_connect_timeout": 10,   # timeout connexion initiale
        "socket_timeout":         30,   # timeout opérations
        "socket_keepalive":       True, # keepalive TCP
            }
    if cfg.get("redis_password"):
        redis_kwargs["password"] = cfg["redis_password"]
    tls_cfg = cfg.get("tls", {})
    if tls_cfg.get("enabled", False):
        redis_kwargs["ssl"]             = True
        redis_kwargs["ssl_certfile"]    = tls_cfg["client_cert"]
        redis_kwargs["ssl_keyfile"]     = tls_cfg["client_key"]
        redis_kwargs["ssl_ca_certs"]    = tls_cfg["ca_cert"]
        redis_kwargs["ssl_cert_reqs"]   = "required"
        log.info("Redis mTLS enabled")
    else:
        log.warning("Redis TLS disabled — traffic is unencrypted")

    # Test DB connection on startup
    try:
        conn = get_db(cfg)
        with conn:
            with conn.cursor() as cur:
                cur.execute("SELECT COUNT(*) AS n FROM domain")
                n = cur.fetchone()["n"]
        log.info(f"MariaDB connected — {n} domain(s) in database")
    except Exception as e:
        log.error(f"Cannot connect to MariaDB: {e}")
        sys.exit(1)

    r = redis.Redis(**redis_kwargs)
    r.ping()
    log.info(f"Connected to Redis at {cfg['redis_host']}:{cfg['redis_port']}")

    pubsub = r.pubsub()
    pubsub.subscribe("mail.commands")
    log.info("Subscribed to mail.commands — waiting for messages…")
    # Announce reconnection so the backend can flush the pending queue
    import socket as _socket
    r.publish("agent.ready", json.dumps({"service": "mail", "host": _socket.gethostname()}))
    log.info(f"Published agent.ready for service 'mail'")

    for message in pubsub.listen():
        if message["type"] != "message":
            continue
        try:
            msg = json.loads(message["data"])
        except (json.JSONDecodeError, KeyError):
            log.warning("Malformed message received, ignored")
            continue

        # mTLS handles authentication — no token check needed
        # Only clients with a valid CA-signed cert can connect to Redis

        msg_id  = msg.get("id", "unknown")
        action  = msg.get("action", "")
        payload = msg.get("payload", {})

        log.info(f"Received action '{action}' (id={msg_id})")

        handler = HANDLERS.get(action)
        if not handler:
            ack = {"id": msg_id, "success": False, "error": f"Unknown action '{action}'"}
        else:
            ack = handler(payload, cfg)
            ack["id"] = msg_id

        r.publish("mail.acks", json.dumps(ack))
        log.info(f"ACK for '{action}' (id={msg_id}): success={ack.get('success')}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NetAdmin Mail Agent (PostfixAdmin/MySQL)")
    parser.add_argument("--config", default="/etc/netadmin/config.yaml")
    args = parser.parse_args()

    if not Path(args.config).exists():
        log.error(f"Config file not found: {args.config}")
        sys.exit(1)

    agent_cfg = {**DEFAULT_CFG, **load_config(args.config)}
    log.info(
        f"Mail agent starting "
        f"(db={agent_cfg['db_user']}@{agent_cfg['db_host']}/{agent_cfg['db_name']})"
    )

    try:
        run(agent_cfg)
    except KeyboardInterrupt:
        log.info("Mail agent stopped")
    except Exception as e:
        log.error(f"Fatal error: {e}")
        sys.exit(1)