#!/bin/bash
set -e

# Définition des couleurs pour les sorties
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
NC='\033[0m' # Pas de couleur

# Fonction d'affichage de l'aide
usage() {
  echo "Usage: $0 [-p port] [-a address]"
  echo "  -p Port d'écoute pour WireGuard (par défaut: 8022)"
  echo "  -a Adresse IP et masque pour l'interface VPN (par défaut: 10.0.0.1/24)"
  exit 1
}

# Valeurs par défaut
PORT=8022
WG_ADDRESS="10.0.0.1/24"

# Traitement des options
while getopts "p:a:h" opt; do
  case $opt in
    p)
      PORT=$OPTARG
      ;;
    a)
      WG_ADDRESS=$OPTARG
      ;;
    h)
      usage
      ;;
    \?)
      usage
      ;;
  esac
done

# Validation du port (doit être un entier entre 1 et 65535)
if ! [[ "$PORT" =~ ^[0-9]+$ ]] || [ "$PORT" -lt 1 ] || [ "$PORT" -gt 65535 ]; then
  echo -e "${RED}[ERREUR] Le port spécifié ($PORT) est invalide.${NC}"
  exit 1
fi

# Vérification des privilèges root
if [ "$(id -u)" -ne 0 ]; then
  echo -e "${RED}[ERREUR] Ce script doit être exécuté en tant que root.${NC}"
  exit 1
fi

# Fonction pour afficher des messages d'information
log_info() {
  echo -e "${BLUE}[INFO] $1${NC}"
}

# Mise à jour et installation de WireGuard
log_info "Mise à jour du système et installation de WireGuard..."
apt update && apt install -y wireguard

WG_DIR="/etc/wireguard"
mkdir -p "$WG_DIR"
chmod 700 "$WG_DIR"

# Vérifier si les clés existent déjà afin d'éviter de les écraser
if [[ -f "$WG_DIR/server_private.key" ]] && [[ -f "$WG_DIR/server_public.key" ]]; then
  log_info "Les clés du serveur existent déjà. Utilisation des clés existantes."
else
  log_info "Génération des clés du serveur..."
  wg genkey | tee "$WG_DIR/server_private.key" | wg pubkey > "$WG_DIR/server_public.key"
  chmod 600 "$WG_DIR/server_private.key" "$WG_DIR/server_public.key"
fi

SERVER_PRIV=$(cat "$WG_DIR/server_private.key")
SERVER_PUB=$(cat "$WG_DIR/server_public.key")

# Création du fichier de configuration du serveur
log_info "Création du fichier de configuration du serveur ($WG_DIR/wg0.conf)..."
cat > "$WG_DIR/wg0.conf" <<EOF
[Interface]
Address = $WG_ADDRESS
ListenPort = $PORT
PrivateKey = $SERVER_PRIV
# Activation du forwarding IP lors du démarrage de l'interface
PostUp   = sysctl -w net.ipv4.ip_forward=1
PostDown = sysctl -w net.ipv4.ip_forward=0
EOF
chmod 600 "$WG_DIR/wg0.conf"

# Activation et démarrage de wg-quick
log_info "Activation de wg-quick pour le démarrage automatique..."
systemctl enable wg-quick@wg0

log_info "Démarrage de l'interface WireGuard..."
systemctl start wg-quick@wg0

log_info "Autorise ssh que dans le réseau vpn"
iptables -A INPUT -s 10.0.0.0/24 -d 10.0.0.0/24 -j ACCEPT
iptables -A INPUT -d 10.0.0.0/24 -j DROP
iptables-save > /etc/iptables/rules.v4

echo -e "${GREEN}[SUCCÈS] Installation du serveur VPN terminée.${NC}"
echo -e "${GREEN}[INFO] Clé publique du serveur (à utiliser pour la configuration client) :${NC}"
echo "$SERVER_PUB"