#!/bin/bash
#
# Nuke Telemetry Firewall Blocker
# Blocks Foundry telemetry at the firewall level using iptables or nftables
#
# Usage:
#   sudo bash firewall_block_nuke.sh          # Apply blocks
#   sudo bash firewall_block_nuke.sh --restore # Remove blocks
#   sudo bash firewall_block_nuke.sh --status  # Check current rules
#

set -e

# Known Foundry telemetry IPs (from packet capture analysis)
LEARN_FOUNDRY_IP="52.50.232.31"      # learn.foundry.com (AWS Ireland)
HONEYCOMB_IP="52.205.16.9"            # api.honeycomb.io (AWS Virginia)

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

# Logging
LOG_FILE="/var/log/nuke_firewall_block.log"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | sudo tee -a "$LOG_FILE" >/dev/null
    echo -e "$1"
}

check_root() {
    if [ "$EUID" -ne 0 ]; then
        echo -e "${RED}Error: This script must be run as root${NC}"
        echo "Usage: sudo bash $0"
        exit 1
    fi
}

detect_firewall() {
    if command -v nft &>/dev/null && nft list tables 2>/dev/null | grep -q "inet filter"; then
        echo "nftables"
    elif command -v iptables &>/dev/null; then
        echo "iptables"
    else
        echo "none"
    fi
}

apply_iptables_blocks() {
    log "${GREEN}Applying iptables blocks...${NC}"

    # Check if rules already exist
    if iptables -C OUTPUT -d "$LEARN_FOUNDRY_IP" -j REJECT --reject-with icmp-host-unreachable 2>/dev/null; then
        log "${YELLOW}iptables rules already exist, skipping...${NC}"
    else
        # Block learn.foundry.com
        iptables -A OUTPUT -d "$LEARN_FOUNDRY_IP" -j REJECT --reject-with icmp-host-unreachable
        log "Blocked $LEARN_FOUNDRY_IP (learn.foundry.com)"

        # Block api.honeycomb.io
        iptables -A OUTPUT -d "$HONEYCOMB_IP" -j REJECT --reject-with icmp-host-unreachable
        log "Blocked $HONEYCOMB_IP (api.honeycomb.io)"

        # Block sentry.foundry.com by domain (using string match)
        iptables -A OUTPUT -p tcp --dport 443 -m string --string "sentry.foundry.com" --algo bm -j REJECT --reject-with tcp-reset
        log "Blocked sentry.foundry.com (HTTPS string match)"
    fi

    # Make persistent
    if [ -d "/etc/iptables" ]; then
        iptables-save > /etc/iptables/iptables.rules
        log "Rules saved to /etc/iptables/iptables.rules"

        # Enable iptables service if available
        if systemctl is-enabled iptables.service &>/dev/null; then
            systemctl enable iptables.service
        fi
    else
        log "${YELLOW}Warning: /etc/iptables/ directory not found. Rules may not persist across reboots.${NC}"
        log "${YELLOW}Consider installing iptables-persistent or enabling iptables.service${NC}"
    fi
}

apply_nftables_blocks() {
    log "${GREEN}Applying nftables blocks...${NC}"

    # Create table and chain if they don't exist
    nft add table inet filter 2>/dev/null || true
    nft add chain inet filter output { type filter hook output priority 0 \; } 2>/dev/null || true

    # Check if rules already exist
    if nft list chain inet filter output 2>/dev/null | grep -q "$LEARN_FOUNDRY_IP"; then
        log "${YELLOW}nftables rules already exist, skipping...${NC}"
    else
        # Block learn.foundry.com
        nft add rule inet filter output ip daddr "$LEARN_FOUNDRY_IP" counter reject
        log "Blocked $LEARN_FOUNDRY_IP (learn.foundry.com)"

        # Block api.honeycomb.io
        nft add rule inet filter output ip daddr "$HONEYCOMB_IP" counter reject
        log "Blocked $HONEYCOMB_IP (api.honeycomb.io)"

        # Note: nftables doesn't have direct string matching like iptables
        # For sentry.foundry.com, rely on DNS/hosts file blocking
        log "${YELLOW}Note: sentry.foundry.com blocking requires DNS/hosts file (nftables limitation)${NC}"
    fi

    # Make persistent
    nft list ruleset > /etc/nftables.conf
    log "Rules saved to /etc/nftables.conf"

    # Enable nftables service
    if systemctl is-enabled nftables.service &>/dev/null; then
        systemctl enable nftables.service
    fi
}

