#!/bin/bash

# Wireguard to Docker Bridge Setup Script
# This script sets up a custom bridge for Docker that is accessible via Wireguard VPN

# Exit on any error
set -e

# Configuration values - you can adjust these
BRIDGE_NAME="br-vpn-docker"
BRIDGE_SUBNET="172.20.0.0/16"
BRIDGE_IP="172.20.0.1"
DOCKER_NETWORK_NAME="private-net"
WG_INTERFACE="wg0"
EXTERNAL_INTERFACE="eth0"  # Your server's internet-facing interface

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

# Function to print status messages
status() {
    echo -e "${GREEN}[+]${NC} $1"
}

warning() {
    echo -e "${YELLOW}[!]${NC} $1"
}

error() {
    echo -e "${RED}[!]${NC} $1"
    exit 1
}

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

# Step 1: Install dependencies
status "Installing required packages..."
apt-get update
apt-get install -y wireguard netplan.io docker.io iptables bridge-utils

# Step 2: Check if Wireguard is already configured
if [ ! -f "/etc/wireguard/${WG_INTERFACE}.conf" ]; then
    warning "Wireguard configuration not found. You'll need to configure Wireguard separately."
    warning "See: https://www.wireguard.com/quickstart/"
fi

# Step 3: Remove existing Docker network if it exists
status "Checking for existing Docker network..."
if docker network inspect "${DOCKER_NETWORK_NAME}" &>/dev/null; then
    status "Removing existing Docker network '${DOCKER_NETWORK_NAME}'..."
    
    # Check for running containers using this network
    CONTAINERS=$(docker ps --filter network=${DOCKER_NETWORK_NAME} -q)
    if [ -n "$CONTAINERS" ]; then
        warning "Containers are using the network. Stopping them..."
        docker ps --filter network=${DOCKER_NETWORK_NAME} -q | xargs -r docker stop
    fi
    
    # Remove the network
    docker network rm "${DOCKER_NETWORK_NAME}"
fi

# Step 4: Configure the bridge with Netplan
status "Creating Netplan configuration for bridge '${BRIDGE_NAME}'..."
cat > /etc/netplan/02-${BRIDGE_NAME}.yaml <<EOF
network:
  version: 2
  bridges:
    ${BRIDGE_NAME}:
      addresses: [${BRIDGE_IP}/16]
      dhcp4: no
      dhcp6: no
EOF

# Apply Netplan configuration
status "Applying Netplan configuration..."
netplan apply

# Verify bridge is created
if ! ip link show "${BRIDGE_NAME}" &>/dev/null; then
    error "Failed to create bridge interface"
fi

status "Bridge '${BRIDGE_NAME}' created successfully"

# Step 5: Create Docker network using the bridge
status "Creating Docker network '${DOCKER_NETWORK_NAME}' using bridge '${BRIDGE_NAME}'..."
docker network create --driver bridge \
  --opt "com.docker.network.bridge.name"="${BRIDGE_NAME}" \
  --subnet "${BRIDGE_SUBNET}" \
  --gateway "${BRIDGE_IP}" \
  "${DOCKER_NETWORK_NAME}" || error "Failed to create Docker network"

status "Docker network created successfully"

# Step 6: Create Wireguard iptables scripts
status "Creating Wireguard iptables configuration scripts..."

# Create the up script
cat > /etc/wireguard/${WG_INTERFACE}-up.sh <<EOF
#!/bin/bash

# Enable IP forwarding
sysctl -w net.ipv4.ip_forward=1

# Allow established and related connections
iptables -A FORWARD -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT

# Allow forwarding between Wireguard and Docker bridge
iptables -I FORWARD -i ${WG_INTERFACE} -o ${BRIDGE_NAME} -j ACCEPT
iptables -I FORWARD -i ${BRIDGE_NAME} -o ${WG_INTERFACE} -j ACCEPT

# Add rules to Docker's user chain - this is the key improvement!
iptables -I DOCKER-USER -i ${WG_INTERFACE} -o ${BRIDGE_NAME} -j ACCEPT
iptables -I DOCKER-USER -i ${BRIDGE_NAME} -o ${WG_INTERFACE} -j ACCEPT

