#!/bin/bash
################################################################################
# Script Name: nginx-geoip-setup.sh
# Version: 1.0
# Description: Automate nginx GeoIP2 enrichment for IP intelligence.
#              Installs GeoIP2 module, downloads GeoIP databases (MaxMind or
#              DB-IP free), creates nginx geoip2 variable mappings and an
#              enriched log format.
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - nginx installed and running
#   - Root access
#   - MaxMind account OR use --db-ip for no-signup alternative
#
# Usage:
#   sudo ./nginx-geoip-setup.sh --db-ip
#   sudo ./nginx-geoip-setup.sh --account-id 123456 --license-key ABCDEF
#   sudo ./nginx-geoip-setup.sh --dry-run
#   sudo ./nginx-geoip-setup.sh --remove
#   sudo ./nginx-geoip-setup.sh --skip-packages
#
################################################################################

set -euo pipefail

# --- Configuration ---
GEOIP_CONF="/etc/GeoIP.conf"
GEOIP_DB_DIR="/usr/share/GeoIP"
NGINX_GEOIP_CONF="/etc/nginx/conf.d/geoip2.conf"
NGINX_LOG_CONF="/etc/nginx/enriched-log-format.conf"
CRON_WEEKLY="/etc/cron.weekly/geoip-db-update"
CRON_MONTHLY="/etc/cron.monthly/geoip-db-update"
MARKER_START="# nginx-geoip-managed-start"
MARKER_END="# nginx-geoip-managed-end"
TIMESTAMP=$(date +%s)

DRY_RUN=false
REMOVE=false
SKIP_PACKAGES=false
USE_DBIP=false
ACCOUNT_ID="${MAXMIND_ACCOUNT_ID:-}"
LICENSE_KEY="${MAXMIND_LICENSE_KEY:-}"

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

info()  { echo -e "${GREEN}[OK]${NC} $*"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
step()  { echo -e "${CYAN}[STEP]${NC} $*"; }

usage() {
    cat <<EOF
Usage: sudo $(basename "$0") [OPTIONS]

Sets up nginx GeoIP2 enrichment using MaxMind GeoLite2 or DB-IP free databases.
Creates geoip2 variable mappings and an enriched log format for nginx.

Options:
  --db-ip                Use DB-IP free databases (no signup required)
  --account-id ID        MaxMind account ID (not needed with --db-ip)
  --license-key KEY      MaxMind license key (not needed with --db-ip)
  --dry-run              Show what would be done without making changes
  --remove               Remove managed configs and update script (leaves packages/databases)
  --skip-packages        Skip package installation
  -h, --help             Show this help

Environment variables:
  MAXMIND_ACCOUNT_ID     MaxMind account ID (overridden by --account-id)
  MAXMIND_LICENSE_KEY    MaxMind license key (overridden by --license-key)

Database providers:
  MaxMind GeoLite2       Free, requires account signup at maxmind.com
  DB-IP Lite             Free, no signup, direct download (--db-ip)

Examples:
  sudo $(basename "$0") --db-ip
  sudo $(basename "$0") --account-id 123456 --license-key ABCDEF
  sudo $(basename "$0") --dry-run
  sudo $(basename "$0") --remove
EOF
    exit 0
}

# --- Argument parsing ---
while [[ $# -gt 0 ]]; do
    case "$1" in
        --db-ip)         USE_DBIP=true; shift ;;
        --account-id)    ACCOUNT_ID="$2"; shift 2 ;;
        --license-key)   LICENSE_KEY="$2"; shift 2 ;;
        --dry-run)       DRY_RUN=true; shift ;;
        --remove)        REMOVE=true; shift ;;
        --skip-packages) SKIP_PACKAGES=true; shift ;;
        -h|--help)       usage ;;
        *)               echo "Unknown option: $1"; usage ;;
    esac
done

# --- Checks ---
if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}Error: Run as root (sudo)${NC}" >&2
    exit 1
fi

if ! command -v nginx &>/dev/null; then
    echo -e "${RED}Error: nginx not found${NC}" >&2
    exit 1
fi

# --- Distro detection ---
detect_distro() {
    if [[ -f /etc/os-release ]]; then
        # shellcheck source=/dev/null
        . /etc/os-release
        case "$ID" in
            debian|ubuntu|linuxmint) echo "debian" ;;
            rhel|centos|rocky|alma|fedora|ol) echo "rhel" ;;
            *) echo "unknown" ;;
        esac
    else
        echo "unknown"
    fi
}

DISTRO=$(detect_distro)

backup_file() {
    local file="$1"
    if [[ -f "$file" ]]; then
        cp "$file" "${file}.bak.${TIMESTAMP}"
        warn "Backed up: ${file} → ${file}.bak.${TIMESTAMP}"
    fi
}

