#!/bin/bash -e

APP=$(basename $0)
LOCKFILE="/tmp/$APP.lock"

trap "rm -f ${LOCKFILE}; exit" INT TERM EXIT
if ! ln -s $APP $LOCKFILE 2>/dev/null; then
    echo "ERROR: script LOCKED" >&2
    exit 15
fi

function usage {
  echo "Usage: $0 [<options>] [command [arg]]"
  echo "Options:"
  echo " -i : Init (Create server keys and configs)"
  echo " -c : Create new user"
  echo " -d : Delete user"
  echo " -L : Lock user"
  echo " -U : Unlock user"
  echo " -p : Print user config"
  echo " -q : Print user QR code"
  echo " -u <user> : User identifier (uniq field for vpn account)"
  echo " -s <server> : Server host for user connection"
  echo " -I : Interface (default auto)"
  echo " -h : Usage"
  exit 1
}

unset USER
umask 0077

HOME_DIR="/etc/amnezia/amneziawg"
SERVER_NAME="awg0"
SERVER_IP_PREFIX="10.20.30"
SERVER_PORT=39547
SERVER_INTERFACE=$(ip -4 route ls | grep default | grep -Po '(?<=dev )(\S+)' | head -1)

while getopts ":icdpqhLUu:I:s:" opt; do
  case $opt in
     i) INIT=1 ;;
     c) CREATE=1 ;;
     d) DELETE=1 ;;
     L) LOCK=1 ;;
     U) UNLOCK=1 ;;
     p) PRINT_USER_CONFIG=1 ;;
     q) PRINT_QR_CODE=1 ;;
     u) USER="$OPTARG" ;;
     I) SERVER_INTERFACE="$OPTARG" ;;
     h) usage ;;
     s) SERVER_ENDPOINT="$OPTARG" ;;
    \?) echo "Invalid option: -$OPTARG" ; exit 1 ;;
     :) echo "Option -$OPTARG requires an argument" ; exit 1 ;;
  esac
done

