#!/usr/bin/bash

set -f
PROGNAME=${0##*/}
TMPDIR=/dev/shm
DESTDIR=/var/lib/geoip
SOURCEDIR=
MAXMIND_KEY=
DATABASE=GeoLite2-Country,GeoLite2-City,GeoLite2-ASN
PRECMD=
POSTCMD=
SKIP_PRECMD=
SKIP_POSTCMD=
TTL=$((86400/2))
PROXY=
MODE=elastic

function exit_usage() {
    local status=${1:-0}
    [[ "$status" != "0" ]] && exec >&2
    echo "\
Usage: $PROGNAME [OPTION...]
Update Maxmind geoip databases

Available options:
  -c, --config  Configuration file
  -m, --mode    Download mode; maxmind (need key), elastic or baseurl:url
  -p, --pre     Pre command to run at start, exit != 0 to abort
  -P, --post    Post command to run after install
  --skip-pre    Skip pre command
  --skip-post   Skip post command
  -k, --key     Maxmind key
  -d, --db      Maxmind database(s), comma separated, eg: GeoLite2-Country
  -D, --dir     Target directory
  -S, --srcdir  Source directory, do not download
  -T, --ttl     Abort is target directory is yonger than TTL
  -q, --quiet   Do not print info when logging on stderr.
  -h, --help    Display this help."
    exit "$status"
}

function _log() {
    local now=$(date +%Y-%m-%dT%H:%M:%S.%3N%:z)
    local LOGGER_SEVERITY=${LOGGER_SEVERITY:-info}
    local STDOUT_LABEL=${STDOUT_LABEL:-INFO}
    [[ -t 0 ]] || logger -t "$PROGNAME[$$]" -p "$LOGGER_SEVERITY" -- "$*"
    [[ -n $QUIET && $STDOUT_LABEL == INFO ]] && return 0
    echo "$now $STDOUT_LABEL $PROGNAME: $*"
}

function info() { LOGGER_SEVERITY=info STDOUT_LABEL=INFO _log "$@"; }
function warning() { LOGGER_SEVERITY=warning STDOUT_LABEL=WARNING _log "$@"; }
function error() { LOGGER_SEVERITY=err STDOUT_LABEL=ERROR _log "$@"; }
function fatal() { LOGGER_SEVERITY=crit STDOUT_LABEL=FATAL _log "$@"; exit 2; }

# not used here but available in pre/post command when using bash functions
function exec_quiet() {(
    [[ -n $QUIET ]] && exec 1>/dev/null
    "$@"
)}

function mycurl() {
    curl ${PROXY:+-x "$PROXY"} -fsS --connect-timeout 10 -L "$@"
}

# $1: database
function getdb() {
    if [[ -n $SOURCEDIR ]]; then
        if [[ -f $SOURCEDIR/$1.tar.gz ]]; then
            cp -a "$SOURCEDIR/$1.tar.gz" ./
        elif [[ -f $SOURCEDIR/$1.mmdb ]]; then
            cp -a "$SOURCEDIR/$1.mmdb" ./
        else
            error "No source found for database $1"
            return 1
        fi
    else
        local fn="getdb_$MODE"
        declare -F -f "$fn" >/dev/null || return 1
        "$fn" "$1"
    fi
}

# $1: database name
function getdb_maxmind() {
    mycurl -o "$1.tar.gz" -m 300 \
        "https://download.maxmind.com/app/geoip_download?suffix=tar.gz&edition_id=$1&license_key=$MAXMIND_KEY"
}

# $1: database name
ELASTIC_MANIFEST=
function getdb_elastic() {
    if [[ -z $ELASTIC_MANIFEST ]]; then
        ELASTIC_MANIFEST=$(mycurl -m 30 "https://geoip.elastic.co/v1/database?elastic_geoip_service_tos=agree" |jq .)
        [[ -z $ELASTIC_MANIFEST ]] && fatal 'Failed to download elastic geoip manifest'
    fi
    local url=$(echo "$ELASTIC_MANIFEST" |jq -r --arg db "$1" '.[] |select(.name==$db+".tgz") |.url')
    [[ -z $url ]] && { error "Cannot get an URL for database $1"; return 1; }
    mycurl -o "$1.tar.gz" -m 300 "$url"
}

# $1: database name
function getdb_baseurl() {
    local ext
    for ext in tar.gz mmdb; do
        mycurl -I -m 30 "$BASEURL/$1.$ext" &>/dev/null || continue
        mycurl -o "$1.$ext" -m 300 "$BASEURL/$1.$ext"
        return $?
    done
}

function on_exit() {
    if [[ -d "$TMPDIR/$PROGNAME" ]]; then
        rm -rf "$TMPDIR/$PROGNAME"
    fi
}

QUIET=${QUIET:-}

while (( $# > 0 )); do
    case "$1" in
        -c|--config) source "$2" || fatal 'Failed to source config file'; shift ;;
        -m|--mode) MODE=$2; shift ;;
        -p|--pre) PRECMD=$2; shift ;;
        -P|--post) POSTCMD=$2; shift ;;
        --skip-pre) SKIP_PRECMD=1 ;;
        --skip-post) SKIP_POSTCMD=1 ;;
        -k|--key) MAXMIND_KEY=$2; shift ;;
        -d|--db) DATABASE=$2; shift ;;
        -D|--dir) DESTDIR=$2; shift ;;
        -S|--srcdir) SOURCEDIR=$2; shift ;;
        -T|--ttl) TTL=$2; shift ;;
        -q|--quiet) QUIET=1 ;;
        -h|--help) exit_usage ;;
        *) exit_usage 1 ;;
    esac
    shift
