#!/usr/bin/env bash
# =============================================================================
# Script Name : ip_manager.sh
# Description : Manage customer IPs in IPSET sets (allow / restrict)
#               Called by BOSS4 billing system on customer state changes.
#
# Commands (also PARAMETERS):
#   NEW        IP             — new customer: add to ALLOW set
#   DELETE     IP             — delete/prototype: remove from BOTH sets
#   RESTRICT   IP             — no payment: add to BOTH sets (allow + restrict)
#   SUSPEND    IP             — manual suspend: remove from BOTH sets
#   RESUME     IP             — resume: remove from restrict, add to allow (if not there)
#   UPDATE     IP_OLD IP_NEW  — IP/tariff change: replace OLD with NEW in each set
#
# Usage:
#   ./ip_manager.sh COMMAND IP              # for NEW/DELETE/RESTRICT/SUSPEND/RESUME
#   ./ip_manager.sh UPDATE  IP_OLD IP_NEW   # for UPDATE
#
# Author  : syr4ok (Andrii Syrovatko)
# Version : 2.0.1r
# =============================================================================

# --- STRICT MODE ---
set -uo pipefail
IFS=$'\n\t'

# --- Configuration Loader ---
CONFIG_FILE="$(dirname "$0")/ip_manager.conf"

if [[ -f "$CONFIG_FILE" ]]; then
    # shellcheck source=/dev/null
    source "$CONFIG_FILE"
else
    echo "Error: Configuration file not found. Create ip_manager.conf from example."
    exit 1
fi

# --- Environment & Tools ---
IPSET_BIN=$(which ipset 2>/dev/null || true)
MAIL_BIN=$(which mail 2>/dev/null || true)

if [[ -z "$IPSET_BIN" ]]; then
    echo "Error: 'ipset' utility not found. Please install it."
    exit 1
fi

if [[ -z "$MAIL_BIN" ]]; then
    echo "Warning: 'mail' utility not found. Notifications will be disabled."
fi

# --- LOG COLORS ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

# INIT — ensure log directory and file exist
mkdir -p "${LOG_DIR}" 2>/dev/null || true
CAN_LOG_TO_FILE=0
if touch "${LOG_FILE}" 2>/dev/null; then
    CAN_LOG_TO_FILE=1
    echo -e "${BLUE}Info: Log file (${LOG_FILE}) created or exists. Logging to file enabled.${NC}"
else
    echo -e "${YELLOW}Warning: Cannot write to log file (${LOG_FILE}). Logging to file disabled.${NC}"
fi
# --- LOGGING ---
log() {
    local level="$1"; shift
    local msg="$*"
    local ts
    ts=$(date '+%Y-%m-%d %H:%M:%S')

    # Clear colors for file logging (keep colors in console, but remove them from file logs)
    local plain_msg
    plain_msg=$(echo -e "${msg}" | sed 's/\x1B\[[0-9;]*[mK]//g')

    # Write logs to file if possible
    if [[ "${CAN_LOG_TO_FILE}" -eq 1 ]]; then
        echo "${ts} [${level}] ${plain_msg}" >> "${LOG_FILE}"
    fi

    # Type logs tto console forever
    echo -e "${ts} [${level}] ${msg}"
}

log_info()  { log "INFO " "${GREEN}${*}${NC}"; }
log_warn()  { log "WARN " "${YELLOW}${*}${NC}"; }
log_error() { log "ERROR" "${RED}${*}${NC}"; }
log_debug() { log "DEBUG" "${CYAN}${*}${NC}"; }
log_sep()   {
    log "INFO " "${BLUE}------------------------------------------------------------------------------${NC}"
}

# LOCK — single instance with wait queue (flock-based)
acquire_lock() {
    local call_args
    call_args="$(IFS=" "; echo "$*")"
    exec 9>"${LOCK_FILE}"
    log_debug "Waiting for lock (timeout: ${LOCK_TIMEOUT}s) PID=$$"
    if ! flock -w "${LOCK_TIMEOUT}" 9; then
        log_error "Could not acquire lock after ${LOCK_TIMEOUT}s — another instance hung? Aborting!"

        if [[ -n "${MAIL_BIN}" ]]; then
            echo -e "WARN: $0\nProcess for [${call_args}] timed out by lockfile (${LOCK_FILE})\nafter allowed ${LOCK_TIMEOUT}s waiting.\nexit 9" \
                | "${MAIL_BIN}" -s "WARN: Process timed out by lockfile after ${LOCK_TIMEOUT}s" "$MAIL_RECEIVER"
        fi
        exit 9
    fi
    echo $$ >&9
    log_debug "Lock acquired by PID=$$"
}

release_lock() {
    flock -u 9
    exec 9>&-
    log_debug "Lock released by PID=$$"
}

