#!/bin/bash
#
# coolervpn: your own WireGuard VPN on Ubuntu 22.04 / 24.04, in one command.
#
#   Single VPN   phone/laptop -> server -> internet
#   Double VPN   phone/laptop -> entry server -> exit server -> internet
#
# Run as root:  sudo bash coolervpn.sh
# Non-interactive: --single | --double-exit | --double-entry <token>
#
# MIT License. Inspired by Nyr/wireguard-install; written from scratch.

set -euo pipefail

WG_DIR=/etc/wireguard
STATE=$WG_DIR/coolervpn.env
CLIENT_DIR=$WG_DIR/clients
CLIENT_NET=10.7.0            # clients get 10.7.0.2 ... 10.7.0.254
LINK_NET=10.66.66            # entry <-> exit link in Double mode
RT_TABLE=51820

red()   { printf '\033[31m%s\033[0m\n' "$*"; }
green() { printf '\033[32m%s\033[0m\n' "$*"; }
bold()  { printf '\033[1m%s\033[0m\n' "$*"; }
die()   { red "Error: $*" >&2; exit 1; }

# ---------------------------------------------------------------- checks

check_system() {
	[[ $EUID -eq 0 ]] || die "run this with sudo: sudo bash $0"
	[[ -r /etc/os-release ]] || die "can't detect the OS"
	. /etc/os-release
	[[ ${ID:-} == ubuntu ]] || die "this script supports Ubuntu 22.04 and 24.04 only (found ${PRETTY_NAME:-unknown})"
	case ${VERSION_ID:-} in
		22.04|24.04) ;;
		*) die "this script supports Ubuntu 22.04 and 24.04 only (found $VERSION_ID)" ;;
	esac
	if systemd-detect-virt -cq 2>/dev/null; then
		red "Warning: this looks like a container (OpenVZ/LXC). WireGuard may not work here."
	fi
}

public_ip() {
	local ip=""
	for url in https://api.ipify.org https://ifconfig.me/ip https://icanhazip.com; do
		ip=$(curl -4 -fsS --max-time 5 "$url" 2>/dev/null | tr -d '[:space:]') || true
		[[ $ip =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]] && { echo "$ip"; return; }
	done
	ip -4 route get 1.1.1.1 | awk '{for(i=1;i<=NF;i++) if($i=="src") print $(i+1)}'
}

public_iface() {
	ip -4 route show default | awk '{for(i=1;i<=NF;i++) if($i=="dev") {print $(i+1); exit}}'
}

ask() {  # ask "question" default -> echoes answer
	local reply
	read -rp "$1 [$2]: " reply </dev/tty || true
	echo "${reply:-$2}"
}

# Which DNS the devices use. Enter = Google.
choose_dns() {   # $1 = address of our own resolver for this mode
	local own=$1 c
	echo >&2
	echo "Which DNS should your devices use?" >&2
	echo "  1) Google       8.8.8.8 (default)" >&2
	echo "  2) Cloudflare   1.1.1.1" >&2
	echo "  3) Quad9        9.9.9.9 (blocks known malicious domains)" >&2
	echo "  4) Private      your own resolver on the server, no DNS company sees your lookups" >&2
	c=$(ask "Choose" "1")
	case $c in
		2) echo "1.1.1.1, 1.0.0.1" ;;
		3) echo "9.9.9.9, 149.112.112.112" ;;
		4) echo "$own" ;;
		*) echo "8.8.8.8, 8.8.4.4" ;;
	esac
}

valid_port() { [[ $1 =~ ^[0-9]+$ ]] && (( $1 >= 1 && $1 <= 65535 )); }

# ---------------------------------------------------------------- install

install_packages() {
	green "Installing WireGuard, Unbound and tools from Ubuntu's repositories..."
	export DEBIAN_FRONTEND=noninteractive
	apt-get update -q
	apt-get install -yq wireguard wireguard-tools qrencode unbound iptables curl unattended-upgrades
	# Security updates install themselves from now on.
	cat > /etc/apt/apt.conf.d/20auto-upgrades <<-EOF
	APT::Periodic::Update-Package-Lists "1";
	APT::Periodic::Unattended-Upgrade "1";
	EOF
}

