#!/bin/bash
##
## install_quota_clone.sh
##
## Configures Dovecot 2.4 quota_clone plugin and enables PostfixAdmin
## quota display by writing real-time usage to the quota2 PostgreSQL table.
##
## Must be run as root on the target mail server.
##
## What this script does:
##   - Backs up the Dovecot and PostfixAdmin directories
##   - Drops the old ADD-semantics mergequota2 PostgreSQL trigger
##   - Creates a new SET-semantics upsertquota2 trigger (validation only)
##   - Creates /etc/dovecot/conf.d/91-quota-clone.conf (dict_server)
##   - Appends the quota_clone plugin block to 90-quota.conf
##   - Enables $CONF['used_quotas'] and $CONF['new_quota_table'] in PostfixAdmin
##   - Installs /usr/local/bin/update_quota2.sh for initial batch fill
##   - Reloads Dovecot
##
## Safe to run multiple times (idempotent).
##

set -uo pipefail

# ── paths ─────────────────────────────────────────────────────────────────────
DOVECOT_LINK="/usr/local/dovecot"
PFA_LINK="/var/www/adm.oopen.de/postfixadmin"
BACKUP_BASE="/root/backup"

# ── colour helpers ────────────────────────────────────────────────────────────
_bold='\033[1m'; _green='\033[0;32m'; _red='\033[0;31m'; _yellow='\033[1;33m'; _nc='\033[0m'

heading() { echo; echo -e "${_bold}${*}${_nc}"; }
step()    { printf '  %-60s' "$*"; }
ok()      { echo -e " [${_green}OK${_nc}]"; }
skip()    { echo -e " [${_yellow}SKIP${_nc}]"; }
failed()  { echo -e " [${_red}FAILED${_nc}]"; }
info()    { echo -e "  ${*}"; }
die()     { echo -e "\n${_red}FATAL: ${*}${_nc}" >&2; exit 1; }

# ── root check ────────────────────────────────────────────────────────────────
[[ $EUID -eq 0 ]] || die "This script must be run as root."

# ── resolve symlinks ──────────────────────────────────────────────────────────
[[ -L "$DOVECOT_LINK" ]] || die "$DOVECOT_LINK is not a symlink."
[[ -L "$PFA_LINK"     ]] || die "$PFA_LINK is not a symlink."

DOVECOT_REAL=$(realpath "$DOVECOT_LINK")
PFA_REAL=$(realpath "$PFA_LINK")
DOVECOT_CONF_D="${DOVECOT_REAL}/etc/dovecot/conf.d"
SQL_CONNECT="${DOVECOT_REAL}/etc/dovecot/sql-connect.conf.ext"
PFA_CONFIG="${PFA_REAL}/config.local.php"

echo
echo -e "${_bold}Dovecot quota_clone + PostfixAdmin quota display${_nc}"
echo
info "Dovecot : $DOVECOT_REAL"
info "PFA     : $PFA_REAL"

# ── read DB credentials from Dovecot's sql-connect.conf.ext ──────────────────
[[ -f "$SQL_CONNECT" ]] || die "Not found: $SQL_CONNECT"
dbuser=$(grep -E '^\s*user\s*='     "$SQL_CONNECT" | head -1 | sed 's/.*=\s*//' | tr -d '[:space:]')
dbpass=$(grep -E '^\s*password\s*=' "$SQL_CONNECT" | head -1 | sed 's/.*=\s*//' | tr -d '[:space:]')
dbname=$(grep -E '^\s*dbname\s*='   "$SQL_CONNECT" | head -1 | sed 's/.*=\s*//' | tr -d '[:space:]')
[[ -n "$dbuser" ]] || die "Could not parse 'user' from $SQL_CONNECT"
[[ -n "$dbpass" ]] || die "Could not parse 'password' from $SQL_CONNECT"
[[ -n "$dbname" ]] || die "Could not parse 'dbname' from $SQL_CONNECT"
info "DB      : $dbuser@$dbname"

# ── backup ────────────────────────────────────────────────────────────────────
heading "Backup"
mkdir -p "$BACKUP_BASE"
BACKUP_DIR="${BACKUP_BASE}/quota_clone_$(date +%Y%m%d_%H%M%S)"
mkdir -p "$BACKUP_DIR"

step "$(basename "$DOVECOT_REAL") ..."
tar -czf "${BACKUP_DIR}/dovecot.tar.gz" \
    -C "$(dirname "$DOVECOT_REAL")" "$(basename "$DOVECOT_REAL")" \
    && ok || { failed; die "Backup of $DOVECOT_REAL failed."; }

step "$(basename "$PFA_REAL") ..."
tar -czf "${BACKUP_DIR}/postfixadmin.tar.gz" \
    -C "$(dirname "$PFA_REAL")" "$(basename "$PFA_REAL")" \
    && ok || { failed; die "Backup of $PFA_REAL failed."; }

