#!/bin/sh
#
# wg-connect - bring a WireGuard tunnel up/down on BusyBox systems.
# Usage: wg-connect up   [config]   (defaults to ./peer.conf)
#        wg-connect down
#
# Requires: wg, ip, grep, sed, cut (all typically available on BusyBox).

set -e

STATE_FILE="/tmp/wg-connect.state"
IFACE_DEFAULT="wg0"

usage() {
    echo "Usage: wg-connect up <config> | wg-connect down [name]" >&2
    echo "  up <config>   config name (looked up in /etc/wireguard/<name>.conf)" >&2
    echo "  up <path>     full or relative path to a .conf file" >&2
    echo "  down [name]   tear down a tunnel (default: wg0)" >&2
    exit 1
}

die() {
    echo "wg-connect: $*" >&2
    exit 1
}

# Resolve a config argument to a path:
#   - empty       → usage error
#   - contains /  → use as-is (already a path)
#   - *.conf      → try cwd, then /etc/wireguard/
#   - plain name  → /etc/wireguard/<name>.conf
resolve_config() {
    _name="$1"

    [ -z "$_name" ] && { usage; exit 1; }

    # Full or relative path given
    case "$_name" in
        */* ) echo "$_name"; return 0 ;;
    esac

    # Explicit .conf filename - try cwd first, then /etc/wireguard
    case "$_name" in
        *.conf )
            [ -f "$_name" ] && { echo "$_name"; return 0; }
            [ -f "/etc/wireguard/$_name" ] && { echo "/etc/wireguard/$_name"; return 0; }
            die "config not found: $_name (tried . and /etc/wireguard/)"
            ;;
    esac

    # Bare name → /etc/wireguard/<name>.conf
    _path="/etc/wireguard/${_name}.conf"
    [ -f "$_path" ] || die "config not found: $_path"
    echo "$_path"
}

# -- config parser -------------------------------------------------------
# Reads a WireGuard .conf file and sets global vars:
#   IFACE, ADDRESS, DNS, LISTEN_PORT, PRIVATE_KEY
# Each Peer is stored as positional-ish vars: we only extract what we need
# for route/DNS handling, since wg setconf reads the raw config directly.

parse_config() {
    _cfg="$1"

    [ -f "$_cfg" ] || die "config file not found: $_cfg"

    # Reset globals
    ADDRESS=""
    DNS=""
    LISTEN_PORT=""
    PRIVATE_KEY=""
    ALLOWED_IPS=""
    ENDPOINTS=""

    while IFS= read -r line || [ -n "$line" ]; do
        # Strip comments and leading/trailing whitespace
        line="${line%%#*}"
        line="${line# }"
        line="${line% }"
        [ -z "$line" ] && continue

        case "$line" in
            \[*\] )
                _section="${line#\[}"
                _section="${_section%\]}"
                ;;
            Address\ *=\ * )
                ADDRESS="${line#*= }"
                ADDRESS="${ADDRESS# }"
                ;;
            DNS\ *=\ * )
                DNS="${line#*= }"
                DNS="${DNS# }"
                ;;
            ListenPort\ *=\ * )
                LISTEN_PORT="${line#*= }"
                LISTEN_PORT="${LISTEN_PORT# }"
                ;;
            PrivateKey\ *=\ * )
                PRIVATE_KEY="${line#*= }"
                PRIVATE_KEY="${PRIVATE_KEY# }"
                ;;
            AllowedIPs\ *=\ * )
                _ips="${line#*= }"
                _ips="${_ips# }"
                if [ -z "$ALLOWED_IPS" ]; then
                    ALLOWED_IPS="$_ips"
                else
                    ALLOWED_IPS="$ALLOWED_IPS, $_ips"
                fi
                ;;
            Endpoint\ *=\ * )
                _ep="${line#*= }"
                _ep="${_ep# }"
                _ep="${_ep%:*}"     # strip port, keep only IP
                if [ -z "$ENDPOINTS" ]; then
                    ENDPOINTS="$_ep"
                else
                    ENDPOINTS="$ENDPOINTS $_ep"
                fi
                ;;
        esac
    done < "$_cfg"
}

# -- helpers -------------------------------------------------------------

is_iface_up() {
    ip link show "$1" >/dev/null 2>&1
}

save_default_route() {
    # Returns the current IPv4 default route line, if any
    ip route show default 2>/dev/null || true
}

add_routes() {
    # Add routes for AllowedIPs. Handles IPv4 only for BusyBox compat.
    # 0.0.0.0/0 is handled specially: it replaces the default route.
    _iface="$1"
    _ips="$2"
    _old_IFS="$IFS"

    IFS=","
    for _net in $_ips; do
        _net="${_net# }"
        _net="${_net% }"
        [ -z "$_net" ] && continue

        # Skip IPv6 - BusyBox ip may not support it
        case "$_net" in
            *:* ) continue ;;
        esac

        if [ "$_net" = "0.0.0.0/0" ] || [ "$_net" = "0.0.0.0" ]; then
            ip route replace default dev "$_iface" 2>/dev/null || \
            ip route add default dev "$_iface" 2>/dev/null || true
        else
            ip route add "$_net" dev "$_iface" 2>/dev/null || true
        fi
    done
    IFS="$_old_IFS"
}

setup_dns() {
    _dns="$1"
    [ -z "$_dns" ] && return 0

    cp /etc/resolv.conf "$DNS_BACKUP"

    # Write minimal resolv.conf
    {
        echo "nameserver $_dns"
    } > /etc/resolv.conf
}

restore_dns() {
    if [ -f "$DNS_BACKUP" ]; then
        cp "$DNS_BACKUP" /etc/resolv.conf
        rm -f "$DNS_BACKUP"
    fi
}

# -- commands ------------------------------------------------------------

cmd_up() {
    _cfg="$(resolve_config "${1:-}")"

    parse_config "$_cfg"

    # Interface name comes from config filename (matching wg-quick behaviour)
    _iface_name="${_cfg##*/}"          # strip path
    _iface_name="${_iface_name%.conf}" # strip .conf
    IFACE="${IFACE:-$_iface_name}"

    STATE_FILE="/tmp/wg-connect.${IFACE}.state"

    if is_iface_up "$IFACE"; then
        die "interface $IFACE already exists - run 'wg-connect down $IFACE' first"
    fi

    # Check for port conflicts with any existing WireGuard interface
    if [ -n "$LISTEN_PORT" ]; then
        _used_ports="$(wg show interfaces 2>/dev/null | while read -r _wgif; do
            wg show "$_wgif" listen-port 2>/dev/null
        done)"
        for _p in $_used_ports; do
            [ "$_p" = "$LISTEN_PORT" ] && die "port $LISTEN_PORT is already in use by another WireGuard interface"
        done
    fi

    # Trap to roll back partial setup on failure.
    _up_done=""
    trap 'rollback' INT TERM EXIT

    rollback() {
        trap - INT TERM EXIT
        if [ -n "$_up_done" ]; then
            return 0
        fi
        # Undo whatever was created, in reverse order.
        restore_dns 2>/dev/null || true
        is_iface_up "$IFACE" && ip link del "$IFACE" 2>/dev/null || true
        for _ep in $ENDPOINT_ROUTES; do
            [ -z "$_ep" ] && continue
            ip route del "$_ep" 2>/dev/null || true
        done
        rm -f "$STATE_FILE"
    }

    # Save current default route and extract its gateway
    DEFAULT_ROUTE="$(save_default_route)"
    _gw=""
    case "$DEFAULT_ROUTE" in
        *\ via\ * ) _gw="${DEFAULT_ROUTE##* via }"; _gw="${_gw%% *}" ;;
    esac

    # Add explicit routes for WireGuard endpoints through the physical
    # gateway BEFORE we replace the default route, otherwise the encrypted
    # UDP packets have no path to the server.
    ENDPOINT_ROUTES=""
    for _ep in $ENDPOINTS; do
        [ -z "$_ep" ] && continue
        case "$_ep" in
            *:* ) continue ;;   # skip IPv6 endpoints
        esac
        if [ -n "$_gw" ]; then
            ip route add "$_ep" via "$_gw" 2>/dev/null || true
        fi
        ENDPOINT_ROUTES="${ENDPOINT_ROUTES}${_ep} "
    done

    # Create the interface
    ip link add "$IFACE" type wireguard

    # Configure WireGuard - strip wg-quick-only fields (Address, DNS, etc.)
    # that wg setconf rejects. These are handled manually below.
    _wg_conf="/tmp/wg-connect.$$.conf"
    grep -v -E '^[[:space:]]*(Address|DNS|MTU|Table|PreUp|PostUp|PreDown|PostDown|SaveConfig)[[:space:]]*=' "$_cfg" > "$_wg_conf"
    wg setconf "$IFACE" "$_wg_conf"
    rm -f "$_wg_conf"

    # Assign address (append /32 if no CIDR given)
    case "$ADDRESS" in
        */* ) ;;
        * ) ADDRESS="${ADDRESS}/32" ;;
    esac
    ip addr add "$ADDRESS" dev "$IFACE"

    # Bring interface up
    ip link set "$IFACE" up

    # Add routes (including default route replacement)
    add_routes "$IFACE" "$ALLOWED_IPS"

    # Set up DNS
    DNS_BACKUP="/tmp/resolv.conf.wg.bak"
    setup_dns "$DNS"

    # Write state file for cmd_down
    {
        echo "IFACE=$IFACE"
        echo "DNS_BACKUP=$DNS_BACKUP"
        echo "DEFAULT_ROUTE=$DEFAULT_ROUTE"
        echo "ENDPOINT_ROUTES=$ENDPOINT_ROUTES"
    } > "$STATE_FILE"

    # Mark success - disable rollback
    _up_done=1
    trap - INT TERM EXIT

    echo "wg-connect: $IFACE is up"
}

cmd_down() {
    # Resolve state file: explicit name, or default "peer" (from peer.conf)
    _down_name="${1:-}"

    if [ -n "$_down_name" ]; then
        case "$_down_name" in
            *.conf ) _down_name="${_down_name%.conf}" ;;
        esac
        case "$_down_name" in
            */* ) _down_name="${_down_name##*/}" ;;
        esac
        STATE_FILE="/tmp/wg-connect.${_down_name}.state"
        IFACE="$_down_name"
    fi

    if [ -f "$STATE_FILE" ]; then
        # Read state
        while IFS='=' read -r key value; do
            case "$key" in
                IFACE) IFACE="$value" ;;
                DNS_BACKUP) DNS_BACKUP="$value" ;;
                DEFAULT_ROUTE) DEFAULT_ROUTE="$value" ;;
                ENDPOINT_ROUTES) ENDPOINT_ROUTES="$value" ;;
            esac
        done < "$STATE_FILE"

        # Restore DNS
        restore_dns

        # Remove endpoint routes
        for _ep in $ENDPOINT_ROUTES; do
            [ -z "$_ep" ] && continue
            ip route del "$_ep" 2>/dev/null || true
        done
    else
        # No state file - try to clean up a partially-created interface
        if [ -z "$IFACE" ]; then
            IFACE="$IFACE_DEFAULT"
        fi
        if ! is_iface_up "$IFACE"; then
            die "no state file found and $IFACE does not exist - nothing to tear down"
        fi
        echo "wg-connect: no state file, cleaning up leftover $IFACE" >&2
    fi

    # Delete the interface (this also removes its routes and address)
    if is_iface_up "$IFACE"; then
        ip link del "$IFACE"
    fi

    # Restore the previous default route if there was one
    if [ -n "$DEFAULT_ROUTE" ]; then
        ip route add $DEFAULT_ROUTE 2>/dev/null || true
    fi

    rm -f "$STATE_FILE"

    echo "wg-connect: $IFACE is down"
}

# -- main ----------------------------------------------------------------

case "${1:-}" in
    up)   shift; cmd_up "$@" ;;
    down) shift; cmd_down "$@" ;;
    *)    usage ;;
esac
