#!/bin/sh

set -eu

APPDIR=/usr/share/rekono
DEFAULTS=/etc/default/rekono
CONFDIR=/etc/rekono
NGINX_CONF=/etc/rekono/nginx.conf
PROGRESS=/run/rekono/progress
KBX_DUMP=/usr/libexec/rekono/rekono-kbx-dump
UPGRADE_GUIDE=https://github.com/pablosnt/rekono/blob/main/UPGRADE.md
WORKERS='tasks monitor findings executions'
UNITS='rekono-nginx.service rekono-frontend.service rekono-backend.service
       rekono-telegram-bot.service
       rekono-worker@tasks.service rekono-worker@monitor.service
       rekono-worker@findings.service rekono-worker@executions.service
       rekono-initializer.service rekono-valkey.service'


if [ -r "${DEFAULTS}" ]; then
    . "${DEFAULTS}"
fi



if [ -t 1 ] && [ -z "${NO_COLOR:-}" ]; then
    RESET=$(printf '\033[0m')
    BOLD=$(printf '\033[1m')
    DIM=$(printf '\033[2m')
    RED=$(printf '\033[1;31m')
    GREEN=$(printf '\033[1;32m')
    YELLOW=$(printf '\033[1;33m')
    CYAN=$(printf '\033[1;36m')
    TTY=yes
else
    RESET='' BOLD='' DIM='' RED='' GREEN='' YELLOW='' CYAN=''
    TTY=no
fi

port() {
    sed -n 's/^[[:space:]]*listen[[:space:]]\+127\.0\.0\.1:\([0-9]\+\).*/\1/p' "${NGINX_CONF}" | head -n 1
}

url() {
    echo "https://127.0.0.1:$(port)"
}

say() {
    printf '%s\n' "$*"
}

die() {
    printf '%s[-]%s %s\n' "${RED}" "${RESET}" "$*" >&2
    exit 1
}

control() {
    error="$(systemctl --no-ask-password "$@" 2>&1)" && return 0
    if [ "$(id -u)" -eq 0 ] && [ -n "${error}" ]; then
        printf '%s\n' "${error}" >&2
    fi
    return 1
}

ESCALATE_ACTION='manage the Rekono services'

escalate() {
    if [ "$(id -u)" -eq 0 ]; then
        die "Could not ${ESCALATE_ACTION}."
    fi
    printf '%s[i]%s Root privileges are needed to %s.\n' "${YELLOW}" "${RESET}" "${ESCALATE_ACTION}" >&2
    if [ -t 0 ] && command -v sudo >/dev/null 2>&1; then
        exec sudo -- "$0" "$@"
    elif command -v pkexec >/dev/null 2>&1; then
        exec pkexec "$0" "$@"
    fi
    die "Not allowed to ${ESCALATE_ACTION}, and no way to ask for permission."
}

step() {
    label="$1"
    shift
    log=$(mktemp)
    if [ "${TTY}" = no ]; then
        printf '  %s ... ' "${label}"
        if "$@" >"${log}" 2>&1; then
            printf 'ok\n'
            rm -f "${log}"
            return 0
        fi
        printf 'FAILED\n'
        sed 's/^/      /' "${log}" >&2
        rm -f "${log}"
        return 1
    fi

    "$@" >"${log}" 2>&1 &
    pid=$!
    frames='- \ | /'
    while kill -0 "${pid}" 2>/dev/null; do
        for frame in ${frames}; do
            printf '\r  %s%s%s %s' "${CYAN}" "${frame}" "${RESET}" "${label}"
            sleep 0.1
        done
    done
    if wait "${pid}"; then
        printf '\r  %s+%s %s\033[K\n' "${GREEN}" "${RESET}" "${label}"
        rm -f "${log}"
        return 0
    fi
    printf '\r  %sx%s %s\033[K\n' "${RED}" "${RESET}" "${label}"
    sed 's/^/      /' "${log}" >&2
    rm -f "${log}"
    return 1
}

