#!/usr/bin/env bash
set -Eeuo pipefail
IFS=$'\n\t'
export LC_ALL=C
umask 077

readonly TABLE_FAMILY="inet"
readonly TABLE_NAME="ochenstarik_smm"
readonly CHAIN_NAME="links"

fail() { printf '%s\n' "policy helper: $*" >&2; exit 78; }
firewall_unavailable() { printf '%s\n' "mesh.firewall-unavailable" >&2; exit 79; }

testing="${SMM_POLICY_TESTING:-0}"
if [[ "$testing" != "1" ]]; then
    [[ ${EUID:-$(id -u)} -eq 0 ]] || fail "root is required"
    readonly STATE_FILE="/var/lib/ochenstarik-server-monitor-manager/mesh/nodes.tsv"
    readonly RECONCILE_MARKER="/var/lib/ochenstarik-server-monitor-manager/mesh/reconcile-requested"
    readonly RECONCILE_LOCK="/var/lib/ochenstarik-server-monitor-manager/mesh/reconcile-requested.lock"
    readonly FLOCK_COMMAND="/usr/bin/flock"
else
    readonly STATE_FILE="${SMM_POLICY_STATE_FILE:-/dev/null}"
    readonly RECONCILE_MARKER="${SMM_POLICY_RECONCILE_MARKER:-${TMPDIR:-/tmp}/smm-policy-reconcile-requested}"
    readonly RECONCILE_LOCK="${SMM_POLICY_RECONCILE_LOCK:-${RECONCILE_MARKER}.lock}"
    readonly FLOCK_COMMAND="${SMM_POLICY_FLOCK:-/usr/bin/flock}"
fi

node_pattern='^[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?$'
ipv4_pattern='^([0-9]{1,3}\.){3}[0-9]{1,3}$'
generation_pattern='^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$'

validate_node_id() {
    [[ "$1" =~ $node_pattern ]] || fail "invalid node id"
}

lookup_node_ip() {
    local node_id="$1" ip
    [[ -r "$STATE_FILE" ]] || fail "mesh node state is unavailable"
    ip="$(awk -F '\t' -v node="$node_id" '$1 == node { print $2; exit }' "$STATE_FILE")"
    [[ "$ip" =~ $ipv4_pattern ]] || fail "node has no valid mesh address: $node_id"
    printf '%s\n' "$ip"
}