info "Saved to: $BACKUP_DIR"

# ── PostgreSQL ────────────────────────────────────────────────────────────────
heading "PostgreSQL"

step "Drop old mergequota2 trigger + function ..."
psql -U"$dbuser" "$dbname" -c "
    DROP TRIGGER IF EXISTS mergequota2 ON quota2;
    DROP FUNCTION IF EXISTS merge_quota2();
" > /dev/null 2>&1 \
    && ok || { failed; die "Could not drop old mergequota2 trigger."; }

# Single-quoted heredoc: $$ is passed literally to psql (no shell expansion).
step "Create upsertquota2 trigger + function ..."
psql -U"$dbuser" "$dbname" <<'SQL' > /dev/null 2>&1
CREATE OR REPLACE FUNCTION upsert_quota2() RETURNS trigger
    LANGUAGE plpgsql AS $$
    BEGIN
        IF NEW.bytes    IS NULL OR NEW.bytes    < 0 THEN NEW.bytes    := 0; END IF;
        IF NEW.messages IS NULL OR NEW.messages < 0 THEN NEW.messages := 0; END IF;
        RETURN NEW;
    END;
    $$;

ALTER FUNCTION public.upsert_quota2() OWNER TO postfix;

DO $$
BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'upsertquota2') THEN
        CREATE TRIGGER upsertquota2
            BEFORE INSERT ON quota2
            FOR EACH ROW EXECUTE PROCEDURE upsert_quota2();
    END IF;
END;
$$;
SQL
[[ $? -eq 0 ]] && ok || { failed; die "Could not create upsertquota2 trigger."; }

# ── Dovecot configuration ─────────────────────────────────────────────────────
heading "Dovecot configuration"

# 91-quota-clone.conf — always (re)written, content is idempotent
QUOTA_CLONE_CONF="${DOVECOT_CONF_D}/91-quota-clone.conf"
step "Write 91-quota-clone.conf ..."
cat > "$QUOTA_CLONE_CONF" <<EOF
##
## Quota Clone Configuration (Dovecot 2.4+)
##
## Writes quota usage to the quota2 table in real time so PostfixAdmin
## can display mailbox fill levels via \$CONF['used_quotas'] = 'YES'.
##
dict_server {
  dict quota_clone_pgsql {
    driver = sql
    sql_driver = pgsql
    pgsql localhost {
      parameters {
        user = $dbuser
        password = $dbpass
        dbname = $dbname
      }
    }
    dict_map priv/quota/storage {
      sql_table = quota2
      username_field = username
      value_field bytes {
      }
    }
    dict_map priv/quota/messages {
      sql_table = quota2
      username_field = username
      value_field messages {
      }
    }
  }
}
EOF
[[ $? -eq 0 ]] && ok || { failed; die "Could not write $QUOTA_CLONE_CONF"; }

# 90-quota.conf — append only if not already present
QUOTA_CONF="${DOVECOT_CONF_D}/90-quota.conf"
step "Append quota_clone block to 90-quota.conf ..."
if grep -q "quota_clone_pgsql" "$QUOTA_CONF" 2>/dev/null; then
    skip
else
    cat >> "$QUOTA_CONF" <<'EOF'

##
## Quota Clone - write quota usage to quota2 table for PostfixAdmin
##
mail_plugins {
  quota_clone = yes
}

quota_clone {
  dict proxy {
    name = quota_clone_pgsql
  }
}
EOF
    [[ $? -eq 0 ]] && ok || { failed; die "Could not append to $QUOTA_CONF"; }
fi

# ── PostfixAdmin ──────────────────────────────────────────────────────────────
heading "PostfixAdmin"

[[ -f "$PFA_CONFIG" ]] || die "Not found: $PFA_CONFIG"

# Sets $CONF['key'] = 'val' in config.local.php.
# - Already correct value → SKIP
# - Present with different value → update in place (perl, only non-commented lines)
# - Not present → append
pfa_set() {
    local key="$1" val="$2"
    step "\$CONF['${key}'] = '${val}' ..."
    if grep -qF "\$CONF['${key}'] = '${val}';" "$PFA_CONFIG"; then
        skip
        return
    fi
    if grep -qF "\$CONF['${key}']" "$PFA_CONFIG"; then
        # present with wrong value — update active (non-commented) line
        perl -i -p -e "s|^(\s*\\\$CONF\['${key}'\]\s*=\s*)'[^']*'|\${1}'${val}'|" "$PFA_CONFIG" \
            && ok || { failed; die "Could not update \$CONF['${key}'] in $PFA_CONFIG"; }
    else
        # not present — append
        { echo ""; echo "\$CONF['${key}'] = '${val}';"; } >> "$PFA_CONFIG" \
            && ok || { failed; die "Could not append \$CONF['${key}'] to $PFA_CONFIG"; }
    fi
}

