vpn-setup.sh

#!/usr/bin/env bash

# Interactive IPv4 WireGuard setup.  SSH is allowed only on the detected
# public interface; it is explicitly denied when arriving through wg0.

set -Eeuo pipefail
IFS=$'\n\t'
umask 077

die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; }
ask() { local label="$1" default="$2" value; read -r -p "$label [$default]: " value; printf '%s' "${value:-$default}"; }
required() { local label="$1" value; while :; do read -r -p "$label: " value; [[ -n "$value" ]] && { printf '%s' "$value"; return; }; done; }

[[ $EUID -eq 0 ]] || die 'Run this script as root.'
command -v openssl >/dev/null 2>&1 || die 'openssl is required.'

detect_os() {
    [[ -f /etc/os-release ]] || die 'Cannot detect operating system.'
    . /etc/os-release
    OS_ID="$ID"
    OS_VERSION="${VERSION_ID:-unknown}"
    OS_FAMILY="${ID_LIKE:-}"
}

is_debian_family() { [[ "$OS_FAMILY" == *debian* || "$OS_ID" == debian || "$OS_ID" == ubuntu ]]; }
is_rhel_family() { [[ "$OS_FAMILY" == *rhel* || "$OS_ID" == rocky || "$OS_ID" == almalinux || "$OS_ID" == rhel || "$OS_ID" == fedora ]]; }