check_port() {
    p="$(port)"
    [ -n "${p}" ] || die "Could not read the listening port from ${NGINX_CONF}"
    systemctl is-active --quiet rekono-nginx.service && return 0
    ss -ltnH "sport = :${p}" 2>/dev/null | grep -q . || return 0
    printf '%s[-]%s Port %s is already in use:\n\n' "${RED}" "${RESET}" "${p}" >&2
    ss -ltnp "sport = :${p}" 2>/dev/null | sed 's/^/      /' >&2
    printf '\n      Stop that process, or change the port in %s\n' "${NGINX_CONF}" >&2
    exit 1
}

explain_failure() {
    systemctl status --no-pager --lines 0 "$1" 2>&1 | sed -n '1,6p'

    if [ -n "$(systemctl show "$1" --property LoadError --value 2>/dev/null)" ]; then
        printf 'systemd cannot load this unit. Clear its leftover state with:\n'
        printf '  sudo systemctl reset-failed %s\n' "$1"
        printf '  sudo systemctl daemon-reload\n'
        return 0
    fi

    logs="$(journalctl --unit "$1" --no-pager --lines 25 2>/dev/null || true)"
    if [ -n "${logs}" ]; then
        printf '%s\n' "${logs}"
    else
        printf 'Run "sudo journalctl --unit %s" to see why.\n' "$1"
    fi
}

wait_unit() {
    unit="$1"
    budget="${2:-120}"
    elapsed=0
    stalled=0
    while [ "${elapsed}" -lt "${budget}" ]; do
        case "$(systemctl is-active "${unit}" 2>/dev/null || true)" in
            active)
                return 0
                ;;
            failed)
                explain_failure "${unit}" >&2
                return 1
                ;;
            *)
                if systemctl list-jobs --no-legend 2>/dev/null \
                    | awk -v u="${unit}" '$2 == u { found = 1 } END { exit !found }'; then
                    stalled=0
                else
                    stalled=$((stalled + 1))
                    if [ "${stalled}" -ge 5 ]; then
                        printf '%s was never started.\n' "${unit}" >&2
                        explain_failure "${unit}" >&2
                        return 1
                    fi
                fi
                ;;
        esac
        elapsed=$((elapsed + 1))
        sleep 1
    done
    printf 'Timed out after %ss waiting for %s.\n' "${budget}" "${unit}" >&2
    explain_failure "${unit}" >&2
    return 1
}

wait_phase() {
    phase="$1"
    budget="${2:-600}"
    elapsed=0
    while [ "${elapsed}" -lt "${budget}" ]; do
        if [ -r "${PROGRESS}" ] && grep -qx "${phase}" "${PROGRESS}" 2>/dev/null; then
            return 0
        fi
        case "$(systemctl is-active rekono-initializer.service 2>/dev/null || true)" in
            failed)
                explain_failure rekono-initializer.service >&2
                return 1
                ;;
            active)
                # The unit finished without recording this phase, which means
                # the phase list here and the unit file have drifted apart.
                if [ -r "${PROGRESS}" ] && ! grep -qx "${phase}" "${PROGRESS}" 2>/dev/null; then
                    printf 'Initialization finished without the %s phase.\n' "${phase}" >&2
                    return 1
                fi
                ;;
        esac
        elapsed=$((elapsed + 1))
        sleep 1
    done
    printf 'Timed out after %ss waiting for the %s phase.\n' "${budget}" "${phase}" >&2
    explain_failure rekono-initializer.service >&2
    return 1
}

