#!/bin/sh /etc/rc.common 
# AuthShield init (procd)
# Enforces bans for LuCI/Dropbear by dropping at top of input_lan/input_wan.
# Generates /var/run/authshield.nft as a *script* include for fw4.

USE_PROCD=1
START=60
STOP=15
NAME=authshield

WATCH_BIN="/usr/sbin/authshield.sh"
FW_INCLUDE="/var/run/authshield.nft"
BAN_TRACK_FILE="/var/run/authshield.bans" # File storing ban history
CIRCUIT_STATUS_FILE="/var/run/authshield.circuit" # Circuit breaker state
SET_V4="authshield_penalty_v4"
SET_V6="authshield_penalty_v6"
SET_CIRCUIT="authshield_circuit_ports"

# ---- helpers ----

get_uciv() {
    uci -q get "$1"
}

# Normalize a space/semicolon/comma separated port list into a sorted, deduped CSV (e.g. "80,443")
ports_to_csv() {
    printf '%s\n' "$1" \
    | tr ' ,;' '\n' \
    | awk '
        NF {
            gsub(/[^0-9]/, "")
            if ($0 != "") {
                if (!seen[$0]++) {
                    if (out != "") out = out "," $0
                    else out = $0
                }
            }
        }
        END { print out }
    '
}

# Compute effective ports (normalize + add 22 if Dropbear monitoring is enabled)
get_effective_ports() {
    local ports watch_dropbear

    ports="$(get_uciv authshield.@settings[0].ports)"
    [ -n "$ports" ] || ports="80 443"

    watch_dropbear="$(get_uciv authshield.@settings[0].watch_dropbear)"
    [ -n "$watch_dropbear" ] || watch_dropbear=0

    if [ "$watch_dropbear" = "1" ]; then
        case " $ports " in
            *" 22 "*) : ;;
            *) ports="$ports 22" ;;
        esac
    fi

    ports="$(ports_to_csv "$ports")"
    echo "$ports"
}


ensure_fw_include() {
    # ensure we have a firewall include pointing to $FW_INCLUDE with type 'script'
    local found idx
    found=0
    idx=0
    while :; do
        local path type
        path="$(get_uciv firewall.@include[$idx].path)" || break
        type="$(get_uciv firewall.@include[$idx].type)"
        if [ "$path" = "$FW_INCLUDE" ]; then
            [ "$type" = "script" ] || uci set firewall.@include[$idx].type='script'
            found=1
            break
        fi
        idx=$((idx+1))
    done

    if [ "$found" -eq 0 ]; then
        uci add firewall include >/dev/null
        uci set firewall.@include[-1].type='script'
        uci set firewall.@include[-1].path="$FW_INCLUDE"
    fi
    uci commit firewall
}