enable_forwarding() {
	cat > /etc/sysctl.d/99-coolervpn.conf <<-EOF
	net.ipv4.ip_forward = 1
	net.ipv4.conf.all.rp_filter = 2
	EOF
	sysctl -q -p /etc/sysctl.d/99-coolervpn.conf
}

# Unbound: a private DNS resolver on the server, reachable only through the tunnel.
setup_unbound() {
	local listen=$1 allow=$2
	mkdir -p /etc/unbound/unbound.conf.d
	cat > /etc/unbound/unbound.conf.d/coolervpn.conf <<-EOF
	server:
	    interface: $listen
	    ip-freebind: yes
	    port: 53
	    do-ip6: no
	    access-control: 0.0.0.0/0 refuse
	$(for net in $allow; do echo "    access-control: $net allow"; done)
	    hide-identity: yes
	    hide-version: yes
	    qname-minimisation: yes
	    harden-glue: yes
	    harden-dnssec-stripped: yes
	    use-caps-for-id: yes
	    prefetch: yes
	    cache-min-ttl: 60
	    log-queries: no
	    verbosity: 0
	EOF
	unbound-checkconf >/dev/null || die "Unbound config check failed"
	systemctl enable --now unbound >/dev/null 2>&1
	systemctl restart unbound
}

# iptables rules are added on interface up and removed on down (wg-quick runs these).
# -I (insert) keeps them ahead of the REJECT rules some cloud images ship (Oracle).
write_rules() {
	local file=$1 mode=$2 port=$3 pub=$4
	{
		echo "#!/bin/bash"
		echo "# generated by coolervpn ($mode). Called as: $file up|down"
		echo 'A=-I; [[ ${1:-up} == down ]] && A=-D'
		echo 'run() { iptables "$@" 2>/dev/null || true; }'
		case $mode in
		single)
			echo "run \$A INPUT -p udp --dport $port -j ACCEPT"
			echo "run \$A INPUT -i wg0 -p udp --dport 53 -j ACCEPT"
			echo "run \$A INPUT -i wg0 -p tcp --dport 53 -j ACCEPT"
			echo "run \$A FORWARD -i wg0 -o $pub -j ACCEPT"
			echo "run \$A FORWARD -i $pub -o wg0 -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT"
			echo "run \$A FORWARD -i wg0 -o wg0 -j DROP"
			echo "run -t nat \$A POSTROUTING -s $CLIENT_NET.0/24 -o $pub -j MASQUERADE"
			echo "run -t mangle \$A FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu"
			;;
		exit)
			echo "run \$A INPUT -p udp --dport $port -j ACCEPT"
			echo "run \$A INPUT -i wg1 -p udp --dport 53 -j ACCEPT"
			echo "run \$A INPUT -i wg1 -p tcp --dport 53 -j ACCEPT"
			echo "run \$A FORWARD -i wg1 -o $pub -j ACCEPT"
			echo "run \$A FORWARD -i $pub -o wg1 -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT"
			echo "run -t nat \$A POSTROUTING -s $CLIENT_NET.0/24 -o $pub -j MASQUERADE"
			echo "run -t nat \$A POSTROUTING -s $LINK_NET.0/24 -o $pub -j MASQUERADE"
			echo "run -t mangle \$A FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu"
			;;
		entry)
			echo "run \$A INPUT -p udp --dport $port -j ACCEPT"
			# Clients may only go through the exit tunnel. Never straight out of this server.
			echo "run \$A FORWARD -i wg0 -o $pub -j DROP"
			echo "run \$A FORWARD -i wg0 -o wg0 -j DROP"
			echo "run \$A FORWARD -i wg0 -o wg1 -j ACCEPT"
			echo "run \$A FORWARD -i wg1 -o wg0 -j ACCEPT"
			echo "run -t mangle \$A FORWARD -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu"
			# Policy routing: everything from clients goes to table $RT_TABLE.
			# If the exit tunnel is down, that table only has an "unreachable" route: kill switch.
			echo 'if [[ ${1:-up} == up ]]; then'
			echo "  ip rule add to $CLIENT_NET.0/24 lookup main priority 100 2>/dev/null || true"
			echo "  ip rule add from $CLIENT_NET.0/24 lookup $RT_TABLE priority 110 2>/dev/null || true"
			echo "  ip route replace unreachable default metric 4096 table $RT_TABLE"
			echo 'else'
			echo "  ip rule del to $CLIENT_NET.0/24 lookup main priority 100 2>/dev/null || true"
			echo "  ip rule del from $CLIENT_NET.0/24 lookup $RT_TABLE priority 110 2>/dev/null || true"
			echo "  ip route flush table $RT_TABLE 2>/dev/null || true"
			echo 'fi'
			;;
		esac
	} > "$file"
	chmod 700 "$file"
}

