#!/bin/bash
set -e

exec 3>&1
export WG_HIDE_KEYS=never
netns1="wg-test-$$-1"
netns2="wg-test-$$-2"
pretty() { echo -e "\x1b[32m\x1b[1m[+] ${1:+NS$1: }${2}\x1b[0m" >&3; }
pp() { pretty "" "$*"; "$@"; }
maybe_exec() { if [[ $BASHPID -eq $$ ]]; then "$@"; else exec "$@"; fi; }
n1() { pretty 1 "$*"; maybe_exec ip netns exec $netns1 "$@"; }
n2() { pretty 2 "$*"; maybe_exec ip netns exec $netns2 "$@"; }
ip1() { pretty 1 "ip $*"; ip -n $netns1 "$@"; }
ip2() { pretty 2 "ip $*"; ip -n $netns2 "$@"; }
sleep() { read -t "$1" -N 0 || true; }
waitiface() { pretty "${1//*-}" "wait for $2 to come up"; ip netns exec "$1" bash -c "while [[ \$(< \"/sys/class/net/$2/operstate\") != up ]]; do read -t .1 -N 0 || true; done;"; }

cleanup() {
        set +e
        exec 2>/dev/null
        ip1 link del dev wg0
        ip2 link del dev wg0
        local to_kill="$(ip netns pids $netns1) $(ip netns pids $netns2)"
        [[ -n $to_kill ]] && kill $to_kill
        pp ip netns del $netns1
        pp ip netns del $netns2
        exit
}

trap cleanup EXIT

ip netns del $netns1 2>/dev/null || true
ip netns del $netns2 2>/dev/null || true
pp ip netns add $netns1
pp ip netns add $netns2

key1="$(pp wg genkey)"
key2="$(pp wg genkey)"
pub1="$(pp wg pubkey <<<"$key1")"
pub2="$(pp wg pubkey <<<"$key2")"
psk="$(pp wg genpsk)"
[[ -n $key1 && -n $key2 && -n $psk ]]

configure_peers() {
        ip1 addr add 192.168.241.1/24 dev wg0
        ip2 addr add 192.168.241.2/24 dev wg0

        n1 wg set wg0 \
                private-key <(echo "$key1") \
                listen-port 1 \
                peer "$pub2" \
                        preshared-key <(echo "$psk") \
                        allowed-ips 192.168.241.2/32,fd00::2/128
        n2 wg set wg0 \
                private-key <(echo "$key2") \
                listen-port 2 \
                peer "$pub1" \
                        preshared-key <(echo "$psk") \
                        allowed-ips 192.168.241.1/32,fd00::1/128

        ip1 link set up dev wg0
        ip2 link set up dev wg0
}

n1 bash -c 'echo 1 > /proc/sys/net/ipv6/conf/all/disable_ipv6'
n2 bash -c 'echo 1 > /proc/sys/net/ipv6/conf/all/disable_ipv6'
n1 bash -c 'echo 1 > /proc/sys/net/ipv6/conf/default/disable_ipv6'
n2 bash -c 'echo 1 > /proc/sys/net/ipv6/conf/default/disable_ipv6'

ip1 link add dev wg0 type wireguard
ip2 link add dev wg0 type wireguard
configure_peers

ip1 link add dummy0 type dummy
ip1 addr add 10.0.0.10/24 dev dummy0
ip1 link set dummy0 up

ip1 link add veth1 type veth peer name veth2
ip1 link set veth2 netns $netns2

ip1 addr add 10.0.0.1/24 dev veth1
ip1 addr add 10.0.0.2/24 dev veth1
ip2 addr add 10.0.0.3/24 dev veth2

ip1 link set veth1 up
ip2 link set veth2 up
waitiface $netns1 veth1
waitiface $netns2 veth2

n1 iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
n2 iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

n2 wg set wg0 peer "$pub1" endpoint 10.0.0.1:1
n2 ping -W 1 -c 5 -f 192.168.241.1
[[ $(n2 wg show wg0 endpoints) == "$pub1        10.0.0.1:1" ]]

n1 conntrack -L
n2 conntrack -L

n2 wg set wg0 peer "$pub1" endpoint 10.0.0.2:1
n2 ping -W 1 -c 5 -f 192.168.241.1
[[ $(n2 wg show wg0 endpoints) == "$pub1        10.0.0.2:1" ]]

n1 conntrack -L
n2 conntrack -L