# =====================================================
# REMOVE MODE
# =====================================================
if [[ "$REMOVE" == "true" ]]; then
    step "Removing GeoIP2 configuration"

    for file in "$NGINX_GEOIP_CONF" "$NGINX_LOG_CONF" "$CRON_WEEKLY" "$CRON_MONTHLY"; do
        if [[ -f "$file" ]]; then
            if [[ "$DRY_RUN" == "true" ]]; then
                echo "  Would remove: ${file}"
            else
                rm -f "$file"
                info "Removed: ${file}"
            fi
        else
            warn "Not found: ${file} (already removed?)"
        fi
    done

    if [[ "$DRY_RUN" == "true" ]]; then
        echo "  Would run: nginx -t"
        echo "  Would run: systemctl reload nginx"
    else
        step "Testing nginx configuration"
        if nginx -t 2>&1; then
            info "nginx config valid"
        else
            echo -e "${RED}[ERROR] nginx config test failed${NC}" >&2
            exit 1
        fi

        step "Reloading nginx"
        systemctl reload nginx
        info "nginx reloaded"
    fi

    echo ""
    echo -e "${BOLD}GeoIP2 configuration removed.${NC}"
    echo "  Packages and databases left in place."
    echo "  To remove databases: rm -rf ${GEOIP_DB_DIR}/*.mmdb"
    exit 0
fi

# =====================================================
# INSTALL MODE
# =====================================================

# --- Prompt for credentials if not using DB-IP ---
if [[ "$USE_DBIP" != "true" ]]; then
    if [[ -z "$ACCOUNT_ID" ]]; then
        read -rp "MaxMind Account ID (or Ctrl+C and rerun with --db-ip): " ACCOUNT_ID
    fi
    if [[ -z "$LICENSE_KEY" ]]; then
        read -rp "MaxMind License Key: " LICENSE_KEY
    fi

    if [[ -z "$ACCOUNT_ID" || -z "$LICENSE_KEY" ]]; then
        echo -e "${RED}Error: MaxMind credentials required (or use --db-ip for no-signup alternative)${NC}" >&2
        exit 1
    fi
fi

# =====================================================
# Step 1: Install packages
# =====================================================
if [[ "$SKIP_PACKAGES" == "true" ]]; then
    warn "Skipping package installation (--skip-packages)"
else
    step "Installing GeoIP2 packages"

    if [[ "$DISTRO" == "debian" ]]; then
        if [[ "$USE_DBIP" == "true" ]]; then
            PACKAGES="libnginx-mod-http-geoip2 mmdb-bin curl"
        else
            PACKAGES="libnginx-mod-http-geoip2 geoipupdate mmdb-bin"
        fi
        if [[ "$DRY_RUN" == "true" ]]; then
            echo "  Would run: apt-get update && apt-get install -y ${PACKAGES}"
        else
            apt-get update -qq
            # shellcheck disable=SC2086
            apt-get install -y $PACKAGES
            info "Packages installed: ${PACKAGES}"
        fi
    elif [[ "$DISTRO" == "rhel" ]]; then
        if [[ "$USE_DBIP" == "true" ]]; then
            PACKAGES="nginx-mod-http-geoip2 mmdb-bin curl"
        else
            PACKAGES="nginx-mod-http-geoip2 geoipupdate mmdb-bin"
        fi
        if [[ "$DRY_RUN" == "true" ]]; then
            echo "  Would run: dnf install -y ${PACKAGES}"
        else
            # shellcheck disable=SC2086
            dnf install -y $PACKAGES
            info "Packages installed: ${PACKAGES}"
        fi
    else
        echo -e "${RED}Error: Unsupported distro. Install packages manually and rerun with --skip-packages${NC}" >&2
        exit 1
    fi
fi

# =====================================================
# Step 2: Write GeoIP.conf (MaxMind only)
# =====================================================
if [[ "$USE_DBIP" != "true" ]]; then
    step "Writing MaxMind configuration to ${GEOIP_CONF}"

    GEOIP_CONF_CONTENT="${MARKER_START}
AccountID ${ACCOUNT_ID}
LicenseKey ${LICENSE_KEY}
EditionIDs GeoLite2-City GeoLite2-ASN GeoLite2-Country
DatabaseDirectory ${GEOIP_DB_DIR}
${MARKER_END}"

    if [[ "$DRY_RUN" == "true" ]]; then
        echo "  Would create: ${GEOIP_CONF}"
    else
        backup_file "$GEOIP_CONF"
        echo "$GEOIP_CONF_CONTENT" > "$GEOIP_CONF"
        chmod 600 "$GEOIP_CONF"
        info "Created: ${GEOIP_CONF}"
    fi
fi

