#!/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 table.
##
## Supports PostgreSQL and MySQL/MariaDB backends.
## Must be run as root on the target mail server.
##
## What this script does:
##   - Backs up the Dovecot and PostfixAdmin directories
##   - Creates the quota2 table if it does not exist (MySQL only)
##   - Drops the old ADD-semantics mergequota2 trigger (if present)
##   - 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"
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."

echo
echo -e "${_bold}Dovecot quota_clone + PostfixAdmin quota display${_nc}"
echo

# ── Dovecot path ──────────────────────────────────────────────────────────────
[[ -L "$DOVECOT_LINK" ]] || die "$DOVECOT_LINK is not a symlink."
DOVECOT_REAL=$(realpath "$DOVECOT_LINK")
DOVECOT_CONF_D="${DOVECOT_REAL}/etc/dovecot/conf.d"
SQL_CONNECT="${DOVECOT_REAL}/etc/dovecot/sql-connect.conf.ext"
info "Dovecot : $DOVECOT_REAL"

# ── PostfixAdmin path (interactive) ──────────────────────────────────────────
_pfa_default=""
mapfile -t _pfa_candidates < <(find /var/www -maxdepth 3 -name "postfixadmin" -type l 2>/dev/null | sort)
if [[ ${#_pfa_candidates[@]} -eq 1 ]]; then
    _pfa_default="${_pfa_candidates[0]}"
fi

while true; do
    if [[ -n "$_pfa_default" ]]; then
        read -rp "  PostfixAdmin symlink [${_pfa_default}]: " PFA_LINK
        PFA_LINK="${PFA_LINK:-$_pfa_default}"
    else
        read -rp "  PostfixAdmin symlink (e.g. /var/www/hostname/htdocs/postfixadmin): " PFA_LINK
    fi
    if [[ -L "$PFA_LINK" ]]; then
        break
    elif [[ -d "$PFA_LINK" ]]; then
        echo -e "  ${_yellow}Warning: $PFA_LINK is a directory, not a symlink — continuing anyway.${_nc}"
        break
    else
        echo -e "  ${_red}Not found or not a symlink: $PFA_LINK — please try again.${_nc}"
    fi
done

PFA_REAL=$(realpath "$PFA_LINK")
info "PFA     : $PFA_REAL"

# config.local.php is either in PFA_REAL itself (symlink → PFA root)
# or one level up (symlink → public/ document root)
if [[ -f "${PFA_REAL}/config.local.php" ]]; then
    PFA_CONFIG="${PFA_REAL}/config.local.php"
elif [[ -f "$(dirname "${PFA_REAL}")/config.local.php" ]]; then
    PFA_CONFIG="$(dirname "${PFA_REAL}")/config.local.php"
else
    die "config.local.php not found in $PFA_REAL nor in $(dirname "$PFA_REAL")"
fi
info "Config  : $PFA_CONFIG"

# ── DB credentials from sql-connect.conf.ext ─────────────────────────────────
[[ -f "$SQL_CONNECT" ]] || die "Not found: $SQL_CONNECT"

sql_driver=$(grep -E '^\s*sql_driver\s*=' "$SQL_CONNECT" | head -1 | sed 's/.*=\s*//' | tr -d '[:space:]')
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 "$sql_driver" ]] || die "Could not parse 'sql_driver' from $SQL_CONNECT"
[[ -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"

case "$sql_driver" in
    pgsql)
        info "DB      : PostgreSQL  $dbuser@$dbname"
        dbsocket=""
        ;;
    mysql)
        # socket path: line looks like  mysql /run/mysqld/mysqld.sock {
        dbsocket=$(grep -E '^\s*mysql\s+/' "$SQL_CONNECT" | head -1 | awk '{print $2}')
        [[ -n "$dbsocket" ]] || die "Could not parse MySQL socket path from $SQL_CONNECT"
        info "DB      : MariaDB/MySQL  $dbuser@$dbname  (socket: $dbsocket)"
        ;;
    *)
        die "Unsupported sql_driver '$sql_driver' in $SQL_CONNECT (expected: pgsql or mysql)"
        ;;