validate_rule() {
    local action="$1"
    case "$action" in
        link-connect) [[ $# -eq 6 ]] || fail "invalid link-connect argument count" ;;
        link-disconnect) [[ $# -eq 5 ]] || fail "invalid link-disconnect argument count" ;;
        link-status) [[ $# -eq 5 ]] || fail "invalid link-status argument count" ;;
        *) fail "unsupported action" ;;
    esac
    validate_node_id "$2"
    validate_node_id "$3"
    [[ "$2" != "$3" ]] || fail "source and destination must differ"
    [[ "$4" == "tcp" || "$4" == "udp" ]] || fail "invalid protocol"
    [[ "$5" =~ ^[0-9]+$ ]] && (( 10#$5 >= 1 && 10#$5 <= 65535 )) || fail "invalid port"
    if [[ "$action" == "link-connect" ]]; then
        [[ "$6" =~ ^[0-9]+$ ]] && (( 10#$6 >= 0 && 10#$6 <= 525600 )) || fail "invalid TTL"
    fi
}

run_nft() {
    if [[ "$testing" == "1" ]]; then
        printf 'nft'
        printf ' %q' "$@"
        printf '\n'
    else
        /usr/sbin/nft "$@"
    fi
}

inspect_firewall() {
    local listing
    if [[ "$testing" == "1" ]]; then
        [[ "${SMM_POLICY_FIREWALL_UNAVAILABLE:-0}" != "1" ]] || firewall_unavailable
        [[ -z "${SMM_POLICY_FIREWALL_ERROR:-}" ]] \
            || fail "could not inspect nftables Link policy: $SMM_POLICY_FIREWALL_ERROR"
        return 0
    fi
    if listing="$(/usr/sbin/nft -a list chain "$TABLE_FAMILY" "$TABLE_NAME" "$CHAIN_NAME" 2>&1)"; then
        printf '%s\n' "$listing"
        return 0
    fi
    if grep -Eiq 'No such file or directory|does not exist' <<<"$listing"; then
        firewall_unavailable
    fi
    fail "could not inspect nftables Link policy: ${listing%%$'\n'*}"
}

ensure_firewall_available() {
    inspect_firewall >/dev/null
}

rule_exists() {
    local comment="$1" listing
    listing="$(inspect_firewall)"
    grep -Fq "comment \"$comment\"" <<<"$listing"
}

connect_rule() {
    local source_id="$1" target_id="$2" protocol="$3" port="$4"
    local source_ip target_ip comment
    ensure_firewall_available
    source_ip="$(lookup_node_ip "$source_id")"
    target_ip="$(lookup_node_ip "$target_id")"
    comment="smm:${source_id}:${target_id}:${protocol}:${port}"
    rule_exists "$comment" && return 0
    run_nft add rule "$TABLE_FAMILY" "$TABLE_NAME" "$CHAIN_NAME" \
        ip saddr "$source_ip" ip daddr "$target_ip" "$protocol" dport "$port" \
        counter accept comment "$comment"
}

disconnect_rule() {
    local source_id="$1" target_id="$2" protocol="$3" port="$4" comment handle
    ensure_firewall_available
    # Resolve both identities before touching firewall state.
    lookup_node_ip "$source_id" >/dev/null
    lookup_node_ip "$target_id" >/dev/null
    comment="smm:${source_id}:${target_id}:${protocol}:${port}"
    if [[ "$testing" == "1" ]]; then
        printf 'nft-delete-comment %q\n' "$comment"
        return
    fi
    while IFS= read -r handle; do
        [[ "$handle" =~ ^[0-9]+$ ]] || continue
        run_nft delete rule "$TABLE_FAMILY" "$TABLE_NAME" "$CHAIN_NAME" handle "$handle"
    done < <(
        /usr/sbin/nft -a list chain "$TABLE_FAMILY" "$TABLE_NAME" "$CHAIN_NAME" \
            | grep -F "comment \"$comment\"" \
            | sed -n 's/.* # handle \([0-9][0-9]*\)$/\1/p'
    )
}

status_rule() {
    local source_id="$1" target_id="$2" protocol="$3" port="$4" comment
    ensure_firewall_available
    # A status is meaningful only for identities currently known to the mesh.
    lookup_node_ip "$source_id" >/dev/null
    lookup_node_ip "$target_id" >/dev/null
    comment="smm:${source_id}:${target_id}:${protocol}:${port}"
    if rule_exists "$comment"; then
        printf '%s\n' active
    else
        printf '%s\n' disabled
    fi
}

reconcile_status() {
    local generation
    exec 9>"$RECONCILE_LOCK"
    chmod 0600 "$RECONCILE_LOCK"
    "$FLOCK_COMMAND" -x 9
    if [[ ! -f "$RECONCILE_MARKER" ]]; then
        printf '%s\n' complete
        return
    fi
    IFS= read -r generation <"$RECONCILE_MARKER" || fail "invalid reconciliation marker"
    [[ "$generation" =~ $generation_pattern ]] || fail "invalid reconciliation marker"
    printf 'requested:%s\n' "$generation"
}

reconcile_complete() {
    local expected_generation="$1" current_generation=''
    [[ "$expected_generation" =~ $generation_pattern ]] || fail "invalid reconciliation generation"
    exec 9>"$RECONCILE_LOCK"
    chmod 0600 "$RECONCILE_LOCK"
    "$FLOCK_COMMAND" -x 9
    if [[ -f "$RECONCILE_MARKER" ]]; then
        IFS= read -r current_generation <"$RECONCILE_MARKER" || fail "invalid reconciliation marker"
        [[ "$current_generation" =~ $generation_pattern ]] || fail "invalid reconciliation marker"
        if [[ "$current_generation" == "$expected_generation" ]]; then
            rm -f -- "$RECONCILE_MARKER"
        fi
    fi
    printf '%s\n' complete
}

action="${1:-}"
case "$action" in
    reconcile-status)
        [[ $# -eq 1 ]] || fail "invalid reconcile-status argument count"
        reconcile_status
        exit 0
        ;;
    reconcile-complete)
        [[ $# -eq 2 ]] || fail "invalid reconcile-complete argument count"
        reconcile_complete "$2"
        exit 0
        ;;
esac
validate_rule "$@"
case "$action" in
    link-connect) connect_rule "$2" "$3" "$4" "$5" ;;
    link-disconnect) disconnect_rule "$2" "$3" "$4" "$5" ;;
    link-status) status_rule "$2" "$3" "$4" "$5" ;;
esac