write_fw_include() {
    # Read config from UCI (use @settings[0] consistently)
    local ports penalty circuit_enable circuit_penalty
    ports="$(get_effective_ports)"
    penalty="$(get_uciv authshield.@settings[0].penalty)"
    circuit_enable="$(get_uciv authshield.@settings[0].circuit_enable)"
    circuit_penalty="$(get_uciv authshield.@settings[0].circuit_penalty)"

    [ -n "$ports" ] || ports="80 443"
    [ -n "$penalty" ] || penalty="60"
    [ -n "$circuit_enable" ] || circuit_enable=0
    [ -n "$circuit_penalty" ] || circuit_penalty=3600

    local ports_csv
    ports_csv="$ports"

    cat > "$FW_INCLUDE" <<EOF
#!/bin/sh
# Auto-generated by /etc/init.d/authshield – DO NOT EDIT.
set -eu

PORTS_CSV="$ports_csv"
PENALTY="$penalty"
SET_V4="$SET_V4"
SET_V6="$SET_V6"
SET_CIRCUIT="$SET_CIRCUIT"
CIRCUIT_ENABLE="$circuit_enable"
CIRCUIT_PENALTY="$circuit_penalty"

# Delete previous rules
del_old_rules() {
  local chain="\$1"
  nft -a list chain inet fw4 "\$chain" 2>/dev/null | \
    awk '/@'\$SET_V4'/ || /@'\$SET_V6'/ && / dport / {print \$NF}' | \
    tr -d ';' | while read -r h; do
      [ -n "\$h" ] && nft delete rule inet fw4 "\$chain" handle "\$h" 2>/dev/null || true
    done
}

# Delete circuit breaker rule from input_wan
del_circuit_rule() {
  nft -a list chain inet fw4 input_wan 2>/dev/null | \
    awk '/@'\$SET_CIRCUIT'/ && / dport / {print \$NF}' | \
    tr -d ';' | while read -r h; do
      [ -n "\$h" ] && nft delete rule inet fw4 input_wan handle "\$h" 2>/dev/null || true
    done
}

# Ensure table exists (use shell to avoid nft parse errors)
if ! nft list table inet fw4 >/dev/null 2>&1; then
  nft add table inet fw4
fi

# Remove any old rules we inserted in 'input' to avoid duplication
del_old_rules input || true

# Add sets with pure nft syntax (no shell redirects inside)
nft -f - <<NFE
add set inet fw4 \$SET_V4 { type ipv4_addr; flags timeout; timeout \${PENALTY}s; }
add set inet fw4 \$SET_V6 { type ipv6_addr; flags timeout; timeout \${PENALTY}s; }
NFE

# Insert the early-drop rules BEFORE conntrack established/related accept
nft insert rule inet fw4 input index 1 tcp dport {\${PORTS_CSV}} ip  saddr @\${SET_V4} counter drop 2>/dev/null || true
nft insert rule inet fw4 input index 1 tcp dport {\${PORTS_CSV}} ip6 saddr @\${SET_V6} counter drop 2>/dev/null || true

# Circuit breaker setup (if enabled)
if [ "\$CIRCUIT_ENABLE" = "1" ]; then
  # Remove old circuit breaker rule
  del_circuit_rule || true
  
  # Create circuit breaker port set with timeout capability
  nft add set inet fw4 \$SET_CIRCUIT "{ type inet_service; flags timeout; timeout \${CIRCUIT_PENALTY}s; }" 2>/dev/null || true
  
  # Add rule to input_wan that drops traffic when ports are in the circuit set
  # Insert at index 1 to be before any accept rules
  if nft list chain inet fw4 input_wan >/dev/null 2>&1; then
    nft insert rule inet fw4 input_wan index 1 tcp dport @\${SET_CIRCUIT} counter drop 2>/dev/null || true
  fi
fi

exit 0
EOF
    chmod +x "$FW_INCLUDE"
}

regen_rules_and_reload_fw() {
    write_fw_include
    ensure_fw_include
    /etc/init.d/firewall reload >/dev/null 2>&1 || true
}

kill_leftovers() {
    # best-effort cleanup of stray watcher shells
    local pids
    pids="$(pgrep -f "$WATCH_BIN" 2>/dev/null || true)"
    [ -z "$pids" ] && return 0
    kill $pids 2>/dev/null || true
    sleep 1
    pids="$(pgrep -f "$WATCH_BIN" 2>/dev/null || true)"
    [ -z "$pids" ] || kill -9 $pids 2>/dev/null || true
}

# Delete previous rules
del_old_rules() {
  local chain="$1"
  nft -a list chain inet fw4 "$chain" 2>/dev/null |     awk '/@'$SET_V4'/ || /@'$SET_V6'/ && / dport / {print $NF}' |     tr -d ';' | while read -r h; do
      [ -n "$h" ] && nft delete rule inet fw4 "$chain" handle "$h" 2>/dev/null || true
    done
}

# Remove circuit breaker rule
del_circuit_rule() {
    nft -a list chain inet fw4 input_wan 2>/dev/null | \
        awk '/@'$SET_CIRCUIT'/ && / dport / {print $NF}' | \
        tr -d ';' | while read -r handle; do
            [ -n "$handle" ] && nft delete rule inet fw4 input_wan handle "$handle" 2>/dev/null || true
        done
}

# ---- procd lifecycle ----