[ $# -lt 1 ] && usage

function reload_server {
    awg syncconf ${SERVER_NAME} <(awg-quick strip ${SERVER_NAME})
}

function get_new_ip {
    declare -A IP_EXISTS

    for IP in $(grep -i 'Address\s*=\s*' keys/*/*.conf | sed 's/\/[0-9]\+$//' | grep -Po '\d+$')
    do
        IP_EXISTS[$IP]=1
    done

    for IP in {2..255}
    do
        [ ${IP_EXISTS[$IP]} ] || break
    done

    if [ $IP -eq 255 ]; then
        echo "ERROR: can't determine new address" >&2
        exit 3
    fi

    echo "${SERVER_IP_PREFIX}.${IP}/32"
}

function encode {
    python3 encode.py ${USER} > keys/${USER}/${USER}.vpn
}


function add_user_to_server {
    if [ ! -f "keys/${USER}/public.key" ]; then
        echo "ERROR: User not exists" >&2
        exit 1
    fi

    local USER_PUB_KEY=$(cat "keys/${USER}/public.key")
    local USER_PSK_KEY=$(cat "keys/$USER/psk.key")
    local USER_IP=$(grep -i Address "keys/${USER}/${USER}.conf" | sed 's/Address\s*=\s*//i; s/\/.*//')

    if grep "# BEGIN ${USER}$" "$SERVER_NAME.conf" >/dev/null ; then
        echo "User already exists"
        exit 0
    fi

cat <<EOF >> "$SERVER_NAME.conf"
# BEGIN ${USER}
[Peer]
PublicKey = ${USER_PUB_KEY}
AllowedIPs = ${USER_IP}
PresharedKey = ${USER_PSK_KEY}
# END ${USER}
EOF

    ip -4 route add ${USER_IP}/32 dev ${SERVER_NAME} || true
}

function remove_user_from_server {
    sed -i "/# BEGIN ${USER}$/,/# END ${USER}$/d" "$SERVER_NAME.conf"
    if [ -f "keys/${USER}/${USER}.conf" ]; then
        local USER_IP=$(grep -i Address "keys/${USER}/${USER}.conf" | sed 's/Address\s*=\s*//i; s/\/.*//')
        ip -4 route del ${USER_IP}/32 dev ${SERVER_NAME} || true
    fi
}

function init {
    if [ -z "$SERVER_ENDPOINT" ]; then
        echo "ERROR: Server required" >&2
        exit 1
    fi

    if [ -z "$SERVER_INTERFACE" ]; then
        echo "ERROR: Can't determine server interface" >&2
        echo "DEBUG: 'ip route':"
        ip route
        exit 1
    fi

    echo "Interface: $SERVER_INTERFACE"

    mkdir -p "keys/${SERVER_NAME}"
    echo -n "$SERVER_ENDPOINT" > "keys/.server"

    if [ ! -f "keys/${SERVER_NAME}/private.key" ]; then
        awg genkey | tee "keys/${SERVER_NAME}/private.key" | awg pubkey > "keys/${SERVER_NAME}/public.key"
    fi

    if [ -f "$SERVER_NAME.conf" ]; then
        echo "Server already initialized"
        exit 0
    fi

    SERVER_PVT_KEY=$(cat "keys/$SERVER_NAME/private.key")

cat <<EOF > "$SERVER_NAME.conf"
[Interface]
Address = ${SERVER_IP_PREFIX}.1/32
ListenPort = ${SERVER_PORT}
PrivateKey = ${SERVER_PVT_KEY}
PostUp = iptables -t nat -A POSTROUTING -o ${SERVER_INTERFACE} -j MASQUERADE
PostDown = iptables -t nat -D POSTROUTING -o ${SERVER_INTERFACE} -j MASQUERADE
Jc = 5
Jmin = 50
Jmax = 1000
S1 = 147
S2 = 57
H1 = 1121994835
H2 = 1702292146
H3 = 1975368295
H4 = 1948088518

EOF

    echo net.ipv4.ip_forward=1 >> /etc/sysctl.conf
    sysctl -p

    systemctl enable awg-quick@${SERVER_NAME}
    awg-quick up ${SERVER_NAME} || true

    echo "Server initialized successfully"
    exit 0
}

function create {
    if [ -f "keys/${USER}/${USER}.conf" ]; then
        echo "WARNING: key ${USER}.conf already exists" >&2
        return 0
    fi

    SERVER_ENDPOINT=$(cat "keys/.server")
    USER_IP=$( get_new_ip )

    mkdir "keys/${USER}"
    awg genkey | tee "keys/${USER}/private.key" | awg pubkey > "keys/${USER}/public.key" | awg genpsk > "keys/${USER}/psk.key"

    USER_PVT_KEY=$(cat "keys/${USER}/private.key")
    USER_PUB_KEY=$(cat "keys/${USER}/public.key")
    USER_PSK_KEY=$(cat "keys/${USER}/psk.key")
    SERVER_PUB_KEY=$(cat "keys/$SERVER_NAME/public.key")

cat <<EOF > "keys/${USER}/${USER}.conf"
[Interface]
PrivateKey = ${USER_PVT_KEY}
Address = ${USER_IP}
DNS = 1.1.1.1, 8.8.8.8
Jc = 5
Jmin = 50
Jmax = 1000
S1 = 147
S2 = 57
H1 = 1121994835
H2 = 1702292146
H3 = 1975368295
H4 = 1948088518

[Peer]
PublicKey = ${SERVER_PUB_KEY}
Endpoint = ${SERVER_ENDPOINT}:${SERVER_PORT}
AllowedIPs = 8.8.4.0/24, 8.8.8.0/24, 8.34.0.0/16, 8.35.0.0/16, 23.236.0.0/16, 23.251.0.0/16, 31.13.0.0/16, 34.0.0.0/10, 34.64.0.0/10, 35.184.0.0/13, 35.192.0.0/14, 35.196.0.0/15, 35.198.0.0/16, 35.199.0.0/17, 35.200.0.0/13, 35.208.0.0/12, 45.64.0.0/16, 46.19.0.0/16, 57.144.0.0/16, 64.18.0.0/16, 64.233.0.0/16, 66.102.0.0/16, 66.220.0.0/16, 66.249.0.0/16, 68.66.0.0/16, 70.32.0.0/16, 72.14.0.0/16, 72.52.0.0/16, 74.114.0.0/16, 74.119.0.0/16, 74.125.0.0/16, 81.29.0.0/16, 83.174.0.0/16, 84.15.0.0/16, 87.245.0.0/16, 91.213.0.0/16, 92.46.0.0/16, 94.140.0.0/16, 103.4.0.0/16, 103.4.96.0/22, 104.18.0.0/16, 104.21.0.0/16, 104.31.0.0/16, 104.32.0.0/16, 104.132.0.0/16, 104.133.0.0/16, 104.134.0.0/15, 104.156.0.0/16, 104.237.0.0/16, 108.170.0.0/16, 108.177.0.0/15, 108.177.0.0/16, 108.59.0.0/16, 130.211.0.0/16, 136.112.0.0/12, 142.250.0.0/16, 142.251.0.0/16, 146.148.0.0/16, 157.240.0.0/16, 162.159.0.0/16, 162.216.0.0/16, 162.222.0.0/16, 172.110.0.0/16, 172.217.0.0/16, 172.253.0.0/16, 172.67.0.0/16, 173.194.0.0/16, 173.252.0.0/16, 185.38.0.0/16, 185.53.0.0/16, 185.60.0.0/16, 185.89.0.0/16, 188.114.0.0/16, 190.115.0.0/16, 192.158.0.0/16, 192.178.0.0/15, 193.46.0.0/16, 193.186.0.0/16, 199.36.0.0/16, 199.192.0.0/16, 199.223.0.0/16, 204.15.0.0/16, 207.223.0.0/16, 208.65.0.0/16, 208.68.0.0/16, 208.81.0.0/16, 208.117.0.0/16, 209.85.0.0/16, 209.85.128.0/17, 216.40.0.0/16, 216.58.0.0/16, 216.58.192.0/18, 216.239.0.0/16, 216.239.32.0/19, 217.168.0.0/16, 124.108.0.0/16, 129.134.0.0/16, 130.211.0.0/16, 136.112.0.0/12, 146.148.0.0/16, 162.159.0.0/16, 162.216.0.0/16, 162.222.0.0/16, 173.255.0.0/16, 185.214.0.0/16, 185.38.0.0/16, 185.53.0.0/16, 185.60.0.0/16, 185.89.0.0/16, 190.115.0.0/16, 195.123.0.0/16, 194.110.0.0/16, 176.97.0.0/16, 5.45.0.0/16, 193.46.0.0/16, 68.66.0.0/16, 92.46.0.0/16, 124.108.0.0/16
PersistentKeepalive = 20
PresharedKey = ${USER_PSK_KEY}
EOF
    encode
    add_user_to_server
    reload_server
}

cd $HOME_DIR

if [ $INIT ]; then
    init
    exit 0;
fi

if [ ! -f "keys/$SERVER_NAME/public.key" ]; then
    echo "ERROR: Run init script before" >&2
    exit 2
fi

if [ -z "${USER}" ]; then
    echo "ERROR: User required" >&2
    exit 1
fi

if [ $CREATE ]; then
    create
fi

if [ $DELETE ]; then
    remove_user_from_server
    reload_server
    rm -rf "keys/${USER}"
    exit 0
fi

if [ $LOCK ]; then
    remove_user_from_server
    reload_server
    exit 0
fi

if [ $UNLOCK ]; then
    add_user_to_server
    reload_server
    exit 0
fi

if [ $PRINT_USER_CONFIG ]; then
    cat "keys/${USER}/${USER}.conf"
elif [ $PRINT_QR_CODE ]; then
    qrencode -t ansiutf8 < "keys/${USER}/${USER}.vpn"
fi

exit 0