#!/bin/sh

DATA_DIR="data"

echo "========================================"
echo "  iptables-autoconf - USTC Blacklist   "
echo "========================================"
echo ""

echo "[1/6] Initializing..."
mkdir -p "$DATA_DIR"

echo "      [OK] Data directory: $DATA_DIR"

echo ""
echo "[2/6] Initializing ipset..."
ipset create ustc_blacklist_v4 hash:ip --exist
ipset create ustc_blacklist_v4_net hash:net --exist
ipset create ustc_blacklist_v6 hash:ip --exist
ipset create ustc_blacklist_v6_net hash:net --exist

echo "      Flushing existing ipset entries..."
ipset flush ustc_blacklist_v4
ipset flush ustc_blacklist_v4_net
ipset flush ustc_blacklist_v6
ipset flush ustc_blacklist_v6_net
echo "      [OK] ipset initialized"

echo ""
echo "[3/6] Downloading blacklist from USTC..."
[ -f "$DATA_DIR/blacklist_ustc.txt" ] && rm "$DATA_DIR/blacklist_ustc.txt"
wget -q http://blackip.ustc.edu.cn/list.php?txt -O "$DATA_DIR/blacklist_ustc.txt"

if [ $? -ne 0 ]; then
    echo "      [ERROR] Failed to download blacklist file"
    exit 1
fi

file_size=$(wc -c < "$DATA_DIR/blacklist_ustc.txt")
echo "      [OK] Downloaded ($((file_size/1024)) KB)"

echo ""
echo "[4/6] Processing blacklist data..."
python3 blacklist.py
echo "      [OK] Data processed"

echo ""
echo "[5/6] Adding entries to ipset..."

echo "      - IPv4 single addresses..."
count_v4=0
for addr in $(cat "$DATA_DIR/ipv4_list.txt" 2>/dev/null); do
    ipset add ustc_blacklist_v4 $addr 2>/dev/null && count_v4=$((count_v4 + 1))
done
echo "        Added $count_v4 entries"

echo "      - IPv4 CIDR networks..."
count_v4_net=0
for addr in $(cat "$DATA_DIR/ipv4_net_list.txt" 2>/dev/null); do
    ipset add ustc_blacklist_v4_net $addr 2>/dev/null && count_v4_net=$((count_v4_net + 1))
done
echo "        Added $count_v4_net entries"

echo "      - IPv6 single addresses..."
count_v6=0
for addr in $(cat "$DATA_DIR/ipv6_list.txt" 2>/dev/null); do
    ipset add ustc_blacklist_v6 $addr 2>/dev/null && count_v6=$((count_v6 + 1))
done
echo "        Added $count_v6 entries"

echo "      - IPv6 CIDR networks..."
count_v6_net=0
for addr in $(cat "$DATA_DIR/ipv6_net_list.txt" 2>/dev/null); do
    ipset add ustc_blacklist_v6_net $addr 2>/dev/null && count_v6_net=$((count_v6_net + 1))
done
echo "        Added $count_v6_net entries"

echo "      [OK] All entries added to ipset"

echo ""
echo "[6/6] Configuring iptables rules..."

add_rule() {
    local set_name="$1"
    if iptables -C INPUT -m set --match-set "$set_name" src -j DROP 2>/dev/null; then
        echo "      - $set_name: already exists, skipped"
    else
        iptables -A INPUT -m set --match-set "$set_name" src -j DROP
        echo "      - $set_name: added"
    fi
}

add_rule "ustc_blacklist_v4"
add_rule "ustc_blacklist_v4_net"
add_rule "ustc_blacklist_v6"
add_rule "ustc_blacklist_v6_net"
echo "      [OK] iptables configured"

echo ""
echo "[7/7] Summary"
echo "----------------------------------------"
echo "  Blocked entries:"
echo "    IPv4 single:    $count_v4"
echo "    IPv4 CIDR:      $count_v4_net"
echo "    IPv6 single:    $count_v6"
echo "    IPv6 CIDR:      $count_v6_net"
echo "    ----------------------------"
echo "    Total:          $((count_v4 + count_v4_net + count_v6 + count_v6_net))"
echo ""
echo "  Current iptables rules:"
iptables -t filter -L INPUT --line-numbers -v | grep -E "ustc_blacklist|num=1" | head -10
echo "========================================"
echo "  Update completed successfully!"
echo "========================================"