#!/bin/bash
# =============================================================================
# SCRIPT: setup-firewall.sh
# PURPOSE: Configure UFW firewall to block all incoming connections on Ubuntu
# LAST UPDATED: June 30, 2025
# KEY FEATURE: Blocks everything including ICMP (ping) from external sources
# COMPATIBILITY: Ubuntu 24.04.2 LTS
# =============================================================================

set -euo pipefail

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

# Check if running as root
if [[ $EUID -ne 0 ]]; then
   echo -e "${RED}This script must be run as root (use sudo)${NC}"
   exit 1
fi

echo -e "${GREEN}Starting firewall configuration...${NC}"

# Install UFW if not already installed
if ! command -v ufw &> /dev/null; then
    echo -e "${YELLOW}UFW not found. Installing...${NC}"
    apt-get update
    apt-get install -y ufw
fi

# Reset UFW to defaults
echo -e "${YELLOW}Resetting UFW to defaults...${NC}"
ufw --force reset

# Set default policies
echo -e "${YELLOW}Setting default policies...${NC}"
ufw default deny incoming
ufw default deny outgoing
ufw default deny routed

# Allow outgoing connections (so the system can still make requests)
echo -e "${YELLOW}Allowing outgoing connections...${NC}"
ufw default allow outgoing

# Disable IPv6 if not needed (optional but recommended for security)
echo -e "${YELLOW}Disabling IPv6 support in UFW...${NC}"
sed -i 's/IPV6=yes/IPV6=no/g' /etc/default/ufw

# Block ICMP (ping) explicitly
echo -e "${YELLOW}Blocking ICMP echo requests...${NC}"
# Add before.rules to block ICMP
cat > /etc/ufw/before.rules.tmp << 'EOF'
# Block ICMP echo-request
-A ufw-before-input -p icmp --icmp-type echo-request -j DROP
EOF

# Prepend our ICMP rule to the existing before.rules
if ! grep -q "icmp-type echo-request" /etc/ufw/before.rules; then
    # Find the line with "# ok icmp codes" and insert our rule before it
    sed -i '/# ok icmp codes/i # Block ICMP echo-request\n-A ufw-before-input -p icmp --icmp-type echo-request -j DROP\n' /etc/ufw/before.rules
fi

# Clean up temp file
rm -f /etc/ufw/before.rules.tmp

# Enable logging
echo -e "${YELLOW}Enabling firewall logging...${NC}"
ufw logging medium

# Enable UFW
echo -e "${YELLOW}Enabling UFW...${NC}"
ufw --force enable

# Show status
echo -e "${GREEN}Firewall configuration complete!${NC}"
echo -e "${GREEN}Current firewall status:${NC}"
ufw status verbose

echo -e "\n${YELLOW}Summary of configuration:${NC}"
echo "- All incoming connections: DENIED"
echo "- All outgoing connections: ALLOWED"
echo "- ICMP (ping): BLOCKED"
echo "- IPv6: DISABLED"
echo "- Logging: MEDIUM"

echo -e "\n${YELLOW}Important notes:${NC}"
echo "- SSH access is now blocked! Make sure you have console access."
echo "- To allow SSH later: sudo ufw allow ssh"
echo "- To check firewall status: sudo ufw status"
echo "- To disable firewall: sudo ufw disable"
echo "- Firewall logs: /var/log/ufw.log"

echo -e "\n${GREEN}Firewall is now active and blocking all incoming connections!${NC}"