#!/bin/bash
set -e

IPSET_NAME="china"
CIDR_FILE="/tmp/cn.cidr"

SOURCES=(
    "https://www.ipdeny.com/ipblocks/data/countries/cn.zone"
    "https://raw.githubusercontent.com/mayaxcn/china-ip-list/master/chnroute.txt"
)

echo "[*] Killing Steam..."
pkill -9 -f steam || true

echo "[*] Downloading CN IP ranges..."
for url in "${SOURCES[@]}"; do
    echo "    Trying $url ..."
    wget -q "$url" -O "$CIDR_FILE"
    if [[ $(wc -l < "$CIDR_FILE") -gt 10 ]]; then
        echo "    [+] Success: $(wc -l < "$CIDR_FILE") entries"
        break
    fi
done

if [[ $(wc -l < "$CIDR_FILE") -lt 10 ]]; then
    echo "[!] All sources failed. Aborting."
    exit 1
fi

echo "[*] Creating ipset..."
ipset create "$IPSET_NAME" hash:net 2>/dev/null || ipset flush "$IPSET_NAME"

echo "[*] Populating ipset..."
while IFS= read -r cidr; do
    [[ -z "$cidr" || "$cidr" == \#* ]] && continue
    ipset add "$IPSET_NAME" "$cidr" || true
done < "$CIDR_FILE"

echo "[*] Applying iptables rules..."
iptables -A INPUT -m set --match-set "$IPSET_NAME" src -j DROP
iptables -A OUTPUT -m set --match-set "$IPSET_NAME" dst -j DROP

echo "[*] Applying raw table rules (block UDP hole punch)..."
iptables -t raw -A PREROUTING -m set --match-set "$IPSET_NAME" src -j DROP
iptables -t raw -A OUTPUT -m set --match-set "$IPSET_NAME" dst -j DROP

CIDR6_FILE="/tmp/cn6.cidr"
SOURCES6=(
    "https://raw.githubusercontent.com/herrbischoff/country-ip-blocks/master/ipv6/cn.cidr"
    "https://www.ipdeny.com/ipv6/ipaddresses/blocks/cn.cidr"
)

echo "[*] Downloading CN IPv6 ranges..."
for url in "${SOURCES6[@]}"; do
    echo "    Trying $url ..."
    wget -q "$url" -O "$CIDR6_FILE"
    if [[ $(wc -l < "$CIDR6_FILE") -gt 10 ]]; then
        echo "    [+] Success: $(wc -l < "$CIDR6_FILE") entries"
        break
    fi
done

if [[ $(wc -l < "$CIDR6_FILE") -lt 10 ]]; then
    echo "[!] All IPv6 sources failed. Skipping IPv6 blocking."
else
    ipset create china6 hash:net family inet6 2>/dev/null || ipset flush china6
    while IFS= read -r cidr; do
        [[ -z "$cidr" || "$cidr" == \#* ]] && continue
        ipset add china6 "$cidr" || true
    done < "$CIDR6_FILE"

    ip6tables -A INPUT -m set --match-set china6 src -j DROP
    ip6tables -A OUTPUT -m set --match-set china6 dst -j DROP
    ip6tables -t raw -A PREROUTING -m set --match-set china6 src -j DROP
    ip6tables -t raw -A OUTPUT -m set --match-set china6 dst -j DROP
fi

echo "[*] Saving rules..."
mkdir -p /etc/iptables
ipset save > /etc/ipset.conf
iptables-save > /etc/iptables/iptables.rules
ip6tables-save > /etc/iptables/ip6tables.rules
systemctl enable --now iptables
systemctl enable --now ip6tables 2>/dev/null || true

echo "[+] Done! Chinese IPs are now blocked."