esac

# ── DB helper: run SQL (DDL: CREATE TABLE, triggers) ─────────────────────────
# pgsql: postfix user via psql.
# mysql: OS root → MariaDB root via unix_socket auth (Debian default);
#        falls back to postfix user credentials if that fails.
db_exec() {
    if [[ "$sql_driver" == "pgsql" ]]; then
        psql -U"$dbuser" "$dbname" "$@"
    else
        if mysql --socket="$dbsocket" "$dbname" "$@" 2>/dev/null; then
            return 0
        else
            mysql -u"$dbuser" -p"$dbpass" --socket="$dbsocket" "$dbname" "$@"
        fi
    fi
}

# ── 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") ..."
if tar -czf "${BACKUP_DIR}/dovecot.tar.gz" \
        -C "$(dirname "$DOVECOT_REAL")" "$(basename "$DOVECOT_REAL")"; then
    ok
else
    failed; die "Backup of $DOVECOT_REAL failed."
fi

_pfa_backup_src=$(dirname "$PFA_REAL")
_pfa_backup_name=$(basename "$PFA_REAL")
step "${_pfa_backup_name} ..."
if tar -czf "${BACKUP_DIR}/postfixadmin.tar.gz" \
        -C "$_pfa_backup_src" "$_pfa_backup_name"; then
    ok
else
    failed; die "Backup of $PFA_REAL failed."
fi

info "Saved to: $BACKUP_DIR"

# ── Database ──────────────────────────────────────────────────────────────────
if [[ "$sql_driver" == "pgsql" ]]; then
    heading "PostgreSQL"
else
    heading "MariaDB / MySQL"
fi