pfa_set used_quotas    YES
pfa_set new_quota_table YES

# ── /usr/local/bin/update_quota2.sh ──────────────────────────────────────────
heading "Batch update script"

# Single-quoted PYEOF: no shell expansion inside the Python source.
step "Install /usr/local/bin/update_quota2.sh ..."
cat > /usr/local/bin/update_quota2.sh <<'PYEOF'
#!/usr/bin/env python3
"""
Batch-update the quota2 table for all active mailboxes.

Uses 'doveadm quota get' (must run as root) to read current usage and
writes it to quota2 via INSERT ... ON CONFLICT DO UPDATE.

Usage:
    update_quota2.sh [--dry-run] [--limit N]

Options:
    --dry-run   Print what would be done without touching the database.
    --limit N   Process only the first N mailboxes (useful for testing).
"""

import subprocess
import sys

DOVEADM = '/usr/local/dovecot/bin/doveadm'
PSQL    = 'psql'
DB_USER = 'postfix'
DB_NAME = 'postfix'


def get_users(limit=None):
    query = 'SELECT username FROM mailbox WHERE active = true ORDER BY username;'
    if limit:
        query = ('SELECT username FROM mailbox WHERE active = true '
                 'ORDER BY username LIMIT %d;' % limit)
    result = subprocess.run(
        [PSQL, '-U', DB_USER, DB_NAME, '-tAF', '|', '-c', query],
        capture_output=True, text=True
    )
    return [u.strip() for u in result.stdout.strip().split('\n') if u.strip()]


def get_quota(user):
    """Returns (bytes, messages) or (None, None) if not available."""
    result = subprocess.run(
        [DOVEADM, 'quota', 'get', '-u', user],
        capture_output=True, text=True
    )
    bytes_val = messages_val = None
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) < 4:
            continue
        if parts[2] == 'STORAGE' and parts[1] != 'Type':
            try:
                bytes_val = int(parts[3]) * 1024   # KiB → bytes
            except ValueError:
                bytes_val = 0
        elif parts[2] == 'MESSAGE' and parts[1] != 'Type':
            try:
                messages_val = int(parts[3])
            except ValueError:
                messages_val = 0
    return bytes_val, messages_val


def upsert_quota2(user, bytes_val, messages_val):
    sql = (
        "INSERT INTO quota2 (username, bytes, messages) VALUES ('"
        + user.replace("'", "''")
        + "', " + str(bytes_val) + ", " + str(messages_val) + ") "
        "ON CONFLICT (username) DO UPDATE "
        "SET bytes = EXCLUDED.bytes, messages = EXCLUDED.messages;"
    )
    subprocess.run([PSQL, '-U', DB_USER, DB_NAME, '-c', sql],
                   capture_output=True)


def main():
    dry_run = '--dry-run' in sys.argv
    limit = None
    if '--limit' in sys.argv:
        idx = sys.argv.index('--limit')
        try:
            limit = int(sys.argv[idx + 1])
        except (IndexError, ValueError):
            print('Usage: update_quota2.sh [--dry-run] [--limit N]')
            sys.exit(1)

    users = get_users(limit)
    total = len(users)
    print('Processing %d mailboxes...' % total)

    updated = skipped = 0
    for i, user in enumerate(users, 1):
        bytes_val, messages_val = get_quota(user)
        if bytes_val is None or messages_val is None:
            print('  [%d/%d] SKIP %s (no quota data)' % (i, total, user))
            skipped += 1
            continue
        if dry_run:
            print('  [%d/%d] DRY  %s: %d bytes, %d messages'
                  % (i, total, user, bytes_val, messages_val))
        else:
            upsert_quota2(user, bytes_val, messages_val)
            if i % 50 == 0 or i == total:
                print('  [%d/%d] done' % (i, total))
        updated += 1

    print('Done: %d updated, %d skipped.' % (updated, skipped))


if __name__ == '__main__':
    main()
PYEOF
chmod +x /usr/local/bin/update_quota2.sh
[[ $? -eq 0 ]] && ok || { failed; die "Could not install update_quota2.sh"; }

# ── Dovecot reload ────────────────────────────────────────────────────────────
heading "Dovecot reload"
step "Reloading Dovecot ..."
"${DOVECOT_REAL}/sbin/dovecot" reload \
    && ok || { failed; die "Dovecot reload failed."; }

# ── summary ───────────────────────────────────────────────────────────────────
echo
echo -e "${_bold}Done.${_nc}"
echo
info "Backup : $BACKUP_DIR"
info "Next   : run the initial batch fill of quota2 for all mailboxes:"
info "         /usr/local/bin/update_quota2.sh [--dry-run] [--limit 5]"
echo