#!/bin/bash

# Security Headers Testing Script for Hastebin
# This script tests various security header configurations by running curl commands
# to verify headers are correctly set and the application works properly.

# Colors for output
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
PORT=7777
HOST=localhost
SERVER_START_WAIT=5 # seconds
KILL_WAIT=2 # seconds

# Utility functions
print_header() {
  echo -e "\n${BLUE}$1${NC}"
  echo -e "${BLUE}$(printf '=%.0s' $(seq 1 ${#1}))${NC}"
}

print_success() {
  echo -e "${GREEN}✅ $1${NC}"
}

print_error() {
  echo -e "${RED}❌ $1${NC}"
}

print_info() {
  echo -e "${YELLOW}ℹ️ $1${NC}"
}

# Kill any running server instance
kill_server() {
  pkill -f "node test-local.js" >/dev/null 2>&1
  sleep $KILL_WAIT
}

# Start server with specified environment variables
start_server() {
  print_info "Starting server with: $1"
  
  # Use nohup to ensure the process continues running even if the script is interrupted
  eval "HASTEBIN_STORAGE_TYPE=file $1 node test-local.js > /tmp/hastebin-test.log 2>&1 &"
  
  # Store the PID for later cleanup
  SERVER_PID=$!
  
  # Wait for server to start and log the process ID
  print_info "Started server process with PID: $SERVER_PID, waiting ${SERVER_START_WAIT}s..."
  sleep $SERVER_START_WAIT
  
  # Check if the server is actually running
  if ! ps -p $SERVER_PID > /dev/null; then
    print_error "Server failed to start! Check logs at /tmp/hastebin-test.log"
    cat /tmp/hastebin-test.log | head -n 20
    return 1
  fi
  
  return 0
}

# Check if a header exists and matches expected value
check_header() {
  local header="$1"
  local expected="$2"
  local response="$3"
  
  # Extract the header value from response
  local value=$(echo "$response" | grep -i "^$header:" | sed "s/^$header: //i" | tr -d '\r')
  
  if [ -z "$value" ]; then
    if [ "$expected" == "ABSENT" ]; then
      return 0
    else
      print_error "Header '$header' is missing"
      return 1
    fi
  else
    if [ "$expected" == "ABSENT" ]; then
      print_error "Header '$header' should be absent but found: $value"
      return 1
    elif [ "$expected" == "ANY" ]; then
      return 0
    elif [[ "$value" == *"$expected"* ]]; then
      return 0
    else
      print_error "Header '$header' expected to contain '$expected' but got '$value'"
      return 1
    fi
  fi
}

