#!/usr/bin/env zsh

if [[ "${1}" == '-h' ]] || [[ "${1}" == '--help' ]]; then
    printf 'usage: %s [-R|-D]\n' "${0}"
    printf '  -R|-D: remove rules currently in place (default is to add new rules)\n'
    exit
fi

setopt pipe_fail

local save_file="${HOME}/.cache/qBittorrent-iptables-save"
local op='-A'
if [[ "${1}" == '-D' || "${1}" == '-R' ]]; then
    if ! [[ -f "${save_file}" ]]; then
        echo 'No current rules found!'
        exit 1
    fi
    op="-D"
fi
local iface="$(ip route | grep '^default' | grep -Po '(?<=dev )[^ ]+')"
printf 'Using interface: %s\n' "${iface}"

# <-A|-D> <ex_ip> <in_ip>
function do_rules {
    emulate -L zsh
    PS4='Run: '
    setopt errexit xtrace
    doas iptables -t filter "${1}" FORWARD -i "${iface}" -o wg0-mullvad -j ACCEPT

    doas iptables -t nat "${1}" PREROUTING -d "${2}"/32 -p tcp -m tcp --dport 62000 \
         -j DNAT --to-destination "${3}":62000
    doas iptables -t nat "${1}" PREROUTING -d "${2}"/32 -p udp -m udp --dport 62000 \
         -j DNAT --to-destination "${3}":62000

    doas iptables -t nat "${1}" POSTROUTING -d "${3}"/32 -p tcp -m tcp --sport 62000 \
         -j SNAT --to-source "${2}":62000
    doas iptables -t nat "${1}" POSTROUTING -d "${3}"/32 -p udp -m udp --sport 62000 \
         -j SNAT --to-source "${2}":62000
}

if [[ -f "${save_file}" ]]; then
    local content="$(<"${save_file}")"
    local lines=("${(@f)content}")
    printf 'Old rules found for\nex_ip: %s\nin_ip: %s\n' "${lines[1]}" "${lines[2]}"
    printf 'Removing...\n'
    do_rules -D "${lines[1]}" "${lines[2]}"
    rm -f "${save_file}"
    printf 'Done!\n'
    [[ "${op}" == '-D' ]] && exit
fi

local ex_ip in_ip
ex_ip="$(curl -4 icanhazip.com)" || { echo 'Could not fetch ip!'; exit 1 }
in_ip="$(ip addr show dev wg0-mullvad | \
         awk '/^ *inet [0-9]+/ { print substr($2,0,index($2,"/") - 1) }')" ||
    { echo 'Could not find wireguard iterface address!'; exit 1 }
printf 'Adding rules for:\nex_ip: %s\nin_ip: %s\n' "${ex_ip}" "${in_ip}"
printf '%s\n%s\n' "${ex_ip}" "${in_ip}" >"${save_file}"

do_rules -A "${ex_ip}" "${in_ip}"