# --- HELPERS ---
usage() {
    echo -e "${BOLD}Usage:${NC}"
    echo -e "  $0 ${CYAN}NEW${NC}       IP"
    echo -e "  $0 ${CYAN}DELETE${NC}    IP"
    echo -e "  $0 ${CYAN}RESTRICT${NC}  IP"
    echo -e "  $0 ${CYAN}SUSPEND${NC}   IP"
    echo -e "  $0 ${CYAN}RESUME${NC}    IP"
    echo -e "  $0 ${CYAN}UPDATE${NC}    IP_OLD  IP_NEW"
    echo
}

strip_mask() {
    local val="${1}"
    val="${val// /}"   # trim all spaces
    echo "${val%%/32}"
}

# Validate IPv4 address (without mask)
is_valid_ip() {
    local ip="$1"
    local octet='(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)'
    [[ "${ip}" =~ ^${octet}\.${octet}\.${octet}\.${octet}$ ]]
}

ipset_exists() {
    local set="$1"
    local ip="$2"
    # grep -q closes pipe early after first match → ipset gets SIGPIPE → pipefail returns 141
    # fix: read full ipset output first, then grep from variable
    local dump
    dump=$("${IPSET_BIN}" -L "${set}" 2>/dev/null)
    grep -qw "^${ip}$" <<< "${dump}"
}

ipset_add() {
    local set="$1"
    local ip="$2"

    if ipset_exists "${set}" "${ip}"; then
        log_warn "IP ${ip} already in set [${set}] — skipping add"
        return 0
    fi

    log_debug "ipset add ${set} ${ip}"
    if [[ "${DRY_RUN}" -eq 1 ]]; then
        log_warn "[DRY-RUN] Would ADD ${ip} to [${set}]"
        return 0
    fi
    if ! "${IPSET_BIN}" add "${set}" "${ip}" 2>&1 | tee -a "${LOG_FILE}"; then
        log_error "Failed to ADD ${ip} to [${set}]"
        return 1
    fi
    log_info "Added ${ip} → [${set}]"
}

ipset_del() {
    local set="$1"
    local ip="$2"

    if ! ipset_exists "${set}" "${ip}"; then
        log_warn "IP ${ip} not found in set [${set}] — skipping del"
        return 0
    fi

    log_debug "ipset del ${set} ${ip}"
    if [[ "${DRY_RUN}" -eq 1 ]]; then
        log_warn "[DRY-RUN] Would DEL ${ip} from [${set}]"
        return 0
    fi
    if ! "${IPSET_BIN}" del "${set}" "${ip}" 2>&1 | tee -a "${LOG_FILE}"; then
        log_error "Failed to DEL ${ip} from [${set}]"
        return 1
    fi
    log_info "Removed ${ip} ← [${set}]"
}

ipset_save() {
    log_debug "Saving ipset state to ${IPSET_SAVE_FILE}"
    if [[ "${DRY_RUN}" -eq 1 ]]; then
        log_warn "[DRY-RUN] Would save ipset state to ${IPSET_SAVE_FILE}"
        return 0
    fi
    if ! "${IPSET_BIN}" save > "${IPSET_SAVE_FILE}" 2>&1; then
        log_error "Failed to save ipset state to ${IPSET_SAVE_FILE}"
        return 1
    fi
    log_info "ipset state saved → ${IPSET_SAVE_FILE}"
}

# --- COMMAND HANDLERS ---
cmd_new() {
    local ip="$1"
    log_info "[NEW] New customer IP: ${ip} — adding to allow set"
    ipset_add "${ALLOW_SET}" "${ip}"
    ipset_save
}

cmd_delete() {
    local ip="$1"
    log_info "[DELETE] Removing ${ip} from BOTH sets"
    ipset_del "${ALLOW_SET}"    "${ip}"
    ipset_del "${RESTRICT_SET}" "${ip}"
    ipset_save
}

cmd_restrict() {
    local ip="$1"
    log_info "[RESTRICT] Adding ${ip} to ALLOW + RESTRICT sets (no payment)"
    ipset_add "${ALLOW_SET}"    "${ip}"
    ipset_add "${RESTRICT_SET}" "${ip}"
    ipset_save
}

cmd_suspend() {
    local ip="$1"
    log_info "[SUSPEND] Removing ${ip} from BOTH sets (manual suspend)"
    ipset_del "${ALLOW_SET}"    "${ip}"
    ipset_del "${RESTRICT_SET}" "${ip}"
    ipset_save
}

cmd_resume() {
    local ip="$1"
    log_info "[RESUME] Resuming ${ip}: remove from restrict, ensure in allow"
    ipset_del "${RESTRICT_SET}" "${ip}"
    ipset_add "${ALLOW_SET}"    "${ip}"
    ipset_save
}