open_browser() {
    target="$1"

    user="${SUDO_USER:-}"
    if [ -z "${user}" ] && [ -n "${PKEXEC_UID:-}" ]; then
        user="$(id -un "${PKEXEC_UID}" 2>/dev/null || true)"
    fi
    if [ -z "${user}" ] || [ "${user}" = 'root' ]; then
        detach xdg-open "${target}"
        return 0
    fi

    uid="$(id -u "${user}" 2>/dev/null)" || return 1
    home="$(getent passwd "${user}" | cut -d: -f6)"
    runtime="/run/user/${uid}"

    session="$(loginctl list-sessions --no-legend 2>/dev/null \
        | awk -v u="${user}" '$3 == u { print $1; exit }')"
    display="$(loginctl show-session "${session}" --value --property Display 2>/dev/null)"
    [ -n "${display}" ] || display="${DISPLAY:-:0}"

    set -- HOME="${home}" \
           XDG_RUNTIME_DIR="${runtime}" \
           DBUS_SESSION_BUS_ADDRESS="unix:path=${runtime}/bus" \
           DISPLAY="${display}"
    if [ -S "${runtime}/wayland-0" ]; then
        set -- "$@" WAYLAND_DISPLAY="${WAYLAND_DISPLAY:-wayland-0}"
    fi
    if [ -r "${home}/.Xauthority" ]; then
        set -- "$@" XAUTHORITY="${XAUTHORITY:-${home}/.Xauthority}"
    fi

    detach runuser -u "${user}" -- env "$@" xdg-open "${target}"
    return 0
}

detach() {
    setsid "$@" >/dev/null 2>&1 </dev/null &
}

banner() {
    printf '\n%s+-( %sRekono%s )\n' "${CYAN}" "${BOLD}" "${RESET}${CYAN}"
    printf '|%s\n' "${RESET}"
    printf '%s|%s  %s[*]%s Web UI: %s%s%s\n' "${CYAN}" "${RESET}" "${GREEN}" "${RESET}" "${GREEN}" "$(url)" "${RESET}"
    printf '%s|%s\n' "${CYAN}" "${RESET}"
    printf '%s|%s  Default credentials:\n' "${CYAN}" "${RESET}"
    printf '%s|%s    user:     %srekono%s\n' "${CYAN}" "${RESET}" "${BOLD}" "${RESET}"
    printf '%s|%s    password: %srekono%s\n' "${CYAN}" "${RESET}" "${BOLD}" "${RESET}"
    printf '%s|%s\n' "${CYAN}" "${RESET}"
    printf '%s|%s  %s[i]%s Change this password after your first login.\n' "${CYAN}" "${RESET}" "${YELLOW}" "${RESET}"
    printf '%s|%s\n' "${CYAN}" "${RESET}"
    printf '%s|%s  %sThe certificate is self-signed, so the browser will warn you.%s\n' "${CYAN}" "${RESET}" "${DIM}" "${RESET}"
    printf '%s|%s  %sRun "rekono status" for the files you may want to edit.%s\n' "${CYAN}" "${RESET}" "${DIM}" "${RESET}"
    printf '%s+-%s\n\n' "${CYAN}" "${RESET}"
}

pause_before_closing() {
    [ -t 0 ] || return 0
    printf '\n%s[i]%s Press Enter to close this window.\n' "${CYAN}" "${RESET}"
    read -r _unused || true
}

kbx_choose() {
    directory="$1"

    printf '\n'
    printf '  %sThe data of rekono-kbx could not be imported.%s\n\n' "${BOLD}" "${RESET}"
    printf '  Nothing has been deleted. Your Rekono 1.x data is still in %s%s/data%s\n\n' "${BOLD}" "${directory}" "${RESET}"
    printf '  Starting Rekono now creates an empty database, and that data is not imported into it. You can:\n\n'
    printf '    %s1.%s Stop here, solve the problem above, and run "rekono" again.\n' "${BOLD}" "${RESET}"
    printf '    %s2.%s Stop here and import the data yourself, following %s\n' "${BOLD}" "${RESET}" "${UPGRADE_GUIDE}"
    printf '    %s3.%s Continue, and start Rekono with an empty database. The old data stays on disk, but Rekono stops offering to import it.\n\n' "${BOLD}" "${RESET}"

    if [ ! -t 0 ]; then
        printf '  %s[i]%s Run "rekono" from a terminal to choose, or set REKONO_KBX_SKIP=1 in %s to start without the rekono-kbx data.\n\n' "${YELLOW}" "${RESET}" "${DEFAULTS}"
        return 1
    fi

    printf '  Continue and start Rekono with an empty database? [y/N] '
    read -r answer || answer=''
    case "${answer}" in
        y|Y|yes|YES|Yes)
            "${KBX_DUMP}" discard "${directory}" 2>&1 | sed 's/^/  /'
            printf '\n'
            return 0
            ;;
        *)
            printf '\n  Rekono was not started. Your rekono-kbx data is untouched.\n\n'
            return 1
            ;;
    esac
}