save_state() {
	umask 077
	cat > "$STATE" <<-EOF
	MODE=$MODE
	ENDPOINT=$ENDPOINT
	PORT=$PORT
	DNS="$DNS"
	MTU="${MTU:-}"
	EOF
}

load_state() { [[ -f $STATE ]] && . "$STATE"; }

cloud_firewall_note() {
	echo
	bold "Important: open UDP port $1 in your cloud provider's firewall"
	echo "  Oracle: VCN > Security List > Ingress rule, UDP, port $1"
	echo "  Azure:  VM > Networking > Add inbound port rule, UDP, port $1"
	echo "  AWS:    Security group > Inbound rule, Custom UDP, port $1"
	echo "  Hetzner/DigitalOcean/Vultr: only if you turned their firewall on"
}

# ---------------------------------------------------------------- single / entry

install_client_server() {   # MODE is single or entry
	local pub server_key
	pub=$(public_iface)
	[[ -n $pub ]] || die "can't find the network interface with the default route"

	ENDPOINT=$(ask "Public IP or domain that clients connect to" "$(public_ip)")
	PORT=$(ask "UDP port for WireGuard" "51820")
	valid_port "$PORT" || die "bad port: $PORT"
	if [[ $MODE == single ]]; then
		DNS=$(choose_dns "$CLIENT_NET.1")
	fi

	install_packages
	enable_forwarding
	mkdir -p "$CLIENT_DIR"
	chmod 700 "$WG_DIR" "$CLIENT_DIR"

	server_key=$(wg genkey)
	write_rules "$WG_DIR/rules-wg0.sh" "$MODE" "$PORT" "$pub"

	umask 077
	cat > "$WG_DIR/wg0.conf" <<-EOF
	# coolervpn: $MODE
	[Interface]
	Address = $CLIENT_NET.1/24
	ListenPort = $PORT
	PrivateKey = $server_key
	PostUp = $WG_DIR/rules-wg0.sh up
	PostDown = $WG_DIR/rules-wg0.sh down
	EOF

	if [[ $MODE == single ]]; then
		MTU=""
		[[ $DNS == "$CLIENT_NET.1" ]] && setup_unbound "$CLIENT_NET.1" "$CLIENT_NET.0/24"
	fi

	systemctl enable --now wg-quick@wg0 >/dev/null 2>&1 || die "wg0 failed to start: journalctl -u wg-quick@wg0"
	save_state
}

# ---------------------------------------------------------------- double: exit

