#!/usr/bin/env bash
# test-headers.sh — Security Headers lab (ch.6)
# ---------------------------------------------------------------------------
# Two jobs:
#   ./test-headers.sh --apply      write a strong security-headers snippet (root)
#   ./test-headers.sh [url]        probe a live URL and grade its headers (A–F)
# ---------------------------------------------------------------------------
set -uo pipefail
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
print_header(){ echo -e "\n${GREEN}=== $1 ===${NC}"; }

apply() {
  [[ ${EUID:-$(id -u)} -eq 0 ]] || { echo -e "${RED}--apply needs root${NC}"; exit 1; }
  install -d /etc/nginx/snippets
  cat > /etc/nginx/snippets/security-headers.conf <<'EOF'
add_header X-Content-Type-Options "nosniff" always;
add_header X-Frame-Options "DENY" always;
add_header Referrer-Policy "no-referrer" always;
add_header Permissions-Policy "camera=(), microphone=(), geolocation=(), interest-cohort=()" always;
add_header Cross-Origin-Opener-Policy "same-origin" always;
add_header Cross-Origin-Resource-Policy "same-origin" always;
add_header Content-Security-Policy "default-src 'self'; object-src 'none'; base-uri 'self'; frame-ancestors 'none'; upgrade-insecure-requests" always;
add_header Strict-Transport-Security "max-age=63072000; includeSubDomains; preload" always;
EOF
  echo -e "${GREEN}[+] Wrote /etc/nginx/snippets/security-headers.conf${NC}"
  echo    "    'include snippets/security-headers.conf;' in your server block, then: nginx -t && systemctl reload nginx"
  exit 0
}
[[ "${1:-}" == "--apply" ]] && apply

URL="${1:-http://localhost/}"
print_header "Grading security headers on ${URL}"
H="$(curl -sSIL --max-time 10 "$URL")" || { echo -e "${RED}could not reach ${URL}${NC}"; exit 1; }
val(){ grep -i "^$1:" <<<"$H" | head -1 | cut -d: -f2- | sed 's/^[[:space:]]*//; s/[[:space:]]*$//'; }

score=0; max=0
grade1(){ # name weight [must-contain]
  local name="$1" w="$2" need="${3:-}" v; max=$((max+w)); v="$(val "$name")"
  if [[ -z "$v" ]]; then printf '  \033[31mx %-28s missing\033[0m\n' "$name"; return; fi
  if [[ -n "$need" && "$v" != *"$need"* ]]; then
    printf '  \033[33m~ %-28s weak: %s\033[0m\n' "$name" "$v"; score=$((score + w/2)); return; fi
  printf '  \033[32m+ %-28s %s\033[0m\n' "$name" "$v"; score=$((score + w))
}
grade1 "Content-Security-Policy"   30
grade1 "Strict-Transport-Security" 20 "max-age=6"
grade1 "X-Content-Type-Options"    15 "nosniff"
grade1 "X-Frame-Options"           15
grade1 "Referrer-Policy"           10
grade1 "Permissions-Policy"        10

# penalties
if grep -qiE "^Server:.*[0-9]+\.[0-9]+" <<<"$H"; then
  printf '  \033[31mx %-28s leaks version: %s\033[0m\n' "Server" "$(val Server)"; score=$((score-5)); fi
if grep -qi "unsafe-inline" <<<"$H"; then
  printf '  \033[33m~ %-28s CSP allows unsafe-inline\033[0m\n' "Content-Security-Policy"; score=$((score-5)); fi

(( score<0 )) && score=0
pct=0; (( max>0 )) && pct=$(( score*100/max ))
g="F"; (( pct>=90 )) && g="A"; (( pct>=80 && pct<90 )) && g="B"; (( pct>=70 && pct<80 )) && g="C"; (( pct>=60 && pct<70 )) && g="D"
echo -e "\n  Score: ${pct}/100   Grade: ${GREEN}${g}${NC}"
echo    "  Apply a strong baseline:  sudo $0 --apply"