done

if [[ $MODE == maxmind && -z $MAXMIND_KEY ]]; then
    fatal 'Maxmind key required, check usage'
elif [[ $MODE == baseurl:* ]]; then
    BASEURL=${MODE#*:}
    MODE=baseurl
fi

[[ $DESTDIR == /* ]] || DESTDIR="$PWD/$DESTDIR"
[[ -n $SOURCEDIR && $SOURCEDIR != /* ]] && SOURCEDIR="$PWD/$SOURCEDIR"

if [[ -d $DESTDIR/. ]]; then
    # ignore TTL if DESTDIR is empty
    if ! find "$DESTDIR/" -maxdepth 0 -empty -print -quit |grep -q / &&
            (( ($(date +%s) - $(stat -Lc %Y "$DESTDIR")) < TTL )); then
        info 'Abort due to target directory mtime not old enought'
        exit 0
    fi
else
    mkdir -p "$DESTDIR" || fatal 'Failed to create target directory'
fi

[[ -w $DESTDIR ]] || fatal 'Target directory not writable'

if [[ -n $PRECMD && -z $SKIP_PRECMD ]]; then
    info 'Run pre command'
    if ! REPLY=$(DATABASE="$DATABASE" SOURCEDIR="$SOURCEDIR" \
                    DESTDIR="$DESTDIR" "$PRECMD"); then
        info "Abort due to pre command${REPLY:+ ($REPLY)}"
        exit 0
    fi
fi

trap on_exit EXIT
mkdir -p "$TMPDIR/$PROGNAME" || fatal 'Failed to create temporary directory'
cd "$TMPDIR/$PROGNAME" || fatal 'Failed to cd in temporary directory'

success=0

for db in ${DATABASE//,/ }; do
    info "Download $db"
    if ! getdb "$db"; then
        error "Failed to retrieve database $db"
        continue
    fi
    # assume we have either .tar.gz or .mmdb
    if [[ -f $db.tar.gz ]]; then
        info "Extract $db archive"
        if ! tar xzf "$db.tar.gz"; then
            error "Failed to extract database $db archive"
            continue
        fi
    fi
    (( success++ ))
done

(( success > 0 )) || fatal 'No usable database'

info 'Install databases'
find ./ -name '*.mmdb' -exec cp -a {} "$DESTDIR/" \; ||
    fatal 'Failed to install databases'

if [[ -n $POSTCMD && -z $SKIP_POSTCMD ]]; then
    info 'Run post command'
    DATABASE="$DATABASE" DESTDIR="$DESTDIR" "$POSTCMD" ||
        fatal 'Post comand failed'
fi