install_exit() {
	local pub exit_key entry_key entry_pub exit_pub endpoint token
	MODE=exit
	pub=$(public_iface)
	[[ -n $pub ]] || die "can't find the network interface with the default route"
	endpoint=$(ask "Public IP of THIS (exit) server" "$(public_ip)")
	PORT=$(ask "UDP port for the link between the two servers" "51821")
	valid_port "$PORT" || die "bad port: $PORT"

	install_packages
	enable_forwarding
	mkdir -p "$WG_DIR"; chmod 700 "$WG_DIR"

	exit_key=$(wg genkey);  exit_pub=$(echo "$exit_key" | wg pubkey)
	entry_key=$(wg genkey); entry_pub=$(echo "$entry_key" | wg pubkey)
	write_rules "$WG_DIR/rules-wg1.sh" exit "$PORT" "$pub"

	umask 077
	cat > "$WG_DIR/wg1.conf" <<-EOF
	# coolervpn: double, exit server
	[Interface]
	Address = $LINK_NET.1/24
	ListenPort = $PORT
	PrivateKey = $exit_key
	PostUp = $WG_DIR/rules-wg1.sh up
	PostDown = $WG_DIR/rules-wg1.sh down

	[Peer]
	# entry server (and the clients behind it)
	PublicKey = $entry_pub
	AllowedIPs = $LINK_NET.2/32, $CLIENT_NET.0/24
	EOF

	setup_unbound "$LINK_NET.1" "$LINK_NET.0/24 $CLIENT_NET.0/24"
	systemctl enable --now wg-quick@wg1 >/dev/null 2>&1 || die "wg1 failed to start: journalctl -u wg-quick@wg1"
	ENDPOINT=$endpoint; DNS=$LINK_NET.1; MTU=""
	save_state

	token=$(printf 'cv1|%s|%s|%s|%s' "$endpoint" "$PORT" "$exit_pub" "$entry_key" | base64 -w0)
	echo
	green "Exit server is ready."
	cloud_firewall_note "$PORT"
	echo
	bold "Now run this script on the ENTRY server, choose 2 (Double VPN), then 2 (ENTRY server), and paste this key:"
	echo
	echo "$token"
	echo
	red "This key contains a private key. Paste it only on your entry server, then clear your terminal."
	echo "(You can show it again: sudo bash $0 --show-token)"
	umask 077; echo "$token" > "$WG_DIR/entry-token.txt"
}

# ---------------------------------------------------------------- double: entry

install_entry() {
	local token=${1:-} decoded tag ex_ip ex_port ex_pub en_key
	MODE=entry
	if [[ -z $token ]]; then
		read -rp "Paste the key from the exit server: " token </dev/tty
	fi
	decoded=$(echo "$token" | base64 -d 2>/dev/null) || die "that key isn't valid"
	IFS='|' read -r tag ex_ip ex_port ex_pub en_key <<< "$decoded"
	[[ $tag == cv1 && -n $ex_ip && -n $en_key ]] || die "that key isn't valid"

	# Public DNS or the exit server's private resolver. Either way the
	# lookups leave the internet from the exit server, never from here.
	DNS=$(choose_dns "$LINK_NET.1")
	MTU=1340             # room for two layers of WireGuard
	install_client_server

	umask 077
	cat > "$WG_DIR/wg1.conf" <<-EOF
	# coolervpn: double, link to the exit server
	[Interface]
	Address = $LINK_NET.2/24
	PrivateKey = $en_key
	MTU = 1420
	Table = off
	PostUp = ip route replace default dev wg1 metric 10 table $RT_TABLE
	PostDown = ip route del default dev wg1 metric 10 table $RT_TABLE 2>/dev/null || true

	[Peer]
	# exit server
	PublicKey = $ex_pub
	Endpoint = $ex_ip:$ex_port
	AllowedIPs = 0.0.0.0/0
	PersistentKeepalive = 25
	EOF

	systemctl restart wg-quick@wg0
	systemctl enable --now wg-quick@wg1 >/dev/null 2>&1 || die "wg1 failed to start: journalctl -u wg-quick@wg1"

	sleep 2
	if ping -c 2 -W 2 -I wg1 "$LINK_NET.1" >/dev/null 2>&1; then
		green "Link to the exit server works."
	else
		red "Can't reach the exit server through the tunnel yet."
		echo "Check that UDP port $ex_port is open in the exit server's cloud firewall."
	fi
}

# ---------------------------------------------------------------- clients

next_ip() {
	local used n
	used=$(grep -ho "$CLIENT_NET\.[0-9]*/32" "$WG_DIR/wg0.conf" | cut -d. -f4 | cut -d/ -f1 || true)
	for n in $(seq 2 254); do
		grep -qx "$n" <<< "$used" || { echo "$CLIENT_NET.$n"; return; }
	done
	die "no free addresses left (253 clients max)"
}