if [[ "$sql_driver" == "mysql" ]]; then
    step "Create quota2 table if not exists ..."
    if db_exec -e "
        CREATE TABLE IF NOT EXISTS quota2 (
            username varchar(100) NOT NULL,
            bytes bigint NOT NULL DEFAULT 0,
            messages int NOT NULL DEFAULT 0,
            PRIMARY KEY (username)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
    " > /dev/null 2>&1; then
        ok
    else
        failed; die "Could not create quota2 table."
    fi
fi

step "Drop old mergequota2 trigger ..."
if [[ "$sql_driver" == "pgsql" ]]; then
    if db_exec -c "
        DROP TRIGGER IF EXISTS mergequota2 ON quota2;
        DROP FUNCTION IF EXISTS merge_quota2();
    " > /dev/null 2>&1; then
        ok
    else
        failed; die "Could not drop old mergequota2 trigger."
    fi
else
    if db_exec -e "DROP TRIGGER IF EXISTS mergequota2;" > /dev/null 2>&1; then
        ok
    else
        failed; die "Could not drop old mergequota2 trigger."
    fi
fi

step "Create upsertquota2 trigger ..."
if [[ "$sql_driver" == "pgsql" ]]; then
    # Single-quoted heredoc: $$ reaches psql literally (no shell expansion).
    if db_exec <<'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
    then
        ok
    else
        failed; die "Could not create upsertquota2 trigger."
    fi
else
    # DELIMITER lets the MySQL CLI handle ; inside BEGIN...END in batch mode.
    if db_exec <<'SQL' > /dev/null 2>&1
DROP TRIGGER IF EXISTS upsertquota2;
DELIMITER //
CREATE TRIGGER upsertquota2
    BEFORE INSERT ON quota2
    FOR EACH ROW
    BEGIN
        IF NEW.bytes    IS NULL OR NEW.bytes    < 0 THEN SET NEW.bytes    = 0; END IF;
        IF NEW.messages IS NULL OR NEW.messages < 0 THEN SET NEW.messages = 0; END IF;
    END//
DELIMITER ;
SQL
    then
        ok
    else
        failed; die "Could not create upsertquota2 trigger."
    fi
fi

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

QUOTA_CLONE_CONF="${DOVECOT_CONF_D}/91-quota-clone.conf"

_write_quota_clone_conf() {
    if [[ "$sql_driver" == "pgsql" ]]; then
        cat > "$QUOTA_CLONE_CONF" <<EOF
##
## Quota Clone Configuration (Dovecot 2.4+ / PostgreSQL)
##
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
    else
        cat > "$QUOTA_CLONE_CONF" <<EOF
##
## Quota Clone Configuration (Dovecot 2.4+ / MariaDB / MySQL)
##
dict_server {
  dict quota_clone_mysql {
    driver = sql
    sql_driver = mysql
    mysql $dbsocket {
      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
    fi
}

step "Write 91-quota-clone.conf ..."
if _write_quota_clone_conf; then ok; else failed; die "Could not write $QUOTA_CLONE_CONF"; fi

# dict name differs by driver — used both as grep marker and in quota_clone block
if [[ "$sql_driver" == "pgsql" ]]; then
    _dict_name="quota_clone_pgsql"
else
    _dict_name="quota_clone_mysql"
fi

QUOTA_CONF="${DOVECOT_CONF_D}/90-quota.conf"
step "Append quota_clone block to 90-quota.conf ..."
if grep -q "$_dict_name" "$QUOTA_CONF" 2>/dev/null; then
    skip
else
    if cat >> "$QUOTA_CONF" <<EOF

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

quota_clone {
  dict proxy {
    name = $_dict_name
  }
}
EOF
    then
        ok
    else
        failed; die "Could not append to $QUOTA_CONF"
    fi
fi

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

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
        if perl -i -p -e "s|^(\s*\\\$CONF\['${key}'\]\s*=\s*)'[^']*'|\${1}'${val}'|" "$PFA_CONFIG"; then
            ok
        else
            failed; die "Could not update \$CONF['${key}'] in $PFA_CONFIG"
        fi
    else
        if { echo ""; echo "\$CONF['${key}'] = '${val}';"; } >> "$PFA_CONFIG"; then
            ok
        else
            failed; die "Could not append \$CONF['${key}'] to $PFA_CONFIG"
        fi
    fi
}

pfa_set used_quotas     YES
pfa_set new_quota_table YES

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

step "Install /usr/local/bin/update_quota2.sh ..."

if [[ "$sql_driver" == "pgsql" ]]; then

cat > /usr/local/bin/update_quota2.sh <<'PYEOF'
#!/usr/bin/env python3
"""
Batch-update the quota2 table for all active mailboxes (PostgreSQL).

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

Usage: update_quota2.sh [--dry-run] [--limit N]
"""
import subprocess, sys

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

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

def get_quota(user):
    r = subprocess.run([DOVEADM, 'quota', 'get', '-u', user],
                       capture_output=True, text=True)
    bv = mv = None
    for line in r.stdout.splitlines():
        p = line.split()
        if len(p) < 4: continue
        if p[2] == 'STORAGE' and p[1] != 'Type':
            try: bv = int(p[3]) * 1024
            except ValueError: bv = 0
        elif p[2] == 'MESSAGE' and p[1] != 'Type':
            try: mv = int(p[3])
            except ValueError: mv = 0
    return bv, mv

def upsert(user, bv, mv):
    sql = ("INSERT INTO quota2 (username,bytes,messages) VALUES ('"
           + user.replace("'","''") + "'," + str(bv) + "," + str(mv) + ") "
           "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 = '--dry-run' in sys.argv
    limit = None
    if '--limit' in sys.argv:
        try: limit = int(sys.argv[sys.argv.index('--limit') + 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)
    ok = skip = 0
    for i, user in enumerate(users, 1):
        bv, mv = get_quota(user)
        if bv is None or mv is None:
            print('  [%d/%d] SKIP %s' % (i, total, user)); skip += 1; continue
        if dry:
            print('  [%d/%d] DRY  %s: %d bytes, %d msgs' % (i, total, user, bv, mv))
        else:
            upsert(user, bv, mv)
            if i % 50 == 0 or i == total: print('  [%d/%d] done' % (i, total))
        ok += 1
    print('Done: %d updated, %d skipped.' % (ok, skip))

if __name__ == '__main__': main()
PYEOF

else  # mysql — shell expands DB_USER, DB_PASS, DB_SOCKET, DB_NAME into the script

cat > /usr/local/bin/update_quota2.sh <<PYEOF
#!/usr/bin/env python3
"""
Batch-update the quota2 table for all active mailboxes (MariaDB/MySQL).

Uses 'doveadm quota get' (must run as root) and writes to quota2 via
REPLACE INTO.

Usage: update_quota2.sh [--dry-run] [--limit N]
"""
import subprocess, sys

DOVEADM   = '/usr/local/dovecot/bin/doveadm'
MYSQL     = 'mysql'
DB_USER   = '$dbuser'
DB_PASS   = '$dbpass'
DB_SOCKET = '$dbsocket'
DB_NAME   = '$dbname'

def mysql_query(sql):
    return subprocess.run(
        [MYSQL, '-u', DB_USER, '-p' + DB_PASS, '--socket', DB_SOCKET,
         DB_NAME, '-sNe', sql],
        capture_output=True, text=True)

def get_users(limit=None):
    q = 'SELECT username FROM mailbox WHERE active = 1 ORDER BY username;'
    if limit:
        q = 'SELECT username FROM mailbox WHERE active = 1 ORDER BY username LIMIT %d;' % limit
    r = mysql_query(q)
    return [u.strip() for u in r.stdout.strip().split('\n') if u.strip()]

def get_quota(user):
    r = subprocess.run([DOVEADM, 'quota', 'get', '-u', user],
                       capture_output=True, text=True)
    bv = mv = None
    for line in r.stdout.splitlines():
        p = line.split()
        if len(p) < 4: continue
        if p[2] == 'STORAGE' and p[1] != 'Type':
            try: bv = int(p[3]) * 1024
            except ValueError: bv = 0
        elif p[2] == 'MESSAGE' and p[1] != 'Type':
            try: mv = int(p[3])
            except ValueError: mv = 0
    return bv, mv

def upsert(user, bv, mv):
    sql = ("REPLACE INTO quota2 (username, bytes, messages) VALUES ('"
           + user.replace("'", "\\'") + "', " + str(bv) + ", " + str(mv) + ");")
    mysql_query(sql)

def main():
    dry = '--dry-run' in sys.argv
    limit = None
    if '--limit' in sys.argv:
        try: limit = int(sys.argv[sys.argv.index('--limit') + 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)
    ok = skip = 0
    for i, user in enumerate(users, 1):
        bv, mv = get_quota(user)
        if bv is None or mv is None:
            print('  [%d/%d] SKIP %s' % (i, total, user)); skip += 1; continue
        if dry:
            print('  [%d/%d] DRY  %s: %d bytes, %d msgs' % (i, total, user, bv, mv))
        else:
            upsert(user, bv, mv)
            if i % 50 == 0 or i == total: print('  [%d/%d] done' % (i, total))
        ok += 1
    print('Done: %d updated, %d skipped.' % (ok, skip))

if __name__ == '__main__': main()
PYEOF

fi  # sql_driver

if chmod +x /usr/local/bin/update_quota2.sh; then ok; else failed; die "Could not install update_quota2.sh"; fi

# ── Dovecot reload ────────────────────────────────────────────────────────────
heading "Dovecot reload"
step "Reloading Dovecot ..."
if "${DOVECOT_REAL}/sbin/dovecot" reload; then
    ok
else
    failed; die "Dovecot reload failed."
fi

# ── 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