#!/bin/bash

# VM Network Setup Script
# This script sets up networking for VMs to enable SSH access

set -e

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

log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
    exit 1
}

info() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')] INFO: $1${NC}"
}

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    error "This script must be run as root"
fi

BRIDGE_NAME="br0"
BRIDGE_IP="192.168.100.1/24"
NETWORK="192.168.100.0/24"

log "Setting up VM networking..."

# Check if bridge already exists and has IP
if ip link show "$BRIDGE_NAME" &>/dev/null; then
    if ip addr show "$BRIDGE_NAME" | grep -q "192.168.100.1"; then
        info "Bridge $BRIDGE_NAME already configured with IP"
    else
        log "Adding IP address to existing bridge..."
        ip addr add "$BRIDGE_IP" dev "$BRIDGE_NAME"
    fi
else
    log "Bridge $BRIDGE_NAME not found. It will be created when VMs start."
fi

# Enable IP forwarding
log "Enabling IP forwarding..."
echo 1 > /proc/sys/net/ipv4/ip_forward

# Make IP forwarding persistent
if ! grep -q "net.ipv4.ip_forward=1" /etc/sysctl.conf; then
    echo "net.ipv4.ip_forward=1" >> /etc/sysctl.conf
    log "IP forwarding made persistent in /etc/sysctl.conf"
fi

# Set up NAT for VM network
log "Setting up NAT for VM network..."

# Remove existing rules to avoid duplicates
iptables -t nat -D POSTROUTING -s "$NETWORK" -j MASQUERADE 2>/dev/null || true
iptables -D FORWARD -i "$BRIDGE_NAME" -o "$BRIDGE_NAME" -j ACCEPT 2>/dev/null || true
iptables -D FORWARD -i "$BRIDGE_NAME" -j ACCEPT 2>/dev/null || true
iptables -D FORWARD -o "$BRIDGE_NAME" -j ACCEPT 2>/dev/null || true

# Add new rules
iptables -t nat -A POSTROUTING -s "$NETWORK" -j MASQUERADE
iptables -A FORWARD -i "$BRIDGE_NAME" -o "$BRIDGE_NAME" -j ACCEPT
iptables -A FORWARD -i "$BRIDGE_NAME" -j ACCEPT
iptables -A FORWARD -o "$BRIDGE_NAME" -j ACCEPT

log "NAT rules configured"

# Install and configure dnsmasq for DHCP
if ! command -v dnsmasq &>/dev/null; then
    log "Installing dnsmasq for DHCP..."
    apt update && apt install -y dnsmasq
fi

# Configure dnsmasq for VM network
DNSMASQ_CONF="/etc/dnsmasq.d/vm-network.conf"
log "Configuring DHCP for VM network..."

cat > "$DNSMASQ_CONF" << EOF
# VM Network DHCP Configuration
# Only bind to the bridge interface to avoid conflicts with systemd-resolved
interface=$BRIDGE_NAME
bind-interfaces

# Disable DNS server functionality (only DHCP)
port=0

# DHCP configuration
dhcp-range=192.168.100.10,192.168.100.100,12h
dhcp-option=3,192.168.100.1
dhcp-option=6,8.8.8.8,8.8.4.4

# Disable reading /etc/hosts and /etc/resolv.conf
no-hosts
no-resolv
EOF

# Restart dnsmasq
systemctl restart dnsmasq
systemctl enable dnsmasq

log "DHCP server configured and started"

# Create a script to show VM IPs
cat > "/usr/local/bin/vm-ips" << 'EOF'
#!/bin/bash
echo "VM DHCP Leases:"
echo "==============="
if [ -f /var/lib/dhcp/dhcpd.leases ]; then
    awk '/lease/ { ip = $2 } /client-hostname/ { hostname = $2; gsub(/[";]/, "", hostname) } /binding state active/ { print ip " - " hostname }' /var/lib/dhcp/dhcpd.leases
elif [ -f /var/lib/dhcpcd5/dhcpcd.leases ]; then
    cat /var/lib/dhcpcd5/dhcpcd.leases
else
    echo "DHCP lease file not found. Checking dnsmasq leases..."
    if [ -f /var/lib/dhcp/dhcpd.leases ]; then
        cat /var/lib/dhcp/dhcpd.leases
    else
        echo "No lease information available"
        echo "Try: arp -a | grep 192.168.100"
    fi
fi
EOF

chmod +x /usr/local/bin/vm-ips

log "Network setup completed!"
echo ""
info "Network Configuration Summary:"
info "- Bridge: $BRIDGE_NAME with IP $BRIDGE_IP"
info "- DHCP range: 192.168.100.10 - 192.168.100.100"
info "- DNS servers: 8.8.8.8, 8.8.4.4"
info "- NAT configured for internet access"
echo ""
info "To see VM IP addresses: vm-ips"
info "To check bridge status: ip addr show $BRIDGE_NAME"
info "To see DHCP leases: cat /var/lib/dhcp/dhcpd.leases"
echo ""
warn "Note: VMs need to be restarted to get DHCP IP addresses"