add_client() {
	local name=${1:-} key pubk psk ip conf server_pub
	load_state
	[[ ${MODE:-} == single || ${MODE:-} == entry ]] || die "clients are added on a Single or Entry server"
	if [[ -z $name ]]; then
		name=$(ask "Name for the new device (letters, digits, - and _)" "phone")
	fi
	[[ $name =~ ^[A-Za-z0-9_-]{1,32}$ ]] || die "use only letters, digits, - and _ (max 32)"
	[[ -e $CLIENT_DIR/$name.conf ]] && die "a device called $name already exists"

	key=$(wg genkey); pubk=$(echo "$key" | wg pubkey); psk=$(wg genpsk)
	ip=$(next_ip)
	server_pub=$(wg show wg0 public-key)

	umask 077
	cat >> "$WG_DIR/wg0.conf" <<-EOF

	# BEGIN_PEER $name
	[Peer]
	PublicKey = $pubk
	PresharedKey = $psk
	AllowedIPs = $ip/32
	# END_PEER $name
	EOF
	wg set wg0 peer "$pubk" preshared-key <(echo "$psk") allowed-ips "$ip/32"

	conf=$CLIENT_DIR/$name.conf
	{
		echo "[Interface]"
		echo "PrivateKey = $key"
		echo "Address = $ip/32"
		echo "DNS = $DNS"
		[[ -n ${MTU:-} ]] && echo "MTU = $MTU"
		echo
		echo "[Peer]"
		echo "PublicKey = $server_pub"
		echo "PresharedKey = $psk"
		# ::/0 too, so IPv6 can't leak around the tunnel
		echo "AllowedIPs = 0.0.0.0/0, ::/0"
		echo "Endpoint = $ENDPOINT:$PORT"
		echo "PersistentKeepalive = 25"
	} > "$conf"

	echo
	green "Device \"$name\" added."
	echo "Scan this with the WireGuard app (iPhone/Android: + > Create from QR code):"
	qrencode -t ansiutf8 < "$conf"
	echo "Config file for a computer: $conf"
}

remove_client() {
	local name=${1:-} pubk
	load_state
	if [[ -z $name ]]; then
		list_clients
		name=$(ask "Device to remove" "")
	fi
	[[ -f $CLIENT_DIR/$name.conf ]] || die "no device called $name"
	pubk=$(sed -n "/^# BEGIN_PEER $name\$/,/^# END_PEER $name\$/p" "$WG_DIR/wg0.conf" | awk -F' = ' '/^PublicKey/{print $2}')
	[[ -n $pubk ]] && wg set wg0 peer "$pubk" remove
	sed -i "/^# BEGIN_PEER $name\$/,/^# END_PEER $name\$/d" "$WG_DIR/wg0.conf"
	rm -f "$CLIENT_DIR/$name.conf"
	green "Device \"$name\" removed."
}