# =====================================================
# Step 3: Download GeoIP databases
# =====================================================
if [[ "$USE_DBIP" == "true" ]]; then
    step "Downloading DB-IP free databases to ${GEOIP_DB_DIR}"

    DBIP_MONTH=$(date +%Y-%m)

    if [[ "$DRY_RUN" == "true" ]]; then
        echo "  Would download: dbip-country-lite-${DBIP_MONTH}.mmdb.gz"
        echo "  Would download: dbip-asn-lite-${DBIP_MONTH}.mmdb.gz"
    else
        mkdir -p "$GEOIP_DB_DIR"
        TMPDIR=$(mktemp -d)
        trap 'rm -rf "$TMPDIR"' EXIT

        curl -fsSL "https://download.db-ip.com/free/dbip-country-lite-${DBIP_MONTH}.mmdb.gz" -o "${TMPDIR}/country.mmdb.gz"
        curl -fsSL "https://download.db-ip.com/free/dbip-asn-lite-${DBIP_MONTH}.mmdb.gz" -o "${TMPDIR}/asn.mmdb.gz"
        gunzip -f "${TMPDIR}/country.mmdb.gz"
        gunzip -f "${TMPDIR}/asn.mmdb.gz"
        mv "${TMPDIR}/country.mmdb" "${GEOIP_DB_DIR}/GeoLite2-City.mmdb"
        mv "${TMPDIR}/asn.mmdb" "${GEOIP_DB_DIR}/GeoLite2-ASN.mmdb"
        info "Downloaded: DB-IP Country → GeoLite2-City.mmdb"
        info "Downloaded: DB-IP ASN → GeoLite2-ASN.mmdb"
    fi
else
    step "Downloading MaxMind GeoIP2 databases to ${GEOIP_DB_DIR}"

    if [[ "$DRY_RUN" == "true" ]]; then
        echo "  Would run: mkdir -p ${GEOIP_DB_DIR}"
        echo "  Would run: geoipupdate"
    else
        mkdir -p "$GEOIP_DB_DIR"
        geoipupdate -v
        info "GeoIP2 databases downloaded"
    fi
fi

# =====================================================
# Step 4: Create nginx geoip2 config
# =====================================================
step "Creating nginx GeoIP2 config at ${NGINX_GEOIP_CONF}"

NGINX_GEOIP_CONTENT="${MARKER_START}
# GeoIP2 variable mappings — generated by nginx-geoip-setup.sh

geoip2 ${GEOIP_DB_DIR}/GeoLite2-City.mmdb {
    auto_reload 60m;
    \$geoip2_country_code   country iso_code;
    \$geoip2_country_name   country names en;
    \$geoip2_city_name      city names en;
}

geoip2 ${GEOIP_DB_DIR}/GeoLite2-ASN.mmdb {
    auto_reload 60m;
    \$geoip2_asn            autonomous_system_number;
    \$geoip2_asn_org        autonomous_system_organization;
}
${MARKER_END}"

if [[ "$DRY_RUN" == "true" ]]; then
    echo "  Would create: ${NGINX_GEOIP_CONF}"
else
    backup_file "$NGINX_GEOIP_CONF"
    echo "$NGINX_GEOIP_CONTENT" > "$NGINX_GEOIP_CONF"
    info "Created: ${NGINX_GEOIP_CONF}"
fi

# =====================================================
# Step 5: Create enriched log format
# =====================================================
step "Creating enriched log format at ${NGINX_LOG_CONF}"

NGINX_LOG_CONTENT="${MARKER_START}
# Enriched log format with GeoIP2 data — generated by nginx-geoip-setup.sh

log_format enriched '\$remote_addr - \$remote_user [\$time_local] '
                    '\"\$request\" \$status \$body_bytes_sent '
                    '\"\$http_referer\" \"\$http_user_agent\" '
                    '\$geoip2_country_code \"\$geoip2_asn_org\"';
${MARKER_END}"

if [[ "$DRY_RUN" == "true" ]]; then
    echo "  Would create: ${NGINX_LOG_CONF}"
else
    backup_file "$NGINX_LOG_CONF"
    echo "$NGINX_LOG_CONTENT" > "$NGINX_LOG_CONF"
    info "Created: ${NGINX_LOG_CONF}"
fi

# =====================================================
# Step 6: Patch nginx.conf to use enriched log format
# =====================================================
NGINX_CONF="/etc/nginx/nginx.conf"

if grep -q "enriched-log-format.conf" "$NGINX_CONF" 2>/dev/null; then
    warn "nginx.conf already includes enriched log format — skipping"
elif grep -q "access_log.*enriched" "$NGINX_CONF" 2>/dev/null; then
    warn "nginx.conf already uses enriched format — skipping"