ipv4_to_int() {
    local address="$1" octet value=0
    IFS=. read -r -a octets <<< "$address"
    [[ ${#octets[@]} -eq 4 ]] || return 1
    for octet in "${octets[@]}"; do
        [[ "$octet" =~ ^[0-9]{1,3}$ ]] || return 1
        (( 10#$octet <= 255 )) || return 1
        value=$(( (value << 8) + 10#$octet ))
    done
    printf '%u' "$value"
}

int_to_ipv4() {
    local value="$1"
    printf '%d.%d.%d.%d' $((value >> 24 & 255)) $((value >> 16 & 255)) $((value >> 8 & 255)) $((value & 255))
}

parse_cidr() {
    local cidr="$1" address prefix network mask address_int
    [[ "$cidr" =~ ^([^/]+)/([0-9]{1,2})$ ]] || return 1
    address="${BASH_REMATCH[1]}"
    prefix="${BASH_REMATCH[2]}"
    (( prefix <= 32 )) || return 1
    address_int="$(ipv4_to_int "$address")" || return 1
    if (( prefix == 0 )); then mask=0; else mask=$(( (0xffffffff << (32 - prefix)) & 0xffffffff )); fi
    network=$(( address_int & mask ))
    CIDR_PREFIX="$prefix"
    CIDR_NETWORK="$network"
    CIDR_BROADCAST=$(( network | (0xffffffff ^ mask) ))
}

validate_ipv4() { ipv4_to_int "$1" >/dev/null; }

detect_os

printf '  WireGuard VPN Setup with SSH Restriction\n'
printf '%s\n' '============================================'
printf 'Detected OS: %s %s\n' "$OS_ID" "$OS_VERSION"

printf '\n--- Step 1: Install WireGuard ---\n'
if ! command -v wg >/dev/null 2>&1; then
    if is_debian_family; then
        apt-get update
        DEBIAN_FRONTEND=noninteractive apt-get install -y wireguard wireguard-tools iptables iptables-persistent
    elif is_rhel_family; then
        dnf install -y wireguard-tools iptables iptables-services
        modprobe wireguard 2>/dev/null || true
    else
        die "Unsupported operating system: $OS_ID"
    fi
fi
for command_name in wg ip iptables iptables-save systemctl; do
    command -v "$command_name" >/dev/null 2>&1 || die "$command_name is required."
done

printf '\n--- Step 2: Gather and validate configuration ---\n'
SERVER_PUB_IP="$(required 'Server public IPv4 address')"
VPN_SUBNET="$(ask 'VPN subnet (IPv4 CIDR)' '10.8.0.0/24')"
VPN_ADDRESS="$(ask 'Server VPN IPv4 address (within subnet)' '10.8.0.1')"
VPN_PORT="$(ask 'WireGuard UDP port' '51820')"
VPN_DNS="$(ask 'DNS for VPN clients (comma-separated)' '1.1.1.1,8.8.8.8')"
SSH_DEFAULT_PORT='22'
if [[ -f /etc/ssh/sshd_config ]]; then
    SSH_DEFAULT_PORT="$(awk '$1 == "Port" { print $2; exit }' /etc/ssh/sshd_config)"
    SSH_DEFAULT_PORT="${SSH_DEFAULT_PORT:-22}"
fi
SSH_PORT="$(ask 'SSH port (used for firewall restriction)' "$SSH_DEFAULT_PORT")"
PEER_COUNT="$(ask 'Number of new VPN clients to configure' '1')"

validate_ipv4 "$SERVER_PUB_IP" || die 'Server public address must be a valid IPv4 address.'
parse_cidr "$VPN_SUBNET" || die 'VPN subnet must be a valid IPv4 CIDR.'
VPN_PREFIX="$CIDR_PREFIX"
VPN_NETWORK="$CIDR_NETWORK"
VPN_BROADCAST="$CIDR_BROADCAST"
VPN_ADDRESS_INT="$(ipv4_to_int "$VPN_ADDRESS")" || die 'Invalid VPN address.'
(( (VPN_ADDRESS_INT & (VPN_PREFIX == 0 ? 0 : (0xffffffff << (32 - VPN_PREFIX)) & 0xffffffff)) == VPN_NETWORK )) || die 'VPN address is outside the VPN subnet.'
(( VPN_PREFIX >= 31 || (VPN_ADDRESS_INT != VPN_NETWORK && VPN_ADDRESS_INT != VPN_BROADCAST) )) || die 'VPN address cannot be the network or broadcast address.'
[[ "$VPN_PORT" =~ ^[0-9]+$ ]] && (( VPN_PORT >= 1 && VPN_PORT <= 65535 )) || die 'Invalid WireGuard port.'
[[ "$SSH_PORT" =~ ^[0-9]+$ ]] && (( SSH_PORT >= 1 && SSH_PORT <= 65535 )) || die 'Invalid SSH port.'
[[ "$PEER_COUNT" =~ ^[0-9]+$ ]] && (( PEER_COUNT >= 1 && PEER_COUNT <= 254 )) || die 'Client count must be between 1 and 254.'
IFS=',' read -ra DNS_LIST <<< "$VPN_DNS"
for dns in "${DNS_LIST[@]}"; do validate_ipv4 "${dns//[[:space:]]/}" || die "Invalid DNS IPv4 address: $dns"; done

WG_DIR="/etc/wireguard"
WG_CONF="$WG_DIR/wg0.conf"
CLIENT_DIR="$WG_DIR/clients"
install -d -m 700 "$WG_DIR" "$CLIENT_DIR"

printf '\n--- Step 3: Generate server keys ---\n'
if [[ ! -f "$WG_DIR/server.key" ]]; then
    wg genkey > "$WG_DIR/server.key.new"
    chmod 600 "$WG_DIR/server.key.new"
    mv -f "$WG_DIR/server.key.new" "$WG_DIR/server.key"
fi
SERVER_PRIVATE_KEY="$(<"$WG_DIR/server.key")"
[[ "$SERVER_PRIVATE_KEY" =~ ^[A-Za-z0-9+/=]{40,46}$ ]] || die 'Existing server.key is not a valid WireGuard private key.'
SERVER_PUBLIC_KEY="$(printf '%s\n' "$SERVER_PRIVATE_KEY" | wg pubkey)" || die 'Could not derive server public key.'
printf '%s\n' "$SERVER_PUBLIC_KEY" > "$WG_DIR/server.pub"
chmod 600 "$WG_DIR/server.pub"

printf '\n--- Step 4: Determine public interface ---\n'
DEFAULT_IFACE="$(ip -4 route show default | awk 'NR == 1 { print $5 }')"
[[ "$DEFAULT_IFACE" =~ ^[A-Za-z0-9_.:-]+$ ]] || die 'Could not determine default IPv4 interface.'

printf '\n--- Step 5: Generate client configurations ---\n'
declare -a PEERS=()
declare -A USED_IPS=()
USED_IPS["$VPN_ADDRESS_INT"]=1
if [[ -f "$WG_CONF" ]]; then
    while read -r existing_ip; do
        [[ -n "$existing_ip" ]] && USED_IPS["$(ipv4_to_int "$existing_ip")"]=1
    done < <(awk '/^AllowedIPs = / { sub("/32$", "", $3); print $3 }' "$WG_CONF")
fi

for ((i = 1; i <= PEER_COUNT; i++)); do
    CLIENT_NAME="$(ask "Client $i name" "client-$i")"
    [[ "$CLIENT_NAME" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]*$ ]] || die 'Client names may contain only letters, numbers, dot, underscore, and hyphen.'
    CLIENT_FILE="$CLIENT_DIR/$CLIENT_NAME.conf"
    [[ ! -e "$CLIENT_FILE" ]] || die "Client already exists: $CLIENT_NAME (refusing to overwrite it)."
    CLIENT_IP="$(ask "Client $i VPN address (blank for automatic allocation)" '')"
    if [[ -z "$CLIENT_IP" ]]; then
        for ((candidate = VPN_NETWORK + 1; candidate < VPN_BROADCAST; candidate++)); do
            if (( candidate <= 0xffffffff )) && [[ -z "${USED_IPS[$candidate]+set}" ]]; then CLIENT_IP="$(int_to_ipv4 "$candidate")"; break; fi
        done
    fi
    CLIENT_IP_INT="$(ipv4_to_int "$CLIENT_IP")" || die "Invalid client address: $CLIENT_IP"
    (( (CLIENT_IP_INT & (VPN_PREFIX == 0 ? 0 : (0xffffffff << (32 - VPN_PREFIX)) & 0xffffffff)) == VPN_NETWORK )) || die 'Client address is outside the VPN subnet.'
    (( VPN_PREFIX >= 31 || (CLIENT_IP_INT != VPN_NETWORK && CLIENT_IP_INT != VPN_BROADCAST) )) || die 'Client address cannot be the network or broadcast address.'
    [[ -z "${USED_IPS[$CLIENT_IP_INT]+set}" ]] || die "Duplicate VPN address: $CLIENT_IP"
    USED_IPS["$CLIENT_IP_INT"]=1
    CLIENT_KEY="$(wg genkey)"
    CLIENT_PUB="$(printf '%s\n' "$CLIENT_KEY" | wg pubkey)"
    {
        printf '[Interface]\nPrivateKey = %s\nAddress = %s/32\nDNS = %s\n\n' "$CLIENT_KEY" "$CLIENT_IP" "$VPN_DNS"
        printf '[Peer]\nPublicKey = %s\nEndpoint = %s:%s\nAllowedIPs = 0.0.0.0/0\nPersistentKeepalive = 25\n' "$SERVER_PUBLIC_KEY" "$SERVER_PUB_IP" "$VPN_PORT"
    } > "$CLIENT_FILE.new"
    chmod 600 "$CLIENT_FILE.new"
    mv -f "$CLIENT_FILE.new" "$CLIENT_FILE"
    PEERS+=("$CLIENT_NAME:$CLIENT_IP:$CLIENT_PUB")
    printf 'Client %s configured with IP %s\n' "$CLIENT_NAME" "$CLIENT_IP"
done

printf '\n--- Step 6: Write WireGuard configuration ---\n'
WG_TMP="$WG_CONF.new"
{
    printf '[Interface]\nPrivateKey = %s\nAddress = %s/%s\nListenPort = %s\n' "$SERVER_PRIVATE_KEY" "$VPN_ADDRESS" "$VPN_PREFIX" "$VPN_PORT"
    VPN_NETWORK_IP="$(int_to_ipv4 "$VPN_NETWORK")"
    printf 'PostUp = iptables -C FORWARD -i wg0 -j ACCEPT 2>/dev/null || iptables -A FORWARD -i wg0 -j ACCEPT; iptables -C FORWARD -o wg0 -j ACCEPT 2>/dev/null || iptables -A FORWARD -o wg0 -j ACCEPT; iptables -t nat -C POSTROUTING -s %s/%s -o %s -j MASQUERADE 2>/dev/null || iptables -t nat -A POSTROUTING -s %s/%s -o %s -j MASQUERADE\n' "$VPN_NETWORK_IP" "$VPN_PREFIX" "$DEFAULT_IFACE" "$VPN_NETWORK_IP" "$VPN_PREFIX" "$DEFAULT_IFACE"
    printf 'PostDown = while iptables -D FORWARD -i wg0 -j ACCEPT 2>/dev/null; do :; done; while iptables -D FORWARD -o wg0 -j ACCEPT 2>/dev/null; do :; done; while iptables -t nat -D POSTROUTING -s %s/%s -o %s -j MASQUERADE 2>/dev/null; do :; done\n\n' "$VPN_NETWORK_IP" "$VPN_PREFIX" "$DEFAULT_IFACE"
    if [[ -f "$WG_CONF" ]]; then awk 'seen || /^\[Peer\]/{seen=1} seen' "$WG_CONF"; fi
    for peer in "${PEERS[@]}"; do IFS=: read -r name ip pub <<< "$peer"; printf '[Peer]\nPublicKey = %s\nAllowedIPs = %s/32\n\n' "$pub" "$ip"; done
} > "$WG_TMP"
chmod 600 "$WG_TMP"
mv -f "$WG_TMP" "$WG_CONF"

printf '\n--- Step 7: Configure forwarding and managed firewall rules ---\n'
sysctl -w net.ipv4.ip_forward=1 >/dev/null
install -d -m 755 /etc/sysctl.d
printf 'net.ipv4.ip_forward=1\n' > /etc/sysctl.d/99-wireguard-forwarding.conf

# A dedicated chain makes repeated runs converge and places the wg0 SSH deny
# before any broad wg0 accept rule or rules left by another firewall script.
iptables -N REIMAGE_WG 2>/dev/null || true
iptables -F REIMAGE_WG
while iptables -D INPUT -j REIMAGE_WG 2>/dev/null; do :; done
iptables -I INPUT 1 -j REIMAGE_WG
iptables -A REIMAGE_WG -p udp --dport "$VPN_PORT" -j ACCEPT
iptables -A REIMAGE_WG -i wg0 -p tcp --dport "$SSH_PORT" -j DROP
iptables -A REIMAGE_WG -p tcp --dport "$SSH_PORT" -i "$DEFAULT_IFACE" -j ACCEPT
iptables -A REIMAGE_WG -p tcp --dport "$SSH_PORT" -j DROP
iptables -A REIMAGE_WG -i wg0 -j ACCEPT

printf '\n--- Step 8: Enable and start WireGuard ---\n'
systemctl enable wg-quick@wg0
systemctl restart wg-quick@wg0

printf '\n--- Step 9: Persist firewall rules ---\n'
if is_debian_family; then
    install -d -m 755 /etc/iptables
    iptables-save > /etc/iptables/rules.v4
    systemctl enable netfilter-persistent 2>/dev/null || true
elif is_rhel_family; then
    iptables-save > /etc/sysconfig/iptables
    systemctl enable iptables 2>/dev/null || true
else
    printf 'WARNING: firewall persistence is not configured for %s.\n' "$OS_ID" >&2
fi

printf '\nVPN setup complete.\nServer: %s:%s\nVPN subnet: %s/%s\nClients: %s\n' "$SERVER_PUB_IP" "$VPN_PORT" "$(int_to_ipv4 "$VPN_NETWORK")" "$VPN_PREFIX" "$PEER_COUNT"
printf 'Client configs: %s/\n' "$CLIENT_DIR"
printf 'SSH is allowed only via public interface %s on port %s; wg0 SSH is denied.\n' "$DEFAULT_IFACE" "$SSH_PORT"