# Test functionality by creating and retrieving a document
test_functionality() {
  print_info "Testing document creation and retrieval..."
  
  # Try multiple times with backoff
  for attempt in 1 2 3; do
    # Create a document
    local create_response=$(curl -s -X POST http://$HOST:$PORT/documents -d "Security Test Document")
    echo "Create response: $create_response"
    
    # Extract the key using a more reliable method
    local key=$(echo $create_response | sed -n 's/.*"key":"\([^"]*\)".*/\1/p')
    
    if [ -n "$key" ]; then
      print_info "Created document with key: $key"
      
      # Retrieve the document
      local get_response=$(curl -s http://$HOST:$PORT/raw/$key)
      
      if [ "$get_response" == "Security Test Document" ]; then
        print_success "Document creation and retrieval successful"
        return 0
      else
        print_error "Document retrieval failed - expected 'Security Test Document' but got '$get_response'"
      fi
    else
      print_error "Failed to extract key from response: $create_response"
    fi
    
    # If we reach here, something failed - wait and retry
    sleep_time=$((attempt * 2))
    print_info "Attempt $attempt failed, waiting ${sleep_time}s before retry..."
    sleep $sleep_time
  done
  
  print_error "Failed to create or retrieve document after 3 attempts"
  return 1
}

# Run a single test
run_test() {
  local test_name="$1"
  local env_vars="$2"
  local headers_to_check="$3"
  
  print_header "Running test: $test_name"
  
  # Kill any existing server and start a new one with the specified env
  kill_server
  if ! start_server "$env_vars"; then
    print_error "Could not start server for test '$test_name'"
    return 1
  fi
  
  # Get headers - retry a few times if needed
  local response=""
  for attempt in 1 2 3; do
    response=$(curl -I -s http://$HOST:$PORT/)
    if [[ "$response" == *"HTTP/1."* ]]; then
      break
    fi
    sleep 2
  done
  
  if [[ "$response" != *"HTTP/1."* ]]; then
    print_error "Could not get response from server"
    return 1
  fi
  
  # Check each header
  local failed=0
  
  # Parse the headers to check and their expected values
  local IFS=','
  read -ra HEADER_CHECKS <<< "$headers_to_check"
  for check in "${HEADER_CHECKS[@]}"; do
    local header=$(echo $check | cut -d: -f1)
    local expected=$(echo $check | cut -d: -f2)
    
    if ! check_header "$header" "$expected" "$response"; then
      failed=1
    else
      print_success "Header '$header' check passed"
    fi
  done
  
  # Test functionality
  if ! test_functionality; then
    failed=1
  fi
  
  if [ $failed -eq 0 ]; then
    print_success "Test '$test_name' passed"
    return 0
  else
    print_error "Test '$test_name' failed"
    return 1
  fi
}

# Run tests
run_tests() {
  print_header "🔒 Hastebin Security Headers Test Suite 🔒"
  
  # Parse command line arguments
  local test_filter=""
  if [[ "$1" == --test=* ]]; then
    test_filter="${1#--test=}"
  fi
  
  local passed=0
  local failed=0
  
  # Run selected tests
  
  # Filter by name if specified
  if [[ -z "$test_filter" || "$test_filter" == *"basic"* ]]; then
    if run_test "Basic Security Headers" "NODE_ENV=production" "content-security-policy:ANY,x-content-type-options:nosniff,x-frame-options:DENY,x-xss-protection:1; mode=block,referrer-policy:strict-origin-when-cross-origin,permissions-policy:ANY"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"csp"* ]]; then
    if run_test "Content Security Policy" "NODE_ENV=production HASTEBIN_ENABLE_CSP=true" "content-security-policy:script-src"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"noCsp"* ]]; then
    if run_test "Disabled CSP" "NODE_ENV=production HASTEBIN_ENABLE_CSP=false" "content-security-policy:ABSENT,x-content-type-options:nosniff,x-frame-options:DENY"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"cors"* ]]; then
    if run_test "Cross-Origin Isolation" "NODE_ENV=production HASTEBIN_ENABLE_CROSS_ORIGIN_ISOLATION=true" "cross-origin-embedder-policy:require-corp,cross-origin-resource-policy:same-origin,cross-origin-opener-policy:same-origin"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"hsts"* ]]; then
    if run_test "HTTP Strict Transport Security" "NODE_ENV=production HASTEBIN_ENABLE_HSTS=true" "strict-transport-security:max-age"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"devMode"* ]]; then
    if run_test "Development Mode" "NODE_ENV=development" "content-security-policy:ANY,x-content-type-options:nosniff"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"devBypass"* ]]; then
    if run_test "Development Mode with CSP Bypass" "NODE_ENV=development HASTEBIN_BYPASS_CSP_IN_DEV=true" "content-security-policy:unsafe-inline"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  if [[ -z "$test_filter" || "$test_filter" == *"combined"* ]]; then
    if run_test "Combined Security Settings" "NODE_ENV=production HASTEBIN_ENABLE_CSP=false HASTEBIN_ENABLE_CROSS_ORIGIN_ISOLATION=true HASTEBIN_ENABLE_HSTS=true" "content-security-policy:ABSENT,x-content-type-options:nosniff,x-frame-options:DENY,cross-origin-embedder-policy:require-corp,strict-transport-security:max-age"; then
      passed=$((passed+1))
    else
      failed=$((failed+1))
    fi
  fi
  
  # Cleanup any remaining server process
  kill_server
  
  # Print summary
  print_header "📊 Test Results:"
  print_success "$passed tests passed"
  if [ $failed -gt 0 ]; then
    print_error "$failed tests failed"
    return 1
  fi
  
  return 0
}

# Show help if requested
if [ $# -gt 0 ] && [ "$1" == "--help" ]; then
  echo "Usage: $0 [--test=test1,test2,...]"
  echo "Available tests: basic, csp, noCsp, cors, hsts, devMode, devBypass, combined"
  exit 0
fi

# Run tests with any arguments passed
run_tests "$@"
result=$?

# Exit with status
exit $result