# Add masquerading (NAT) for outgoing connections
iptables -t nat -A POSTROUTING -o ${BRIDGE_NAME} -j MASQUERADE
iptables -t nat -A POSTROUTING -o ${EXTERNAL_INTERFACE} -j MASQUERADE

echo "Wireguard to Docker bridge rules applied successfully"
EOF

# Create the down script
cat > /etc/wireguard/${WG_INTERFACE}-down.sh <<EOF
#!/bin/bash

# Allow established and related connections
iptables -D FORWARD -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT 2>/dev/null || true

# Remove forwarding rules
iptables -D FORWARD -i ${WG_INTERFACE} -o ${BRIDGE_NAME} -j ACCEPT 2>/dev/null || true
iptables -D FORWARD -i ${BRIDGE_NAME} -o ${WG_INTERFACE} -j ACCEPT 2>/dev/null || true

# Remove Docker-User chain rules
iptables -D DOCKER-USER -i ${WG_INTERFACE} -o ${BRIDGE_NAME} -j ACCEPT 2>/dev/null || true
iptables -D DOCKER-USER -i ${BRIDGE_NAME} -o ${WG_INTERFACE} -j ACCEPT 2>/dev/null || true

# Remove NAT rules
iptables -t nat -D POSTROUTING -o ${BRIDGE_NAME} -j MASQUERADE 2>/dev/null || true
iptables -t nat -D POSTROUTING -o ${EXTERNAL_INTERFACE} -j MASQUERADE 2>/dev/null || true

echo "Wireguard to Docker bridge rules removed successfully"
EOF

# Make scripts executable
chmod +x /etc/wireguard/${WG_INTERFACE}-up.sh
chmod +x /etc/wireguard/${WG_INTERFACE}-down.sh

status "Wireguard iptables scripts created"

# Step 7: Configure Wireguard to use these scripts
if [ -f "/etc/wireguard/${WG_INTERFACE}.conf" ]; then
    status "Configuring Wireguard to use the iptables scripts..."
    
    # Check if PostUp and PostDown are already configured
    if grep -q "PostUp" "/etc/wireguard/${WG_INTERFACE}.conf"; then
        warning "PostUp/PostDown directives already exist in Wireguard config"
        warning "Please manually add these lines to your Wireguard config:"
        echo "PostUp = /etc/wireguard/${WG_INTERFACE}-up.sh"
        echo "PostDown = /etc/wireguard/${WG_INTERFACE}-down.sh"
    else
        # Add the PostUp and PostDown commands
        cat >> "/etc/wireguard/${WG_INTERFACE}.conf" <<EOF

# Added by setup script
PostUp = /etc/wireguard/${WG_INTERFACE}-up.sh
PostDown = /etc/wireguard/${WG_INTERFACE}-down.sh
EOF
        status "PostUp/PostDown directives added to Wireguard config"
    fi
else
    warning "Wireguard configuration file not found"
    warning "When you create your Wireguard config, add these lines:"
    echo "PostUp = /etc/wireguard/${WG_INTERFACE}-up.sh"
    echo "PostDown = /etc/wireguard/${WG_INTERFACE}-down.sh"
fi

# Step 8: Create a systemd service for persistent rules
status "Creating systemd service for persistent iptables rules..."

cat > /etc/systemd/system/wireguard-docker-rules.service <<EOF
[Unit]
Description=Wireguard to Docker bridge iptables rules
After=network.target wg-quick@${WG_INTERFACE}.service docker.service
Requires=wg-quick@${WG_INTERFACE}.service

[Service]
Type=oneshot
ExecStart=/etc/wireguard/${WG_INTERFACE}-up.sh
RemainAfterExit=yes
ExecStop=/etc/wireguard/${WG_INTERFACE}-down.sh

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable wireguard-docker-rules.service

status "Systemd service created and enabled"

# Step 9: Fix any existing Docker daemon.json configuration
status "Ensuring Docker has proper iptables management..."

