#!/bin/bash

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

# Function to print colored text
print_color() {
    local color=$1
    local text=$2
    echo -e "${color}${text}${NC}"
}

# Function to check if script is run as root
check_root() {
    if [ "$EUID" -ne 0 ]; then
        print_color $RED "Please run this script as root"
        exit 1
    fi
}

# Function to test DNS resolution
test_dns() {
    print_color $BLUE "\nTesting DNS Resolution..."
    echo "Current DNS servers:"
    cat /etc/resolv.conf | grep nameserver
    
    read -p "Enter a domain to test DNS resolution (default: google.com): " test_domain
    test_domain=${test_domain:-google.com}
    
    if dig +short "$test_domain" > /dev/null; then
        print_color $GREEN "DNS resolution successful for $test_domain"
        echo "IP Address: $(dig +short "$test_domain")"
    else
        print_color $RED "DNS resolution failed for $test_domain"
    fi
}

# Function to test specific ports
test_ports() {
    print_color $BLUE "\nPort Testing Menu"
    echo "1. Test specific port"
    echo "2. Test common Proxmox ports (8006, 3128, 111, 5900-5999)"
    echo "3. Back to main menu"
    
    read -p "Select an option: " port_choice
    
    case $port_choice in
        1)
            read -p "Enter hostname/IP to test: " host
            read -p "Enter port number: " port
            if nc -zv $host $port 2>&1 | grep -q 'succeeded'; then
                print_color $GREEN "Port $port is open on $host"
            else
                print_color $RED "Port $port is closed on $host"
            fi
            ;;
        2)
            print_color $YELLOW "Testing common Proxmox ports..."
            local ports=(8006 3128 111 5900 5999)
            for port in "${ports[@]}"; do
                if nc -zv localhost $port 2>&1 | grep -q 'succeeded'; then
                    print_color $GREEN "Port $port is open"
                else
                    print_color $RED "Port $port is closed"
                fi
            done
            ;;
        3)
            return
            ;;
    esac
}

# Function to check network interfaces
check_interfaces() {
    print_color $BLUE "\nNetwork Interfaces Information:"
    echo "1. Show all interfaces"
    echo "2. Show specific interface details"
    echo "3. Show bridge information"
    echo "4. Back to main menu"
    
    read -p "Select an option: " int_choice
    
    case $int_choice in
        1)
            print_color $YELLOW "All Network Interfaces:"
            ip -c a
            ;;
        2)
            ip -c link show
            read -p "Enter interface name to inspect: " interface
            if ip link show $interface >/dev/null 2>&1; then
                print_color $YELLOW "Details for interface $interface:"
                ip -c addr show $interface
                ethtool $interface 2>/dev/null || print_color $RED "ethtool not available"
            else
                print_color $RED "Interface $interface not found"
            fi
            ;;
        3)
            print_color $YELLOW "Bridge Information:"
            brctl show 2>/dev/null || print_color $RED "brctl not installed. Installing..."
            if ! command -v brctl &> /dev/null; then
                apt-get update && apt-get install -y bridge-utils
                brctl show
            fi
            ;;
        4)
            return
            ;;
    esac
}

# Function to test network connectivity
test_connectivity() {
    print_color $BLUE "\nNetwork Connectivity Tests:"
    echo "1. Basic ping test"
    echo "2. MTR trace"
    echo "3. Bandwidth test"
    echo "4. Back to main menu"
    
    read -p "Select an option: " conn_choice
    
    case $conn_choice in
        1)
            read -p "Enter host to ping (default: 8.8.8.8): " ping_host
            ping_host=${ping_host:-8.8.8.8}
            ping -c 4 $ping_host
            ;;
        2)
            if ! command -v mtr &> /dev/null; then
                print_color $YELLOW "MTR not found. Installing..."
                apt-get update && apt-get install -y mtr-tiny
            fi
            read -p "Enter host to trace (default: 8.8.8.8): " trace_host
            trace_host=${trace_host:-8.8.8.8}
            mtr -r $trace_host
            ;;
        3)
            if ! command -v iperf3 &> /dev/null; then
                print_color $YELLOW "iperf3 not found. Installing..."
                apt-get update && apt-get install -y iperf3
            fi
            print_color $YELLOW "Starting iperf3 server..."
            iperf3 -s &
            sleep 2
            print_color $YELLOW "Testing local bandwidth..."
            iperf3 -c localhost
            killall iperf3
            ;;
        4)
            return
            ;;
    esac
}

# Function to check network card information
check_network_cards() {
    print_color $BLUE "\nNetwork Card Information:"
    echo "1. Show PCI network devices"
    echo "2. Show detailed network card info"
    echo "3. Show driver information"
    echo "4. Back to main menu"
    
    read -p "Select an option: " card_choice
    
    case $card_choice in
        1)
            print_color $YELLOW "PCI Network Devices:"
            lspci | grep -i ethernet
            ;;
        2)
            print_color $YELLOW "Detailed Network Card Information:"
            if ! command -v lshw &> /dev/null; then
                print_color $YELLOW "lshw not found. Installing..."
                apt-get update && apt-get install -y lshw
            fi
            lshw -class network
            ;;
        3)
            print_color $YELLOW "Network Driver Information:"
            for i in $(ls /sys/class/net/); do
                print_color $GREEN "\nDriver info for $i:"
                ethtool -i $i 2>/dev/null || echo "No driver information available"
            done
            ;;
        4)
            return
            ;;
    esac
}

# Function to display firewall status
check_firewall() {
    print_color $BLUE "\nFirewall Status:"
    echo "1. Show Proxmox firewall status"
    echo "2. Show iptables rules"
    echo "3. Back to main menu"
    
    read -p "Select an option: " fw_choice
    
    case $fw_choice in
        1)
            print_color $YELLOW "Proxmox Firewall Status:"
            pvesh get /nodes/$(hostname)/firewall/rules
            ;;
        2)
            print_color $YELLOW "IPTables Rules:"
            iptables -L -v -n
            ;;
        3)
            return
            ;;
    esac
}

# Main menu function
main_menu() {
    while true; do
        print_color $BLUE "\nProxmox Network Diagnostic Tool"
        echo "1. Test DNS Resolution"
        echo "2. Test Ports"
        echo "3. Check Network Interfaces"
        echo "4. Test Network Connectivity"
        echo "5. Check Network Cards"
        echo "6. Check Firewall Status"
        echo "7. Exit"
        
        read -p "Select an option: " choice
        
        case $choice in
            1) test_dns ;;
            2) test_ports ;;
            3) check_interfaces ;;
            4) test_connectivity ;;
            5) check_network_cards ;;
            6) check_firewall ;;
            7) 
                print_color $GREEN "Exiting..."
                exit 0
                ;;
            *)
                print_color $RED "Invalid option"
                ;;
        esac
        
        read -p "Press Enter to continue..."
    done
}

# Script entry point
check_root
main_menu