cmd_update() {
    local ip_old="$1"
    local ip_new="$2"
    log_info "[UPDATE] Replacing ${ip_old} → ${ip_new}"

    local changed=0

    if ipset_exists "${ALLOW_SET}" "${ip_old}"; then
        log_debug "${ip_old} found in [${ALLOW_SET}] — swapping"
        ipset_del "${ALLOW_SET}" "${ip_old}"
        ipset_add "${ALLOW_SET}" "${ip_new}"
        changed=1
    else
        log_warn "${ip_old} not found in [${ALLOW_SET}] — no change in this set"
    fi

    if ipset_exists "${RESTRICT_SET}" "${ip_old}"; then
        log_debug "${ip_old} found in [${RESTRICT_SET}] — swapping"
        ipset_del "${RESTRICT_SET}" "${ip_old}"
        ipset_add "${RESTRICT_SET}" "${ip_new}"
        changed=1
    else
        log_warn "${ip_old} not found in [${RESTRICT_SET}] — no change in this set"
    fi

    if [[ "${changed}" -eq 0 ]]; then
        log_warn "[UPDATE] ${ip_old} was not found in ANY set — nothing was changed"
    fi

    ipset_save
}

# MAIN EXIT HELPER — always logs final status + releases lock before return
main_exit() {
    local code="$1"
    if [[ "${code}" -eq 0 ]]; then
        log_info "Done. Exit 0 (OK)"
    else
        log_error "Finished with errors (code: ${code})"
    fi
    # log_sep
    release_lock
    # sleep 45    # DEBUG: keep lock released for a while to allow testing concurrent runs and lock timeout
    return "${code}"
}

# MAIN
main() {
    local error_code=0

    log_sep
    log_info "Invoked: $0"
    log_info "Arguments ($#): $(IFS=" "; echo "$*")"
    acquire_lock "$@"
    [[ "${DRY_RUN}" -eq 1 ]] && log_warn "=== DRY-RUN MODE — no real ipset actions will be performed ==="

    if [[ $# -lt 2 ]]; then
        log_error "Not enough arguments! At least 2 required, got $#"
        usage
        main_exit 1; return
    fi

    local raw_command="${1^^}"
    local ip_arg1
    local ip_arg2

    ip_arg1=$(strip_mask "${2:-}")
    ip_arg2=$(strip_mask "${3:-}")

    log_info "COMMAND: ${raw_command} | IP1: ${ip_arg1:-—} | IP2: ${ip_arg2:-—}"

    # Validate IP1
    if [[ -z "${ip_arg1}" ]]; then
        log_error "IP1 is empty or blank!"
        main_exit 1; return
    fi
    if ! is_valid_ip "${ip_arg1}"; then
        log_error "IP1 is not a valid IPv4 address: '${ip_arg1}'"
        main_exit 1; return
    fi

    # Validate IP2 only for UPDATE
    if [[ "${raw_command}" == "UPDATE" ]]; then
        if [[ $# -ne 3 ]]; then
            log_error "UPDATE requires exactly 2 IPs: IP_OLD IP_NEW, got $(( $# - 1 )) IP(s)"
            usage
            main_exit 1; return
        fi
        if [[ -z "${ip_arg2}" ]]; then
            log_warn "[UPDATE] IP_NEW is empty or blank — nothing to do"
            main_exit 0; return
        fi
        if ! is_valid_ip "${ip_arg2}"; then
            log_error "IP2 (IP_NEW) is not a valid IPv4 address: '${ip_arg2}'"
            main_exit 1; return
        fi
        if [[ "${ip_arg1}" == "${ip_arg2}" ]]; then
            log_warn "[UPDATE] IP_OLD == IP_NEW (${ip_arg1}) — nothing to do"
            main_exit 0; return
        fi
    fi

    case "${raw_command}" in

        NEW | DELETE | RESTRICT | SUSPEND | RESUME)
            if [[ $# -gt 2 ]]; then
                log_error "Command ${raw_command} accepts exactly 1 IP, but got $(( $# - 1 )): ${*:2}"
                usage
                main_exit 1; return
            fi
            case "${raw_command}" in
                NEW)      cmd_new      "${ip_arg1}" ;;
                DELETE)   cmd_delete   "${ip_arg1}" ;;
                RESTRICT) cmd_restrict "${ip_arg1}" ;;
                SUSPEND)  cmd_suspend  "${ip_arg1}" ;;
                RESUME)   cmd_resume   "${ip_arg1}" ;;
            esac
            ;;

        UPDATE)
            cmd_update "${ip_arg1}" "${ip_arg2}"
            ;;

        *)
            log_error "Unknown command: '${raw_command}'"
            usage
            main_exit 5; return
            ;;
    esac

    main_exit "${error_code}"
}

main "$@"; _ec=$?
[[ "${DRY_RUN}" -eq 1 ]] && exit 0 || exit "${_ec}"