cmd_start() {
    for argument in "$@"; do
        case "${argument}" in
            --keep-open) trap pause_before_closing EXIT ;;
            *)           usage >&2; exit 1 ;;
        esac
    done

    kbx=''
    if [ -x "${KBX_DUMP}" ]; then
        # Only the exit code 1 means that there is nothing to import
        kbx=$("${KBX_DUMP}" detect 2>&1) && found=0 || found=$?
        [ "${found}" -le 1 ] || die "The data of rekono-kbx cannot be imported: ${kbx}"
        [ "${found}" -eq 0 ] || kbx=''
    fi

    check_port

    if [ -n "${kbx}" ] && [ "$(id -u)" -ne 0 ]; then
        ESCALATE_ACTION='import the data of rekono-kbx'
        escalate start "$@"
    fi

    say "${BOLD}Starting Rekono${RESET}"

    # The data has to be imported before the initializer looks for it
    if [ -n "${kbx}" ]; then
        step 'Importing the data of rekono-kbx' "${KBX_DUMP}" run || kbx_choose "${kbx}" || exit 1
    fi

    # Ask for the whole target at once and then follow it
    control start --no-block rekono.target || escalate start "$@"

    step 'Valkey'                 wait_unit rekono-valkey.service       60
    step 'Configuration'          wait_phase configuration             120
    step 'PostgreSQL'             wait_phase database                  300
    step 'Python environment'     wait_phase environment              1900
    step 'Data import'            wait_phase upgrade                   900
    step 'Database migrations'    wait_phase migrations                620
    step 'Tools and wordlists'    wait_phase toolkit                   620
    step 'Default account'        wait_phase administrator             320
    step 'Monitoring'             wait_phase monitor                   620
    for queue in ${WORKERS}; do
        step "Worker: ${queue}"   wait_unit "rekono-worker@${queue}.service" 180
    done
    step 'Telegram bot'           wait_unit rekono-telegram-bot.service 180
    step 'Backend'                wait_unit rekono-backend.service     200
    step 'Frontend'               wait_unit rekono-frontend.service    140
    step 'Nginx'                  wait_unit rekono-nginx.service       140

    banner

    if command -v xdg-open >/dev/null 2>&1; then
        open_browser "$(url)"
    fi
}

cmd_stop() {
    say "${BOLD}Stopping Rekono${RESET}"
    step 'Rekono services' control stop rekono.target ${UNITS} || escalate stop
    say ''
    say 'Rekono is stopped. Its PostgreSQL cluster and Valkey instance are private, so nothing else on this machine is affected.'
}

unit_state() {
    state="$(systemctl is-active "$1" 2>/dev/null || true)"
    case "${state}" in
        active) printf '  %s+%s %-34s %s\n' "${GREEN}" "${RESET}" "$2" "${state}" ;;
        *)      printf '  %s-%s %-34s %s\n' "${DIM}" "${RESET}" "$2" "${state:-unknown}" ;;
    esac
}