start_service() {
    local enabled threshold window penalty ports watch_dropbear ignore_private escalate_enable
    local circuit_enable circuit_threshold circuit_window circuit_penalty

    # Only when enabled do we continue
    enabled="$(get_uciv authshield.@settings[0].enabled)"; [ -n "$enabled" ] || enabled=1
    [ "$enabled" -eq 1 ] || return 0

    # Generate/refresh nftables include & reload firewall
    regen_rules_and_reload_fw

    # Read watcher env from UCI (@settings[0] consistently)
    threshold="$(get_uciv authshield.@settings[0].threshold)";       [ -n "$threshold" ] || threshold=5
    window="$(get_uciv authshield.@settings[0].window)";             [ -n "$window" ] || window=10
    penalty="$(get_uciv authshield.@settings[0].penalty)";           [ -n "$penalty" ] || penalty=60
    global_enable="$(get_uciv authshield.@settings[0].global_enable)";       [ -n "$global_enable" ] || global_enable=1
    global_threshold="$(get_uciv authshield.@settings[0].global_threshold)"; [ -n "$global_threshold" ] || global_threshold=60
    global_window="$(get_uciv authshield.@settings[0].global_window)";       [ -n "$global_window" ] || global_window=43200
    global_penalty="$(get_uciv authshield.@settings[0].global_penalty)";     [ -n "$global_penalty" ] || global_penalty=86400
    ports="$(get_effective_ports)"
    watch_dropbear="$(get_uciv authshield.@settings[0].watch_dropbear)"; [ -n "$watch_dropbear" ] || watch_dropbear=0
    ignore_private="$(get_uciv authshield.@settings[0].ignore_private_ip)"; [ -n "$ignore_private" ] || ignore_private=1
    escalate_enable="$(get_uciv authshield.@settings[0].escalate_enable)"; [ -n "$escalate_enable" ] || escalate_enable=1
    escalate_threshold="$(get_uciv authshield.@settings[0].escalate_threshold)"; [ -n "$escalate_threshold" ] || escalate_threshold=5
    escalate_window="$(get_uciv authshield.@settings[0].escalate_window)";     [ -n "$escalate_window" ] || escalate_window=3600
    escalate_penalty="$(get_uciv authshield.@settings[0].escalate_penalty)";   [ -n "$escalate_penalty" ] || escalate_penalty=86400
    
    # Circuit breaker settings
    circuit_enable="$(get_uciv authshield.@settings[0].circuit_enable)"; [ -n "$circuit_enable" ] || circuit_enable=1
    circuit_threshold="$(get_uciv authshield.@settings[0].circuit_threshold)"; [ -n "$circuit_threshold" ] || circuit_threshold=120
    circuit_window="$(get_uciv authshield.@settings[0].circuit_window)"; [ -n "$circuit_window" ] || circuit_window=43200
    circuit_penalty="$(get_uciv authshield.@settings[0].circuit_penalty)"; [ -n "$circuit_penalty" ] || circuit_penalty=3600

    [ "$enabled" -eq 1 ] || return 0

    procd_open_instance
    procd_set_param command "$WATCH_BIN"
    procd_set_param respawn 5 10 5           # (timeout, retry, max)
    procd_set_param stdout 1
    procd_set_param stderr 1
    procd_set_param env THRESHOLD="$threshold" \
                        WINDOW="$window" \
                        PENALTY="$penalty" \
                        PORTS="$ports" \
                        WATCH_DROPBEAR="$watch_dropbear" \
                        IGNORE_PRIVATE="$ignore_private" \
                        ESCALATE_ENABLE="$escalate_enable" \
                        ESCALATE_THRESHOLD="$escalate_threshold" \
                        ESCALATE_WINDOW="$escalate_window" \
                        ESCALATE_PENALTY="$escalate_penalty" \
                        BAN_TRACK_FILE="$BAN_TRACK_FILE" \
                        GLOBAL_ENABLE="$global_enable" \
                        GLOBAL_THRESHOLD="$global_threshold" \
                        GLOBAL_WINDOW="$global_window" \
                        GLOBAL_PENALTY="$global_penalty" \
                        CIRCUIT_ENABLE="$circuit_enable" \
                        CIRCUIT_THRESHOLD="$circuit_threshold" \
                        CIRCUIT_WINDOW="$circuit_window" \
                        CIRCUIT_PENALTY="$circuit_penalty" \
                        CIRCUIT_STATUS_FILE="$CIRCUIT_STATUS_FILE" \
                        SET_V4="$SET_V4" \
                        SET_V6="$SET_V6"
    procd_close_instance
}

stop_service() {
    # procd will stop our instance; we just clean up stragglers
    kill_leftovers || true

    # remove the firewall include file only if it exists
    [ -f "$FW_INCLUDE" ] && rm -f "$FW_INCLUDE"

    # remove the BAN_TRACK_FILE only if it exists
    [ -f "$BAN_TRACK_FILE" ] && rm -f "$BAN_TRACK_FILE"
    
    # remove the CIRCUIT_STATUS_FILE only if it exists
    [ -f "$CIRCUIT_STATUS_FILE" ] && rm -f "$CIRCUIT_STATUS_FILE"

    # Remove any old rules we inserted in 'input'
    del_old_rules input || true
    
    # Remove circuit breaker rule from input_wan
    del_circuit_rule || true

    # delete all the sets we created
    nft delete set inet fw4 "$SET_V4" 2>/dev/null || true
    nft delete set inet fw4 "$SET_V6" 2>/dev/null || true
    nft delete set inet fw4 "$SET_CIRCUIT" 2>/dev/null || true
}

reload_service() {
    # Re-read UCI, rebuild procd instance, and relaunch with fresh env
    stop
    start
}

service_triggers() {
    # Reload our service when authshield UCI changes
    procd_add_reload_trigger "authshield"
}

# Convenience handler for rc.common `restart`
restart() {
    stop
    start
}