remove_iptables_blocks() {
    log "${YELLOW}Removing iptables blocks...${NC}"

    # Remove specific rules
    iptables -D OUTPUT -d "$LEARN_FOUNDRY_IP" -j REJECT --reject-with icmp-host-unreachable 2>/dev/null || true
    iptables -D OUTPUT -d "$HONEYCOMB_IP" -j REJECT --reject-with icmp-host-unreachable 2>/dev/null || true
    iptables -D OUTPUT -p tcp --dport 443 -m string --string "sentry.foundry.com" --algo bm -j REJECT --reject-with tcp-reset 2>/dev/null || true

    log "iptables blocks removed"

    # Save updated rules
    if [ -d "/etc/iptables" ]; then
        iptables-save > /etc/iptables/iptables.rules
        log "Updated rules saved"
    fi
}

remove_nftables_blocks() {
    log "${YELLOW}Removing nftables blocks...${NC}"

    # Get rule handles and delete
    nft -a list chain inet filter output 2>/dev/null | grep "$LEARN_FOUNDRY_IP" | awk '{print $NF}' | while read handle; do
        nft delete rule inet filter output handle "$handle" 2>/dev/null || true
    done

    nft -a list chain inet filter output 2>/dev/null | grep "$HONEYCOMB_IP" | awk '{print $NF}' | while read handle; do
        nft delete rule inet filter output handle "$handle" 2>/dev/null || true
    done

    log "nftables blocks removed"

    # Save updated rules
    nft list ruleset > /etc/nftables.conf
    log "Updated rules saved"
}

show_status() {
    FIREWALL=$(detect_firewall)

    log "${GREEN}=== Firewall Status ===${NC}"
    log "Active firewall: $FIREWALL"
    echo ""

    if [ "$FIREWALL" = "iptables" ]; then
        log "iptables OUTPUT chain (Nuke-related rules):"
        iptables -L OUTPUT -v -n | grep -E "$LEARN_FOUNDRY_IP|$HONEYCOMB_IP|sentry.foundry.com" || echo "No Nuke blocking rules found"
    elif [ "$FIREWALL" = "nftables" ]; then
        log "nftables rules (Nuke-related):"
        nft list chain inet filter output 2>/dev/null | grep -E "$LEARN_FOUNDRY_IP|$HONEYCOMB_IP" || echo "No Nuke blocking rules found"
    else
        log "${RED}No firewall detected${NC}"
    fi

    echo ""
    log "Testing connectivity to blocked IPs:"
    timeout 2 ping -c 1 "$LEARN_FOUNDRY_IP" &>/dev/null && echo "  $LEARN_FOUNDRY_IP: REACHABLE (NOT BLOCKED)" || echo "  $LEARN_FOUNDRY_IP: BLOCKED ✓"
    timeout 2 ping -c 1 "$HONEYCOMB_IP" &>/dev/null && echo "  $HONEYCOMB_IP: REACHABLE (NOT BLOCKED)" || echo "  $HONEYCOMB_IP: BLOCKED ✓"
}

main() {
    check_root

    FIREWALL=$(detect_firewall)

    if [ "$1" = "--restore" ] || [ "$1" = "--remove" ]; then
        log "${YELLOW}=== Removing Nuke Telemetry Blocks ===${NC}"

        if [ "$FIREWALL" = "iptables" ]; then
            remove_iptables_blocks
        elif [ "$FIREWALL" = "nftables" ]; then
            remove_nftables_blocks
        else
            log "${RED}No firewall detected, nothing to remove${NC}"
            exit 1
        fi

        log "${GREEN}Blocks removed successfully${NC}"

    elif [ "$1" = "--status" ]; then
        show_status

    else
        log "${GREEN}=== Nuke Telemetry Firewall Blocker ===${NC}"
        log "This will block Foundry telemetry at the firewall level"
        log "Known endpoints: learn.foundry.com, api.honeycomb.io, sentry.foundry.com"
        echo ""

        if [ "$FIREWALL" = "none" ]; then
            log "${RED}Error: No firewall detected${NC}"
            log "Install iptables or nftables first:"
            log "  Arch: sudo pacman -S iptables-nft"
            log "  Ubuntu: sudo apt install iptables"
            exit 1
        fi

        log "Detected firewall: $FIREWALL"
        echo ""

        read -p "Continue? (y/n): " confirm
        if [ "$confirm" != "y" ]; then
            log "Aborted"
            exit 0
        fi

        if [ "$FIREWALL" = "iptables" ]; then
            apply_iptables_blocks
        elif [ "$FIREWALL" = "nftables" ]; then
            apply_nftables_blocks
        fi

        log ""
        log "${GREEN}✓ Firewall blocks applied successfully${NC}"
        log ""
        log "Verify with: sudo bash $0 --status"
        log "Remove with: sudo bash $0 --restore"
        log ""
        log "Log file: $LOG_FILE"
    fi
}

main "$@"