# Create or update Docker daemon config file
DOCKER_CONFIG="/etc/docker/daemon.json"
if [ -f "$DOCKER_CONFIG" ]; then
    # Check if we need to remove any previous workaround settings
    if grep -q '"iptables":\s*false' "$DOCKER_CONFIG" || grep -q '"bridge":\s*"none"' "$DOCKER_CONFIG"; then
        status "Updating Docker configuration to restore iptables management..."
        # Use jq if available for proper JSON manipulation
        if command -v jq >/dev/null 2>&1; then
            TMP_CONFIG=$(mktemp)
            jq 'del(.iptables) | del(.bridge)' "$DOCKER_CONFIG" > "$TMP_CONFIG"
            mv "$TMP_CONFIG" "$DOCKER_CONFIG"
        else
            # Simple sed-based approach (less robust but works for basic cases)
            sed -i 's/"iptables":\s*false,\?//g' "$DOCKER_CONFIG"
            sed -i 's/"bridge":\s*"none",\?//g' "$DOCKER_CONFIG"
            # Clean up any syntax issues this might create
            sed -i 's/,\s*}/}/g' "$DOCKER_CONFIG"
            sed -i 's/{,\s*/{/g' "$DOCKER_CONFIG"
            sed -i 's/,,/,/g' "$DOCKER_CONFIG"
        fi
    fi
else
    # Create default config file
    mkdir -p $(dirname "$DOCKER_CONFIG")
    echo '{}' > "$DOCKER_CONFIG"
fi

status "Docker iptables management properly configured"

# Step 10: Enable IP forwarding permanently
status "Enabling IP forwarding permanently..."
echo "net.ipv4.ip_forward=1" > /etc/sysctl.d/99-ip-forward.conf
sysctl -p /etc/sysctl.d/99-ip-forward.conf

# Step 11: Create a Docker restart hook to reapply rules
status "Creating Docker restart hook..."

cat > /usr/local/bin/docker-post-start.sh <<EOF
#!/bin/bash

# Wait for Docker to fully start
sleep 5

# Re-apply Wireguard rules
/etc/wireguard/${WG_INTERFACE}-up.sh

# Log that we ran
logger -t docker-vpn-fix "Re-applied Wireguard rules after Docker restart"
EOF

chmod +x /usr/local/bin/docker-post-start.sh

cat > /etc/systemd/system/docker-vpn-fix.service <<EOF
[Unit]
Description=Fix VPN rules after Docker starts
After=docker.service
Wants=docker.service

[Service]
Type=oneshot
ExecStart=/usr/local/bin/docker-post-start.sh
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable docker-vpn-fix.service

status "Docker restart hook configured"

# Final status
status "Setup complete!"
status "Bridge name: ${BRIDGE_NAME}"
status "Docker network name: ${DOCKER_NETWORK_NAME}"
status "IP subnet: ${BRIDGE_SUBNET}"

# Restart Docker to apply changes
status "Restarting Docker to apply changes..."
systemctl restart docker

# Provide instructions for next steps
if systemctl is-active --quiet "wg-quick@${WG_INTERFACE}"; then
    status "Wireguard is running. Restarting to apply new configuration..."
    systemctl restart "wg-quick@${WG_INTERFACE}"
    systemctl start wireguard-docker-rules.service
    status "Done! Your Wireguard VPN should now have access to Docker containers."
else
    warning "Wireguard is not running. Start it with:"
    echo "  systemctl start wg-quick@${WG_INTERFACE}"
    echo
    warning "Test your setup by:"
    echo "  1. Connect to your Wireguard VPN"
    echo "  2. Try to ping ${BRIDGE_IP}"
    echo "  3. Try to access a container in the ${DOCKER_NETWORK_NAME} network"
fi

status "To connect containers to this network, use:"
echo "  docker run --network ${DOCKER_NETWORK_NAME} your_container_image"
status "For existing containers:"
echo "  docker network connect ${DOCKER_NETWORK_NAME} container_name"

exit 0