#!/bin/sh
baseUrl="https://nuclei.fyi"
downloadPid=""

fail() {
    printf 'nuclei: %s\n' "$1" >&2
    exit 1
}

# Puts the cursor back and stops a download that's still running.
cleanup() {
    [ -n "$downloadPid" ] && kill "$downloadPid" 2>/dev/null
    [ -t 2 ] && printf '\033[?25h' >&2
    rm -f "${tmp:-}"
}

# The bar uses block characters and the Nuclei green where the terminal can show them.
setupBar() {
    case "${LC_ALL:-${LC_CTYPE:-${LANG:-}}}" in
        *UTF-8* | *utf-8* | *UTF8* | *utf8*) full='█' empty='░' ;;
        *) full='#' empty='-' ;;
    esac
    if [ -n "${NO_COLOR:-}" ]; then
        green='' dim='' reset=''
    elif [ "${COLORTERM:-}" = truecolor ] || [ "${COLORTERM:-}" = 24bit ]; then
        green='\033[38;2;77;255;159m' dim='\033[38;2;77;93;84m' reset='\033[0m'
    else
        green='\033[32m' dim='\033[90m' reset='\033[0m'
    fi
}

megabytes() {
    awk -v bytes="$1" 'BEGIN { printf "%.1f", bytes / 1048576 }'
}

drawBar() {
    got=$1 total=$2 width=30
    [ "$got" -gt "$total" ] && got=$total
    filled=$((got * width / total))
    filledPart='' emptyPart='' i=0
    while [ $i -lt $width ]; do
        if [ $i -lt $filled ]; then filledPart="$filledPart$full"; else emptyPart="$emptyPart$empty"; fi
        i=$((i + 1))
    done
    printf "\r  ${green}%s${dim}%s${reset} %3d%%  %s / %s MB " \
        "$filledPart" "$emptyPart" $((got * 100 / total)) "$(megabytes "$got")" "$(megabytes "$total")" >&2
}

# Downloads $1 to $2, drawing a progress bar when someone is watching.
download() {
    if ! command -v curl >/dev/null 2>&1; then
        command -v wget >/dev/null 2>&1 || fail "this needs curl or wget, and neither is installed"
        if [ -t 2 ]; then wget -q --show-progress -O "$2" "$1"; else wget -q -O "$2" "$1"; fi
        return
    fi
    [ -t 2 ] || { curl -fsSL --retry 2 -o "$2" "$1"; return; }

    # A failed size check means the download will fail too, so let it do that quietly.
    headers=$(curl -fsSIL "$1" 2>/dev/null) || { curl -fsSL --retry 2 -o "$2" "$1"; return; }
    total=$(printf '%s\n' "$headers" | tr -d '\r' | awk 'tolower($1) == "content-length:" { size = $2 } END { print size + 0 }')
    if [ "$total" -le 0 ]; then
        # Without a size there's no way to show how far along it is, so let curl draw its own bar.
        curl -fL --retry 2 --progress-bar -o "$2" "$1"
        return
    fi

    setupBar
    printf '\033[?25l' >&2
    curl -fsSL --retry 2 -o "$2" "$1" &
    downloadPid=$!
    while kill -0 "$downloadPid" 2>/dev/null; do
        drawBar "$(wc -c < "$2" 2>/dev/null || echo 0)" "$total"
        sleep 0.1 2>/dev/null || sleep 1
    done
    wait "$downloadPid"
    status=$?
    downloadPid=""
    [ $status -eq 0 ] && drawBar "$total" "$total"
    printf '\033[?25h\n' >&2
    return $status
}

main() {
    set -eu

    case "$(uname -s)" in
        Linux) ;;
        Darwin) fail "there's no Mac version yet. Use Nuclei in your browser instead: $baseUrl/dashboard" ;;
        MINGW* | MSYS* | CYGWIN*) fail "on Windows, run this in PowerShell: irm $baseUrl/install.ps1 | iex" ;;
        *) fail "this system isn't supported yet. Use Nuclei in your browser instead: $baseUrl/dashboard" ;;
    esac
    case "$(uname -m)" in
        x86_64 | amd64) ;;
        *) fail "only 64-bit Intel and AMD are supported so far, not $(uname -m). Use Nuclei in your browser instead: $baseUrl/dashboard" ;;
    esac

    [ -n "${HOME:-}" ] || fail "HOME isn't set, so there's nowhere to install to"
    dir="$HOME/.local/bin"
    mkdir -p "$dir" || fail "can't create $dir"

    tmp="$(mktemp "$dir/.nuclei.XXXXXX")" || fail "can't write to $dir"
    trap cleanup EXIT
    trap 'cleanup; exit 1' HUP INT TERM

    printf 'Downloading Nuclei...\n'
    download "$baseUrl/nuclei" "$tmp" || fail "the download failed. Check your connection and try again"
    chmod 755 "$tmp"

    version="$("$tmp" version 2>/dev/null)" || fail "the downloaded file doesn't run, so nothing was installed"

    mv -f "$tmp" "$dir/nuclei" || fail "can't write to $dir/nuclei"
    trap - EXIT HUP INT TERM

    printf 'Installed %s to %s/nuclei\n' "$version" "$dir"
    case ":${PATH:-}:" in
        *":$dir:"*)
            printf 'Type nuclei to start.\n'
            ;;
        *)
            printf '\n%s is not on your PATH yet. Add this line to your shell startup file,\n' "$dir"
            printf 'such as ~/.bashrc or ~/.zshrc, then open a new terminal:\n\n'
            printf '    export PATH="$HOME/.local/bin:$PATH"\n\n'
            printf 'Or start it now with: %s/nuclei\n' "$dir"
            ;;
    esac
}

main "$@"