else
    step "Patching ${NGINX_CONF} to use enriched log format"

    if [[ "$DRY_RUN" == "true" ]]; then
        echo "  Would add: include ${NGINX_LOG_CONF}; before access_log"
        echo "  Would change: access_log line to use enriched format"
    else
        backup_file "$NGINX_CONF"

        # Insert include before the access_log line, and append 'enriched' to it
        sed -i "/^[[:space:]]*access_log[[:space:]]/ {
            i\\\\tinclude ${NGINX_LOG_CONF};
            s|access_log \(.*\);|access_log \1 enriched;|
        }" "$NGINX_CONF"

        # Only patch if access_log doesn't already have a format name (avoid double-appending)
        # Check if the sed worked
        if grep -q "enriched-log-format.conf" "$NGINX_CONF"; then
            info "Patched: ${NGINX_CONF}"
        else
            warn "Could not find access_log line in ${NGINX_CONF} — add manually:"
            warn "  include ${NGINX_LOG_CONF};"
            warn "  access_log /var/log/nginx/access.log enriched;"
        fi
    fi
fi

# =====================================================
# Step 7: Create cron job for database updates
# =====================================================
if [[ "$USE_DBIP" == "true" ]]; then
    CRON_FILE="$CRON_MONTHLY"
    step "Creating monthly DB-IP update script at ${CRON_FILE}"

    CRON_CONTENT="#!/bin/sh
# Monthly DB-IP database update — generated by nginx-geoip-setup.sh
MONTH=\$(date +%Y-%m)
TMPDIR=\$(mktemp -d)
trap 'rm -rf \"\$TMPDIR\"' EXIT
if curl -fsSL \"https://download.db-ip.com/free/dbip-country-lite-\${MONTH}.mmdb.gz\" -o \"\${TMPDIR}/country.mmdb.gz\" && \\
   curl -fsSL \"https://download.db-ip.com/free/dbip-asn-lite-\${MONTH}.mmdb.gz\" -o \"\${TMPDIR}/asn.mmdb.gz\" && \\
   gunzip \"\${TMPDIR}/country.mmdb.gz\" && \\
   gunzip \"\${TMPDIR}/asn.mmdb.gz\"; then
  mv \"\${TMPDIR}/country.mmdb\" \"${GEOIP_DB_DIR}/GeoLite2-City.mmdb\"
  mv \"\${TMPDIR}/asn.mmdb\" \"${GEOIP_DB_DIR}/GeoLite2-ASN.mmdb\"
  logger -t dbip-update \"DB-IP databases updated\"
else
  logger -t dbip-update -p user.err \"DB-IP download failed — kept existing databases\"
fi"
else
    CRON_FILE="$CRON_WEEKLY"
    step "Creating weekly geoipupdate script at ${CRON_FILE}"

    CRON_CONTENT="#!/bin/sh
# Weekly GeoIP2 database update — generated by nginx-geoip-setup.sh
/usr/bin/geoipupdate -s 2>&1 | logger -t geoipupdate"
fi

if [[ "$DRY_RUN" == "true" ]]; then
    echo "  Would create: ${CRON_FILE}"
else
    backup_file "$CRON_FILE"
    echo "$CRON_CONTENT" > "$CRON_FILE"
    chmod 755 "$CRON_FILE"
    info "Created: ${CRON_FILE}"
fi

# =====================================================
# Step 8: Validate nginx config
# =====================================================
step "Testing nginx configuration"

if [[ "$DRY_RUN" == "true" ]]; then
    echo "  Would run: nginx -t"
else
    if nginx -t 2>&1; then
        info "nginx config valid"
    else
        echo -e "${RED}[ERROR] nginx config test failed${NC}" >&2
        echo "  Restore backups (.bak.${TIMESTAMP}) and check configuration" >&2
        exit 1
    fi
fi

# =====================================================
# Step 9: Reload nginx
# =====================================================
step "Reloading nginx"

if [[ "$DRY_RUN" == "true" ]]; then
    echo "  Would run: systemctl reload nginx"
else
    systemctl reload nginx
    info "nginx reloaded"
fi

# =====================================================
# Summary
# =====================================================
echo ""
echo -e "${BOLD}Done.${NC}"
echo ""
if [[ "$USE_DBIP" == "true" ]]; then
    echo "  Provider:        DB-IP Lite (no signup)"
else
    echo "  Provider:        MaxMind GeoLite2"
    echo "  GeoIP config:    ${GEOIP_CONF}"
fi
echo "  Database dir:    ${GEOIP_DB_DIR}"
echo "  nginx geoip2:    ${NGINX_GEOIP_CONF}"
echo "  Log format:      ${NGINX_LOG_CONF}"
echo "  Cron job:        ${CRON_FILE}"
echo ""
echo -e "${BOLD}Verify:${NC}"
echo "  tail -5 /var/log/nginx/access.log"
echo "  (should show country code and ASN org at end of each line)"
echo ""
echo "  To remove:  sudo $(basename "$0") --remove"