cmd_status() {
    say ''
    say "${BOLD}Services${RESET}"
    if pg_lsclusters --no-header 2>/dev/null | awk -v v="${PG_VERSION}" '$1 == v && $2 == "rekono" { found = 1 } END { exit !found }'; then
        unit_state "postgresql@${PG_VERSION}-rekono.service" \
            "PostgreSQL (cluster ${PG_VERSION}/rekono)"
    else
        printf '  %s-%s %-34s %s\n' "${DIM}" "${RESET}" 'PostgreSQL' 'not created yet'
    fi
    unit_state rekono-valkey.service        'Valkey'
    unit_state rekono-initializer.service   'Initializer'
    unit_state rekono-backend.service       'Backend'
    unit_state rekono-frontend.service      'Frontend'
    unit_state rekono-nginx.service         'Nginx'
    for queue in ${WORKERS}; do
        unit_state "rekono-worker@${queue}.service" "Worker: ${queue}"
    done
    unit_state rekono-telegram-bot.service  'Telegram bot'

    if systemctl is-active --quiet rekono-nginx.service; then
        say ''
        say "${BOLD}Web UI${RESET}"
        printf '  %s%s%s\n' "${GREEN}" "$(url)" "${RESET}"
    fi

    say ''
    say "${BOLD}Edit${RESET}"
    printf '  %-18s %s\n' 'Configuration'   "${CONFDIR}/config.yaml"
    printf '  %-18s %s\n' 'Nginx'           "${NGINX_CONF}"
    printf '  %-18s %s\n' 'Valkey'          "${CONFDIR}/valkey.conf"
    printf '  %-18s %s\n' 'TLS certificate' "${CONFDIR}/tls/"
    printf '  %-18s %s\n' 'Services'        "${DEFAULTS}  (EXECUTION_WORKERS=${EXECUTION_WORKERS:-?})"

    say ''
    say "${BOLD}Data${RESET}"
    printf '  %-18s %s\n' 'Rekono Home'         "${REKONO_HOME}"
    printf '  %-18s %s\n' 'Logs'            '/var/log/rekono/  (journalctl -u "rekono-*")'

    for stale in $(pg_lsclusters --no-header 2>/dev/null | awk -v v="${PG_VERSION}" '$2 == "rekono" && $1 != v { print $1 }'); do
        say ''
        printf '  %s[i]%s Rekono is on PostgreSQL %s. Its old %s/rekono cluster can be removed:\n' "${YELLOW}" "${RESET}" "${PG_VERSION}" "${stale}"
        printf '      sudo pg_dropcluster %s rekono --stop\n' "${stale}"
    done

    if [ -x "${KBX_DUMP}" ] && kbx="$("${KBX_DUMP}" detect 2>/dev/null)"; then
        say ''
        printf '  %s[i]%s Data from rekono-kbx is waiting in %s/data. Rekono imports it the next time you run "rekono".\n' "${YELLOW}" "${RESET}" "${kbx}"
    fi

    if [ -x "${KBX_DUMP}" ] && kbx="$("${KBX_DUMP}" leftover 2>/dev/null)"; then
        say ''
        printf '  %s[i]%s The data of rekono-kbx was imported. Its old directory can be removed:\n' "${YELLOW}" "${RESET}"
        printf '      rm -rf %s\n' "${kbx}"
    fi

    say ''
}

cmd_version() {
    rekono="$(sed -n "s/^VERSION[[:space:]]*=[[:space:]]*[\"']\([^\"']*\)[\"'].*/\1/p" "${APPDIR}/backend/rekono/settings.py" 2>/dev/null | head -n 1)"
    package="$(dpkg-query --showformat='${Version}' --show rekono 2>/dev/null || true)"
    if [ -n "${rekono}" ] && [ -n "${package}" ]; then
        printf 'Rekono %s (Debian package %s)\n' "${rekono}" "${package}"
    elif [ -n "${rekono}${package}" ]; then
        printf 'Rekono %s\n' "${rekono}${package}"
    else
        die 'Could not determine the Rekono version.'
    fi
}

usage() {
    cat <<EOF
Usage: rekono [command]

Commands:
  start     start Rekono and open the web interface (default)
            --keep-open  wait for Enter before returning
  stop      stop Rekono
  status    show service states and the paths you may want to edit
  version   show the Rekono and package versions

Rekono is not started at boot. Run 'rekono' when you want to use it.
EOF
}

command="${1:-start}"
if [ "$#" -gt 0 ]; then
    shift
fi

case "${command}" in
    start)                cmd_start "$@" ;;
    stop)                 cmd_stop ;;
    status)               cmd_status ;;
    version|-V|--version) cmd_version ;;
    -h|--help|help)       usage ;;
    *)                    usage >&2; exit 1 ;;
esac
