#!/usr/bin/env bash
# WireGuard client provisioning for Linux and macOS.
# Based on the workflow of wireguard-provision.ps1: detect dependencies,
# offer automatic/manual installation, ask for client details, generate
# <device>.conf + <device>.png locally, and print the client's public key.
# Usage: chmod +x wireguard-provision.sh && ./wireguard-provision.sh
# Menu: generate new .conf/.png, or display QR from an existing .conf.
# Optional: ./wireguard-provision.sh /path/to/output-directory
# Requirements: Bash 3.2+, WireGuard wg, Python 3, Python venv/pip, qrcode[pil].
# Linux automatic system dependency installation supports apt, dnf, pacman;
# macOS supports Homebrew. Other platforms: install dependencies manually.
# Never publish .conf or .png: both contain the client private key.
set -euo pipefail
umask 077

say() { printf '%s\n' "$*"; }
# Homebrew is commonly installed outside the default PATH on macOS.
refresh_brew_path() {
  [[ "$OS" == Darwin ]] || return 0
  local brew_bin
  for brew_bin in /opt/homebrew/bin/brew /usr/local/bin/brew; do
    if [[ -x "$brew_bin" ]]; then
      eval "$("$brew_bin" shellenv)"
      break
    fi
  done
}
fail() { printf 'ERROR: %s\n' "$*" >&2; exit 1; }
ask() {
  local prompt="$1" default="${2-}" value
  while :; do
    if [[ -n "$default" ]]; then
      read -r -p "$prompt [$default]: " value || exit 0
      value="${value:-$default}"
    else
      read -r -p "$prompt: " value || exit 0
    fi
    [[ -n "$value" ]] && { REPLY_VALUE="$value"; return; }
    say 'Value cannot be empty.'
  done
}
confirm_install() {
  local response
  read -r -p "$1 Install automatically? [Y]es / [M]anual (exit): " response || exit 0
  case "$response" in
    y|Y|yes|YES|Yes) return 0 ;;
    *) say 'Manual installation selected. Install the missing dependencies and rerun this script.'; exit 0 ;;
  esac
}
install_system_deps() {
  local os="$1"
  if [[ "$os" == Darwin ]]; then
    command -v brew >/dev/null 2>&1 || fail 'Homebrew not found. Install Homebrew from https://brew.sh/ and rerun.'
    brew install wireguard-tools python
    refresh_brew_path
  elif command -v apt-get >/dev/null 2>&1; then
    sudo apt-get update
    sudo apt-get install -y wireguard-tools python3 python3-venv python3-pip
  elif command -v dnf >/dev/null 2>&1; then
    sudo dnf install -y wireguard-tools python3 python3-pip
  elif command -v pacman >/dev/null 2>&1; then
    sudo pacman -S --needed --noconfirm wireguard-tools python python-pip
  else
    fail 'No supported package manager found. Install wg, Python 3, Python venv/pip manually.'
  fi
}
valid_ipv4() {
  local ip="$1" octet
  local -a octets
  [[ "$ip" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]] || return 1
  IFS=. read -r -a octets <<< "$ip"
  [[ ${#octets[@]} -eq 4 ]] || return 1
  for octet in "${octets[@]}"; do
    [[ ${#octet} -le 3 ]] || return 1
    (( 10#$octet <= 255 )) || return 1
  done
}
valid_key() {
  local key="$1" decoded
  [[ "$key" =~ ^[A-Za-z0-9+/]{43}=$ ]] || return 1
  # Python validates the decoded length and canonical base64 representation.
  "$PYTHON" - "$key" <<'PY'
import base64, sys
try:
    key = sys.argv[1]
    raw = base64.b64decode(key, validate=True)
    sys.exit(0 if len(raw) == 32 and base64.b64encode(raw).decode() == key else 1)
except Exception:
    sys.exit(1)
PY
}
valid_routes() {
  local entry ip prefix
  local -a routes
  IFS=, read -r -a routes <<< "$1"
  (( ${#routes[@]} > 0 )) || return 1
  for entry in "${routes[@]}"; do
    entry="${entry//[[:space:]]/}"
    [[ "$entry" =~ ^([0-9.]+)/([0-9]{1,2})$ ]] || return 1
    ip="${BASH_REMATCH[1]}"; prefix="${BASH_REMATCH[2]}"
    valid_ipv4 "$ip" || return 1
    (( 10#$prefix <= 32 )) || return 1
  done
}
valid_dns() {
  local entry
  local -a entries
  IFS=, read -r -a entries <<< "$1"
  (( ${#entries[@]} > 0 )) || return 1
  for entry in "${entries[@]}"; do
    entry="${entry//[[:space:]]/}"
    valid_ipv4 "$entry" || return 1
  done
}

# Keep privilege escalation limited to the package manager; never generate
# credentials as root, where they could end up in /root.
[[ "$(id -u)" -ne 0 ]] || fail 'Run as your regular user, not sudo. The script will request sudo only for package installation.'
OS="$(uname -s)"
[[ "$OS" == Linux || "$OS" == Darwin ]] || fail "Unsupported OS: $OS"
refresh_brew_path
OUT="${1:-$PWD/wireguard-profiles}"

find_brew_python() {
  [[ "$OS" == Darwin ]] || return 1
  command -v brew >/dev/null 2>&1 || return 1
  local prefix candidate
  prefix="$(brew --prefix python 2>/dev/null)" || return 1
  for candidate in "$prefix/libexec/bin/python3" "$prefix/bin/python3" "$(brew --prefix)/bin/python3"; do
    if [[ -x "$candidate" ]] && "$candidate" -c 'import sys, venv; assert sys.version_info >= (3, 10)' >/dev/null 2>&1; then
      printf '%s\n' "$candidate"
      return 0
    fi
  done
  return 1
}
ensure_python() {
  PYTHON=''
  local candidate
  if [[ "$OS" == Darwin ]]; then
    candidate="$(find_brew_python)" || candidate=''
    [[ -n "$candidate" ]] && PYTHON="$candidate"
  fi
  if [[ -z "$PYTHON" ]] && command -v python3 >/dev/null 2>&1 && python3 -c 'import sys, venv; assert sys.version_info >= (3, 10)' >/dev/null 2>&1; then
    PYTHON="$(command -v python3)"
  fi
  if [[ -z "$PYTHON" ]]; then
    say 'Python 3 (>=3.10) with venv support is missing.'
    confirm_install 'Python is required.'
    install_system_deps "$OS"
    refresh_brew_path
    if [[ "$OS" == Darwin ]]; then
      candidate="$(find_brew_python)" || candidate=''
      [[ -n "$candidate" ]] && PYTHON="$candidate"
    fi
    if [[ -z "$PYTHON" ]]; then
      command -v python3 >/dev/null 2>&1 || fail 'Python was not found after installation.'
      PYTHON="$(command -v python3)"
    fi
  fi
}
ensure_wg() {
  if ! command -v wg >/dev/null 2>&1; then
    say 'WireGuard wg is missing.'
    confirm_install 'WireGuard tools are required for key generation.'
    install_system_deps "$OS"
    command -v wg >/dev/null 2>&1 || fail 'wg was not found after installation.'
  fi
}
# Recreate only this tool's private venv; never touch system Python or other venvs.
create_qr_venv() {
  local venv="$1"
  if ! "$PYTHON" -m venv "$venv"; then
    return 1
  fi
  [[ -x "$venv/bin/python" ]] && "$venv/bin/python" -m pip --version >/dev/null 2>&1
}
repair_venv_support() {
  local pyver
  if [[ "$OS" == Linux ]] && command -v apt-get >/dev/null 2>&1; then
    pyver="$("$PYTHON" -c 'import sys; print("%d.%d" % sys.version_info[:2])')"
    say "Installing Python venv support for Python $pyver..."
    sudo apt-get update
    if ! sudo apt-get install -y "python${pyver}-venv"; then
      sudo apt-get install -y python3-venv
    fi
  elif [[ "$OS" == Darwin ]] && command -v brew >/dev/null 2>&1; then
    brew install python
    refresh_brew_path
    local brew_python
    brew_python="$(find_brew_python)" || fail 'Homebrew Python not found after installation.'
    PYTHON="$brew_python"
  else
    fail 'Automatic venv repair is unavailable. Install Python venv/ensurepip support manually.'
  fi
}
ensure_qr() {
  ensure_python
  local state_home venv
  state_home="${XDG_STATE_HOME:-$HOME/.local/state}"
  venv="$state_home/wireguard-provision/venv"
  mkdir -p "$(dirname "$venv")"
  chmod 700 "$(dirname "$venv")"

  # A previous failed installation may have left a Python executable but no pip.
  if [[ -x "$venv/bin/python" ]] && ! "$venv/bin/python" -m pip --version >/dev/null 2>&1; then
    say 'Existing private Python environment has no pip. Attempting repair...'
    "$venv/bin/python" -m ensurepip --upgrade >/dev/null 2>&1 || true
  fi
  if [[ ! -x "$venv/bin/python" ]] || ! "$venv/bin/python" -m pip --version >/dev/null 2>&1; then
    say 'Preparing a complete isolated Python environment...'
    # The directory is owned by this script and contains no user profiles.
    rm -rf -- "$venv"
    if ! create_qr_venv "$venv"; then
      say 'Python venv/ensurepip support is missing or incomplete.'
      confirm_install 'Install the system Python venv package?'
      repair_venv_support
      rm -rf -- "$venv"
      create_qr_venv "$venv" || fail 'Virtual environment still lacks pip after repair.'
    fi
  fi
  QR_PYTHON="$venv/bin/python"
  if ! "$QR_PYTHON" -c 'import qrcode; from PIL import Image' >/dev/null 2>&1; then
    say 'Python QR dependencies are missing.'
    confirm_install 'Install qrcode[pil] in the private venv?'
    "$QR_PYTHON" -m pip install --upgrade pip || fail 'pip upgrade failed.'
    "$QR_PYTHON" -m pip install --upgrade 'qrcode[pil]' || fail 'qrcode/Pillow installation failed.'
    "$QR_PYTHON" -c 'import qrcode; from PIL import Image' >/dev/null 2>&1 || fail 'QR dependency verification failed.'
  fi
}

show_qr() {
  local files=() f choice index selected
  [[ -d "$OUT" ]] || { say "No profile directory found: $OUT"; return; }
  # Use globbing rather than ls parsing, preserving filenames with spaces.
  for f in "$OUT"/*.conf; do
    [[ -f "$f" && ! -L "$f" ]] && files+=("$f")
  done
  if (( ${#files[@]} == 0 )); then
    say "No .conf profiles found in: $OUT"
    return
  fi
  say ''
  say 'Available configurations:'
  for ((index=0; index<${#files[@]}; index++)); do
    printf '%d. %s\n' "$((index+1))" "${files[index]##*/}"
  done
  read -r -p "Select profile [1-${#files[@]}] (0 to cancel): " choice || return
  [[ "$choice" =~ ^[0-9]+$ ]] || { say 'Invalid selection.'; return; }
  (( 10#$choice == 0 )) && return
  (( 10#$choice >= 1 && 10#$choice <= ${#files[@]} )) || { say 'Invalid selection.'; return; }
  selected="${files[10#$choice-1]}"
  say ''
  say 'WARNING: The QR code contains the WireGuard private key.'
  read -r -p 'Display this sensitive QR code on screen? [y/N]: ' choice || return
  [[ "$choice" == y || "$choice" == Y ]] || { say 'Cancelled.'; return; }
  say ''
  ensure_qr
  "$QR_PYTHON" - "$selected" <<'PYQR'
import sys
import qrcode

with open(sys.argv[1], encoding='utf-8') as f:
    content = f.read()
# Keep a standard, scannable half-block QR. Do not resize the PNG.
qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_M, border=4)
qr.add_data(content)
qr.make(fit=True)
matrix = qr.get_matrix()
height, width = len(matrix), len(matrix[0])
utf8 = bool(sys.stdout.encoding and 'UTF' in sys.stdout.encoding.upper())

if not utf8:
    print('Terminal is not configured for UTF-8; using standard ASCII QR.')
    qr.print_ascii(invert=True)
else:
    # Two QR module rows per terminal row; black on white.
    glyphs = (' ', '▄', '▀', '█')
    for y in range(0, height, 2):
        top = matrix[y]
        bottom = matrix[y + 1] if y + 1 < height else [False] * width
        line = ''.join(glyphs[(int(t) << 1) | int(b)] for t, b in zip(top, bottom))
        print('\033[30;47m' + line + '\033[0m')
PYQR
  say 'Scan with the WireGuard mobile app. Avoid recording or sharing this terminal.'
}

generate_profile() {
  ensure_wg
  ensure_qr
  SUCCESS=0
  CREATED_CONF=0
  CREATED_PNG=0
  ask 'User / device name (filename, e.g. Laptop-Alice)'; NAME="$REPLY_VALUE"
[[ "$NAME" =~ ^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$ ]] || fail 'Invalid filename. Use 1-64 letters, digits, dots, underscores or dashes.'
ask 'Client tunnel IPv4/CIDR (e.g. 10.77.0.2/32)'; ADDRESS="$REPLY_VALUE"
[[ "$ADDRESS" =~ ^([0-9.]+)/32$ ]] && valid_ipv4 "${BASH_REMATCH[1]}" || fail 'Client address must be a valid IPv4 /32.'
IP="${ADDRESS%/*}"
ask 'DNS servers (comma-separated, e.g. 1.1.1.1, 9.9.9.9)'; DNS="$REPLY_VALUE"
valid_dns "$DNS" || fail 'DNS must contain IPv4 addresses separated by commas.'
ask 'Server endpoint (hostname:port, e.g. vpn.example.net:51820)'; ENDPOINT="$REPLY_VALUE"
[[ "$ENDPOINT" =~ ^[A-Za-z0-9.-]+:([0-9]{1,5})$ ]] || fail 'Endpoint must be hostname-or-IPv4:port.'
PORT="${BASH_REMATCH[1]}"
(( 10#$PORT >= 1 && 10#$PORT <= 65535 )) || fail 'Endpoint port must be 1-65535.'
ask 'AllowedIPs (IPv4 CIDR list)' '0.0.0.0/0'; ALLOWED="$REPLY_VALUE"
valid_routes "$ALLOWED" || fail 'AllowedIPs must be comma-separated IPv4 CIDR entries.'
ask 'WireGuard SERVER public key (Base64, no quotes)'; SERVER_KEY="$REPLY_VALUE"
valid_key "$SERVER_KEY" || fail 'Invalid server public key (32-byte Base64 expected).'

mkdir -p "$OUT"
OUT="$(cd "$OUT" && pwd -P)"
CONF="$OUT/$NAME.conf"; PNG="$OUT/$NAME.png"
[[ ! -e "$CONF" && ! -L "$CONF" && ! -e "$PNG" && ! -L "$PNG" ]] || fail 'Profile already exists. Choose another name; no files overwritten.'

PRIVATE_KEY="$(wg genkey)" || fail 'wg genkey failed.'
PUBLIC_KEY="$(printf '%s\n' "$PRIVATE_KEY" | wg pubkey)" || fail 'wg pubkey failed.'
valid_key "$PRIVATE_KEY" && valid_key "$PUBLIC_KEY" || fail 'WireGuard returned an invalid key.'

# Exclusive creation prevents overwrites, including in race conditions.
# File descriptors are kept open only for the config write.
cleanup() {
  if [[ "${SUCCESS:-0}" != 1 ]]; then
    [[ "${CREATED_CONF:-0}" == 1 ]] && rm -f -- "$CONF"
    [[ "${CREATED_PNG:-0}" == 1 ]] && rm -f -- "$PNG"
  fi
}
trap cleanup EXIT
( set -C; : > "$CONF" ) || fail 'Could not exclusively create config file.'
CREATED_CONF=1
chmod 600 "$CONF"
printf '[Interface]\nPrivateKey = %s\nAddress = %s\nDNS = %s\n\n[Peer]\nPublicKey = %s\nAllowedIPs = %s\nEndpoint = %s\nPersistentKeepalive = 25\n' \
  "$PRIVATE_KEY" "$ADDRESS" "$DNS" "$SERVER_KEY" "$ALLOWED" "$ENDPOINT" > "$CONF"

# PNG is generated locally; no key or config is sent to any online QR service.
# Reserve the filename first to prevent overwriting another file.
( set -C; : > "$PNG" ) || fail 'Could not exclusively create QR file.'
CREATED_PNG=1
chmod 600 "$PNG"
"$QR_PYTHON" - "$CONF" "$PNG" <<'PY'
import sys
import qrcode
with open(sys.argv[1], 'r', encoding='ascii') as f:
    config = f.read()
qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_M, box_size=8, border=4)
qr.add_data(config)
qr.make(fit=True)
qr.make_image(fill_color='black', back_color='white').save(sys.argv[2])
PY
[[ -s "$PNG" ]] || fail 'QR generation failed.'
SUCCESS=1
unset PRIVATE_KEY
say ''
say 'SUCCESS: WireGuard client configuration and QR code created.'
say "Config: $CONF"
say "QR PNG: $PNG"
say "Client PUBLIC key: $PUBLIC_KEY"
say ''
say 'Example RouterOS command (replace <WIREGUARD_INTERFACE> and verify IP is unused):'
printf '/interface wireguard peers add interface=<WIREGUARD_INTERFACE> public-key="%s" allowed-address=%s/32 comment="%s"\n' "$PUBLIC_KEY" "$IP" "$NAME"
say ''
say 'SECURITY: Both files contain the private key. Share securely; never publish them.'
say 'This tool does not check server-side IP conflicts or register peers automatically.'

  trap - EXIT
}

while :; do
  say ''
  say '=== WireGuard Provisioning Manager (Linux / macOS / WSL) ==='
  say '1. Generate new WireGuard configuration + QR PNG'
  say '2. Display existing configuration as terminal QR code'
  say '3. Exit'
  read -r -p 'Select an option [1-3]: ' TASK || { say ''; exit 0; }
  case "$TASK" in
    1) generate_profile ;;
    2) show_qr ;;
    3) say 'Goodbye.'; exit 0 ;;
    *) say 'Invalid selection. Choose 1, 2 or 3.' ;;
  esac
done