#!/bin/bash
set -e

WG_SERVER_HOME="{{ wg_server_home }}"
WG_PEERS_HOME="{{ wg_peers_home }}"
IP_FILE="{{ wg_ip_file }}"
WIREGUARD_SUBNET_PREFIX="{{ wg_subnet_prefix }}"
NAT_SUBNET="{{ nat_subnet }}"
DEFAULT_PORT="{{ wg_port }}"
DEFAULT_DNS="8.8.8.8"

test "${EUID}" -ne 0 && printf "%s\n" "run as root" && exit 1
umask 077

if ! command -v wg &>/dev/null; then
    printf "%s\n" "[err] wireguard not installed"
    exit 1
fi

function usage() {
    printf "%s\n" \
        ""                                                                      \
        "wireguard peer management script"                                      \
        "usage: $(basename ${0}) <action> [-h] <options>"                       \
        ""                                                                      \
        "add-peer           add a peer to the wg network"                       \
        "options:"                                                              \
        "-s <server>        required: wg server endpoint"                       \
        "-n <name>          optional: peer name, default random"                \
        "-p <port>          optional: wg server port, default ${DEFAULT_PORT}"  \
        "-d <dns>           optional: dns server, default ${DEFAULT_DNS}"       \
        ""                                                                      \
        "remove-peer        remove peer from the wg network"                    \
        "options:"                                                              \
        "-n <name>          required: peer name"                                \
        ""                                                                      \
        "usage              print this help message and exit"                   \
        ""                                                                      \
        "configuration files:"                                                  \
        "wg server config dir: ${WG_SERVER_HOME}"                               \
        "wg peers configs dir: ${WG_PEERS_HOME}"
    exit 1
}

function validate_ip() {
    local ip="${1}"
    if ! [[ "${ip}" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
        printf "%s\n" "[err] invalid IP address: ${ip}"
        exit 1
    fi
}

function validate_port() {
    local port="${1}"
    if ! [[ "${port}" =~ ^[0-9]+$ ]] || [[ "${port}" -lt 1 ]] || [[ "${port}" -gt 65535 ]]; then
        printf "%s\n" "[err] invalid port: ${port}"
        exit 1
    fi
}

function get_next_available_ip() {
    (
        flock -x 200
        touch "${IP_FILE}"
        for i in {2..254}; do
            ip="${WIREGUARD_SUBNET_PREFIX}.${i}"
            if ! grep -q "${ip}" "${IP_FILE}"; then
                printf "%s\n" "${ip}"
                printf "%s\n" "${ip}" >> "${IP_FILE}"
                exit 0
            fi
        done
        printf "%s\n" "[err] no available ips in range ${WIREGUARD_SUBNET_PREFIX}.2 - ${WIREGUARD_SUBNET_PREFIX}.254"
        exit 1
    ) 200>"${IP_FILE}.lock"
}

function add_peer() {
    if ! test -d "${WG_SERVER_HOME}" || ! test -f "${WG_SERVER_HOME}/wg0.conf"; then
        printf "%s\n" "[err] no wg server config found; install the server first"
        exit 1
    fi

    if ! test "${server}"; then
        printf "%s\n" "[err] missing -s <server>"
        exit 1
    fi

    mkdir -p "${WG_PEERS_HOME}/${name}" &>/dev/null
    assigned_ip=$(get_next_available_ip)
    (
        cd "${WG_PEERS_HOME}/${name}"
        wg genkey | tee "${name}.key" | wg pubkey > "${name}.pub"

        cat > "${name}.conf" << EOF
# peer ${name}
[Interface]
PrivateKey = $(cat "${name}.key")
Address    = ${assigned_ip}/24
DNS        = ${dns}

[Peer]
PublicKey           = $(wg pubkey < "${WG_SERVER_HOME}/server.key")
PresharedKey        = $(cat "${WG_SERVER_HOME}/psk.key")
Endpoint            = ${server}:${port}
AllowedIPs          = ${WIREGUARD_SUBNET_PREFIX}.0/24, ${NAT_SUBNET}
PersistentKeepalive = 25
EOF
        printf "%s\n" \
            "[inf] generated peer configuration ${name} to server ${server}" \
            "[inf] config: ${WG_PEERS_HOME}/${name}/${name}.conf"
    )

    (
        peer_public_key=$(wg pubkey < "${WG_PEERS_HOME}/${name}/${name}.key")
        cat >> "${WG_SERVER_HOME}/wg0.conf" << EOF

# peer ${name}
[Peer]
PublicKey    = ${peer_public_key}
PresharedKey = $(cat "${WG_SERVER_HOME}/psk.key")
AllowedIPs   = ${assigned_ip}/32
# peer ${name}
EOF
        systemctl restart wg-quick@wg0.service
        printf "%s\n" "[inf] restarted wg-quick@wg0.service"
    )
}

function remove_peer() {
    if ! test "${name}"; then
        printf "%s\n" "[err] missing -n <name>"
        exit 1
    fi

    if grep -qi "^# peer ${name}" "${WG_SERVER_HOME}/wg0.conf"; then
        sed -i "/# peer ${name}/,/# peer ${name}/ s/^/# /" "${WG_SERVER_HOME}/wg0.conf"
        systemctl restart wg-quick@wg0.service
        printf "%s\n" "[inf] removed peer ${name} and restarted wg server"
    else
        printf "%s\n" "[err] no such peer ${name} in ${WG_SERVER_HOME}/wg0.conf"
        exit 1
    fi
}

port="${DEFAULT_PORT}"
dns="${DEFAULT_DNS}"

if test "${1}"; then
    action="${1}"
    shift
    case "${action}" in
        add-peer)
            while getopts "s:n:p:d:h" opts; do
                case "${opts}" in
                    s) server="${OPTARG}"; validate_ip "${server}";;
                    n) name="${OPTARG}";;
                    p) port="${OPTARG}"; validate_port "${port}";;
                    d) dns="${OPTARG}"; validate_ip "${dns}";;
                    h) usage;;
                    *) usage;;
                esac
            done
            rand=$(printf "%s\n" "${RANDOM}" | md5sum | fold -w4 | head -1)
            name="${name:-${rand}}"
            add_peer
            ;;
        remove-peer)
            while getopts "n:h" opts; do
                case "${opts}" in
                    n) name="${OPTARG}";;
                    h) usage;;
                    *) usage;;
                esac
            done
            remove_peer
            ;;
        usage)
            usage
            ;;
        *)
            usage
            ;;
    esac
else
    usage
fi