#!/bin/bash
# ==========================================
# Force IPv4 for Package Managers (APT/DNF/YUM)
# Supports: Ubuntu, Debian, AlmaLinux, Rocky, CentOS (7/8/9)
# Author: MobinHost (Aria Jahangiri Far)
# ==========================================

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

# File Paths
APT_CONF_FILE="/etc/apt/apt.conf.d/99force-ipv4"
DNF_CONF_FILE="/etc/dnf/dnf.conf"
YUM_CONF_FILE="/etc/yum.conf"

# Detect OS
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$ID
else
    echo -e "${RED}[ERROR] Cannot detect Operating System.${NC}"
    exit 1
fi

# Function to update RHEL-based configs (DNF/YUM)
update_rhel_conf() {
    local target_file=$1
    if [[ -f "$target_file" ]]; then
        if grep -q "ip_resolve=" "$target_file"; then
            sudo sed -i 's/ip_resolve=.*/ip_resolve=4/' "$target_file"
        else
            # Add it under the [main] section
            sudo sed -i '/\[main\]/a ip_resolve=4' "$target_file"
        fi
        return 0
    fi
    return 1
}

# Enable ForceIPv4
enable_ipv4() {
    echo -e "${BLUE}Configuring IPv4 priority for ${CYAN}$NAME${BLUE}...${NC}"

    # Debian/Ubuntu (APT)
    if [[ "$OS" == "ubuntu" || "$OS" == "debian" ]]; then
        echo 'Acquire::ForceIPv4 "true";' | sudo tee "$APT_CONF_FILE" > /dev/null
        echo -e "${GREEN}[SUCCESS] APT configured to force IPv4.${NC}"

    # RHEL/AlmaLinux/Rocky/CentOS (DNF or YUM)
    elif [[ "$OS" == "almalinux" || "$OS" == "rocky" || "$OS" == "centos" || "$OS" == "rhel" ]]; then
        update_rhel_conf "$DNF_CONF_FILE"
        update_rhel_conf "$YUM_CONF_FILE"
        echo -e "${GREEN}[SUCCESS] YUM/DNF configured to force IPv4 (ip_resolve=4).${NC}"
    else
        echo -e "${RED}[ERROR] Unsupported distribution: $OS${NC}"
    fi
}

# Disable ForceIPv4
disable_ipv4() {
    echo -e "${BLUE}Restoring default IP settings for ${CYAN}$NAME${BLUE}...${NC}"

    # Debian/Ubuntu (APT)
    if [[ "$OS" == "ubuntu" || "$OS" == "debian" ]]; then
        sudo rm -f "$APT_CONF_FILE"
        echo -e "${YELLOW}[REMOVED] APT IPv4 force removed.${NC}"

    # RHEL/AlmaLinux/Rocky/CentOS (DNF or YUM)
    elif [[ "$OS" == "almalinux" || "$OS" == "rocky" || "$OS" == "centos" || "$OS" == "rhel" ]]; then
        sudo sed -i '/ip_resolve=4/d' "$DNF_CONF_FILE" 2>/dev/null
        sudo sed -i '/ip_resolve=4/d' "$YUM_CONF_FILE" 2>/dev/null
        echo -e "${YELLOW}[REMOVED] YUM/DNF IPv4 force removed.${NC}"
    fi
}

# Check Status
check_status() {
    STATUS_MSG="${RED}DISABLED${NC}"
    if [[ "$OS" == "ubuntu" || "$OS" == "debian" ]]; then
        [[ -f "$APT_CONF_FILE" ]] && STATUS_MSG="${GREEN}ENABLED${NC}"
    elif [[ "$OS" == "almalinux" || "$OS" == "rocky" || "$OS" == "centos" || "$OS" == "rhel" ]]; then
        (grep -q "ip_resolve=4" "$DNF_CONF_FILE" 2>/dev/null || grep -q "ip_resolve=4" "$YUM_CONF_FILE" 2>/dev/null) && STATUS_MSG="${GREEN}ENABLED${NC}"
    fi
    echo -e "Current Status: $STATUS_MSG"
}

# Main Menu
clear
echo -e "${CYAN}======================================${NC}"
echo -e "   Network Force IPv4 Tool (APT/YUM/DNF)"
echo -e "   Operating System: ${YELLOW}$NAME${NC}"
check_status
echo -e "${CYAN}======================================${NC}"
echo "1) Enable ForceIPv4"
echo "2) Disable ForceIPv4"
echo "3) Exit"
echo -e "${CYAN}--------------------------------------${NC}"

read -p "Enter your choice [1-3]: " choice

case "$choice" in
    1) enable_ipv4 ;;
    2) disable_ipv4 ;;
    3) echo "Exiting..."; exit 0 ;;
    *) echo -e "${RED}Invalid choice!${NC}"; exit 1 ;;
esac