#!/usr/bin/env bash

#########################################################################################
#### add-nginx-block-head.sh — Block HEAD requests in Nginx (HestiaCP compatible)    ####
#### Adds a 444 drop rule for HEAD method crawlers/scrapers.                         ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   sudo ./add-nginx-block-head.sh                                                ####
####   sudo ./add-nginx-block-head.sh --dry-run                                      ####
####   sudo ./add-nginx-block-head.sh --remove                                       ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
DRY_RUN=false
REMOVE=false
SNIPPET_NAME="nginx.conf_block_head"

# ── Colors ────────────────────────────────────────────────────────────
if [[ -t 1 ]]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[0;33m'
    BOLD='\033[1m'
    RESET='\033[0m'
else
    RED="" GREEN="" YELLOW="" BOLD="" RESET=""
fi

log()  { echo -e "${GREEN}[OK]${RESET}    $*"; }
warn() { echo -e "${YELLOW}[WARN]${RESET}  $*"; }
err()  { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
info() { echo -e "${BOLD}[INFO]${RESET}  $*"; }

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

Block HTTP HEAD requests in Nginx by adding per-domain snippets that return
444 (drop connection) for HEAD method requests.

Works with HestiaCP — adds custom nginx config snippets to each domain's
server block via the conf/web include pattern.

Options:
  --dry-run       Show what would be done without making changes
  --remove        Remove the block-head snippets from all domains
  -h, --help      Show this help

What it does:
  1. Finds all HestiaCP web domains
  2. Adds a custom nginx snippet to each domain's conf/web directory
  3. Tests Nginx configuration (nginx -t)
  4. Reloads Nginx if the test passes
  5. Rolls back ALL changes if the test fails

The snippet (added inside each domain's server block):
  if (\$request_method = HEAD) {
      return 444;
  }

444 drops the connection silently — no response headers, no body.
Bots get nothing and waste their connection.
EOF
}

# ── Parse args ────────────────────────────────────────────────────────
while [[ $# -gt 0 ]]; do
    case "$1" in
        --dry-run)  DRY_RUN=true ;;
        --remove)   REMOVE=true ;;
        -h|--help)  usage; exit 0 ;;
        *)          err "Unknown option: $1"; usage; exit 1 ;;
    esac
    shift
done

# ── Preflight checks ─────────────────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
    err "Must run as root (sudo)"
    exit 1
fi

if ! command -v nginx &>/dev/null; then
    err "Nginx not found"
    exit 1
fi

if ! command -v v-list-users &>/dev/null; then
    err "HestiaCP not found (v-list-users missing)"
    exit 1
fi

# ── Snippet content ──────────────────────────────────────────────────
SNIPPET_CONTENT='# Block HEAD request crawlers/scrapers
# Added by add-nginx-block-head.sh
# Returns 444 (drop connection) — no response sent to bot
if ($request_method = HEAD) {
    return 444;
}'

# ── Find all HestiaCP domains ────────────────────────────────────────
get_all_domain_dirs() {
    local users
    users=$(v-list-users plain 2>/dev/null | cut -f1)

    for user in $users; do
        local user_conf="/home/${user}/conf/web"
        [[ -d "$user_conf" ]] || continue

        # Find domain directories by looking for nginx.conf files
        for nginx_conf in "${user_conf}"/*/nginx.conf; do
            [[ -f "$nginx_conf" ]] || continue
            dirname "$nginx_conf"
        done
    done
}

# ── Remove mode ───────────────────────────────────────────────────────
if [[ "$REMOVE" == "true" ]]; then
    removed=0

    while IFS= read -r domain_dir; do
        snippet="${domain_dir}/${SNIPPET_NAME}"
        ssl_snippet="${domain_dir}/nginx.ssl.conf_block_head"

        for f in "$snippet" "$ssl_snippet"; do
            if [[ -f "$f" ]]; then
                if [[ "$DRY_RUN" == "true" ]]; then
                    info "Would remove: ${f}"
                else
                    rm -f "$f"
                    log "Removed ${f}"
                fi
                ((removed++)) || true
            fi
        done
    done < <(get_all_domain_dirs)

    if [[ $removed -eq 0 ]]; then
        info "No block-head snippets found — nothing to remove"
        exit 0
    fi

    if [[ "$DRY_RUN" == "true" ]]; then
        info "Would test and reload Nginx"
        exit 0
    fi

    if nginx -t 2>/dev/null; then
        systemctl reload nginx
        log "Nginx reloaded — HEAD requests are now allowed"
    else
        err "Nginx config test failed after removal — check your config"
        exit 1
    fi
    exit 0
fi

# ── Install mode ──────────────────────────────────────────────────────
domain_dirs=()
while IFS= read -r dir; do
    domain_dirs+=("$dir")
done < <(get_all_domain_dirs)

if [[ ${#domain_dirs[@]} -eq 0 ]]; then
    err "No HestiaCP web domains found"
    exit 1
fi

info "Found ${#domain_dirs[@]} domain config(s)"
echo ""

created=0
skipped=0
created_files=()

for domain_dir in "${domain_dirs[@]}"; do
    domain_name=$(basename "$domain_dir")

    # Add snippet for both HTTP and HTTPS server blocks
    for conf_type in "" ".ssl"; do
        if [[ -n "$conf_type" ]]; then
            snippet="${domain_dir}/nginx${conf_type}.conf_block_head"
        else
            snippet="${domain_dir}/${SNIPPET_NAME}"
        fi

        # Check the main config exists for this type
        main_conf="${domain_dir}/nginx${conf_type}.conf"
        [[ -f "$main_conf" ]] || continue

        if [[ -f "$snippet" ]]; then
            info "Already exists: ${snippet}"
            ((skipped++)) || true
            continue
        fi

        if [[ "$DRY_RUN" == "true" ]]; then
            info "Would create: ${snippet}"
            ((created++)) || true
        else
            echo "$SNIPPET_CONTENT" > "$snippet"
            created_files+=("$snippet")
            log "Created ${snippet}"
            ((created++)) || true
        fi
    done
done

echo ""

if [[ $created -eq 0 && $skipped -gt 0 ]]; then
    info "HEAD requests are already blocked on all domains"
    exit 0
fi

if [[ "$DRY_RUN" == "true" ]]; then
    echo ""
    echo "$SNIPPET_CONTENT"
    echo ""
    info "Would create ${created} snippet(s) (${skipped} already exist)"
    info "Would test Nginx config and reload"
    exit 0
fi

# Test Nginx config
info "Testing Nginx configuration..."
if nginx -t 2>&1; then
    echo ""
    log "Config test passed"
    systemctl reload nginx
    log "Nginx reloaded — HEAD requests blocked on ${#domain_dirs[@]} domain(s) (444 drop)"
else
    echo ""
    err "Config test FAILED — rolling back all changes"
    for f in "${created_files[@]}"; do
        rm -f "$f"
        err "Removed ${f}"
    done
    err "Nginx was NOT reloaded — your site is unaffected"
    exit 1
fi

echo ""
info "Verify with: curl -I https://your-site.com"
info "Expected: curl returns empty reply (connection dropped)"
info "To undo:  $(basename "$0") --remove"