list_clients() {
	local f
	echo "Devices:"
	for f in "$CLIENT_DIR"/*.conf; do
		[[ -e $f ]] || { echo "  (none)"; return; }
		echo "  $(basename "$f" .conf)  $(awk -F' = ' '/^Address/{print $2}' "$f")"
	done
}

# ---------------------------------------------------------------- uninstall

uninstall() {
	local confirm
	confirm=$(ask "Remove the VPN, all devices and keys from this server? (yes/no)" "no")
	[[ $confirm == yes ]] || { echo "Nothing removed."; return; }
	for i in wg0 wg1; do
		systemctl disable --now "wg-quick@$i" >/dev/null 2>&1 || true
	done
	rm -f /etc/unbound/unbound.conf.d/coolervpn.conf /etc/sysctl.d/99-coolervpn.conf
	systemctl restart unbound >/dev/null 2>&1 || true
	rm -rf "$CLIENT_DIR" "$WG_DIR"/wg0.conf "$WG_DIR"/wg1.conf "$WG_DIR"/rules-wg*.sh "$STATE" "$WG_DIR/entry-token.txt"
	green "Removed. The WireGuard and Unbound packages stay installed; remove them with: apt remove wireguard unbound"
}

# ---------------------------------------------------------------- menus

show_status() {
	load_state
	echo
	bold "coolervpn is installed: $(case $MODE in single) echo 'Single VPN';; entry) echo 'Double VPN, entry server';; exit) echo 'Double VPN, exit server';; esac)"
	wg show 2>/dev/null | grep -E 'interface|peer|latest handshake|transfer' | sed 's/^/  /' || true
}

menu_installed() {
	show_status
	echo
	if [[ $MODE == exit ]]; then
		echo "  1) Show the key for the entry server"
		echo "  2) Uninstall"
		echo "  3) Exit"
		case $(ask "Choose" "3") in
			1) cat "$WG_DIR/entry-token.txt" ;;
			2) uninstall ;;
		esac
		return
	fi
	echo "  1) Add a device"
	echo "  2) Remove a device"
	echo "  3) List devices"
	echo "  4) Show a device's QR code again"
	echo "  5) Uninstall"
	echo "  6) Exit"
	case $(ask "Choose" "1") in
		1) add_client ;;
		2) remove_client ;;
		3) list_clients ;;
		4) list_clients; local n; n=$(ask "Device" "phone")
		   [[ -f $CLIENT_DIR/$n.conf ]] && qrencode -t ansiutf8 < "$CLIENT_DIR/$n.conf" || red "no device called $n" ;;
		5) uninstall ;;
	esac
}

menu_new() {
	echo
	bold "coolervpn: your own WireGuard VPN"
	echo
	echo "  1) Single VPN    one server"
	echo "  2) Double VPN    two servers in two countries"
	echo
	case $(ask "Choose" "1") in
		1) MODE=single; install_client_server; green "VPN server is ready."; cloud_firewall_note "$PORT"; add_client ;;
		2)
			echo
			echo "Double VPN uses two servers. Set up the EXIT server first."
			echo "  1) This is the EXIT server   (the second one, traffic leaves the internet here)"
			echo "  2) This is the ENTRY server  (your phone connects here, needs the key from the exit server)"
			echo
			case $(ask "Choose" "1") in
				1) install_exit ;;
				2) install_entry; green "Entry server is ready."; cloud_firewall_note "$PORT"; add_client ;;
				*) die "choose 1 or 2" ;;
			esac ;;
		*) die "choose 1 or 2" ;;
	esac
}

# Refuse to touch a WireGuard setup we didn't make (for example Nyr/wireguard-install).
check_existing() {
	[[ -f $STATE ]] && return
	local f
	for f in "$WG_DIR"/wg0.conf "$WG_DIR"/wg1.conf; do
		if [[ -f $f ]]; then
			red "WireGuard is already set up on this server ($f)."
			echo "This script won't overwrite it. Remove the existing setup first"
			echo "(for Nyr's script: run it again and choose \"Remove WireGuard\"),"
			echo "or use a fresh Ubuntu server."
			exit 1
		fi
	done
}

main() {
	check_system
	case ${1:-} in
		--single)       [[ -f $STATE ]] && die "already installed"; check_existing; MODE=single; install_client_server; cloud_firewall_note "$PORT"; add_client "${2:-phone}" ;;
		--double-exit)  [[ -f $STATE ]] && die "already installed"; check_existing; install_exit ;;
		--double-entry) [[ -f $STATE ]] && die "already installed"; check_existing; install_entry "${2:-}"; cloud_firewall_note "$PORT"; add_client "${3:-phone}" ;;
		--add)          add_client "${2:-}" ;;
		--remove)       remove_client "${2:-}" ;;
		--list)         list_clients ;;
		--show-token)   cat "$WG_DIR/entry-token.txt" ;;
		--uninstall)    uninstall ;;
		"")             if [[ -f $STATE ]]; then menu_installed; else check_existing; menu_new; fi ;;
		*)              echo "usage: sudo bash $0 [--single | --double-exit | --double-entry TOKEN | --add NAME | --remove NAME | --list | --uninstall]"; exit 1 ;;
	esac
}

main "$@"