#!/bin/bash
#
# air-netshape — strict-allowlist wrapper around /sbin/tc for
# operator-driven network-condition simulation on the Pi's egress to
# the video destination (GCS / laptop / viewer). Invoked via sudo
# from cast-server; sudoers restricts the cast service user to exactly this
# path (/etc/sudoers.d/air-vpn, Cmnd_Alias AIR_NETSHAPE).
#
# Why a wrapper?
# --------------
# tc(8) lets you configure anything from HTB bandwidth limits to
# netem packet loss to raw qdisc replacement on ANY interface. We
# only want to permit a single, well-defined combination (HTB rate
# cap + netem delay/jitter/loss, scoped to traffic destined for a
# specific IP:port, on a specific interface). The wrapper enforces
# that combination and refuses anything else.
#
# Subcommand allowlist:
#
#   status                                     — dump current qdiscs + filters for audit.
#   clear                                      — tear down every qdisc this tool installed.
#   apply <rate_kbps> <latency_ms> <jitter_ms> <loss_pct> <dst_host> <dst_port>
#                                              — install a ~4G/LTE-style shaping rule scoped
#                                                to traffic destined for <dst_host>:<dst_port>
#                                                on the video interface (AIR_NETSHAPE_IFACE,
#                                                default eth0). Safe default: the control
#                                                path (HTTP on port 80, SSH on 22) is
#                                                NOT shaped because the filter only matches
#                                                the configured video destination.
#
# Every numeric arg is validated against a strict regex. Non-numeric
# / out-of-range / non-host input is refused with exit 2. The rate is
# permitted in [32, 100000] kbps (32 kbps to 100 Mbps); latency in
# [0, 5000] ms; jitter in [0, 1000] ms; loss in [0, 50] %.
#
# The scheme:
#   Root HTB qdisc on the egress interface with a single class rate-
#   capped to <rate_kbps>. A u32 filter matches outbound packets whose
#   destination IP + destination UDP/TCP port are (dst_host, dst_port)
#   and directs them into the shaped class; all other traffic falls
#   through to the default class (unlimited) so the UI and SSH keep
#   working. A netem qdisc sits under the shaped class and adds the
#   requested delay/jitter/loss to the matched flow only.
#
# Idempotent: re-running `apply` with different knobs replaces the
# existing rules; re-running `clear` when nothing is installed is a
# no-op.

set -euo pipefail

TC=/sbin/tc
IFACE="${AIR_NETSHAPE_IFACE:-eth0}"
PROG=air-netshape
# Fixed qdisc handle so the `status` and `clear` paths can find the
# structure this tool wrote, even if the operator rebooted between
# apply and clear.
ROOT_HANDLE="1:"
SHAPED_CLASS="1:10"
DEFAULT_CLASS="1:20"
NETEM_HANDLE="10:"

die() {
    echo "${PROG}: $*" >&2
    exit 2
}

is_nonneg_int_in_range() {
    # $1 = value, $2 = min (inclusive), $3 = max (inclusive)
    [[ "$1" =~ ^[0-9]+$ ]] && [ "$1" -ge "$2" ] && [ "$1" -le "$3" ]
}

is_ipv4() {
    # Strict dotted-quad with octet bounds.
    local ip=$1 IFS=.
    [[ "$ip" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]] || return 1
    local a b c d
    read -r a b c d <<<"$ip"
    for n in "$a" "$b" "$c" "$d"; do
        [ "$n" -le 255 ] || return 1
    done
    return 0
}

is_port() {
    is_nonneg_int_in_range "$1" 1 65535
}

if [ $# -lt 1 ]; then
    die "missing subcommand"
fi

subcmd=$1
shift

case "$subcmd" in
    status)
        if [ $# -ne 0 ]; then die "'status' takes no arguments"; fi
        # Non-fatal if nothing is configured — print what's there and
        # let the caller parse.
        printf 'iface: %s\n' "$IFACE"
        printf -- '--- qdisc ---\n'
        "$TC" qdisc show dev "$IFACE" || true
        printf -- '--- class ---\n'
        "$TC" class show dev "$IFACE" || true
        printf -- '--- filter ---\n'
        "$TC" filter show dev "$IFACE" || true
        exit 0
        ;;
    clear)
        if [ $# -ne 0 ]; then die "'clear' takes no arguments"; fi
        # `tc qdisc del root` wipes our entire structure (class,
        # netem, filter) in one atomic call. Tolerate "RTNETLINK: No
        # such file or directory" when nothing was installed — that
        # is the idempotent-no-op case, not a failure.
        if "$TC" qdisc del dev "$IFACE" root 2>/dev/null; then
            echo "${PROG}: cleared"
        else
            echo "${PROG}: nothing to clear" >&2
        fi
        exit 0
        ;;
    apply)
        if [ $# -ne 6 ]; then
            die "'apply' needs: <rate_kbps> <latency_ms> <jitter_ms> <loss_pct> <dst_host> <dst_port>"
        fi
        RATE=$1; LATENCY=$2; JITTER=$3; LOSS=$4; HOST=$5; PORT=$6
        is_nonneg_int_in_range "$RATE"    32     100000 || die "rate_kbps out of range [32, 100000]"
        is_nonneg_int_in_range "$LATENCY" 0      5000   || die "latency_ms out of range [0, 5000]"
        is_nonneg_int_in_range "$JITTER"  0      1000   || die "jitter_ms out of range [0, 1000]"
        is_nonneg_int_in_range "$LOSS"    0      50     || die "loss_pct out of range [0, 50]"
        is_ipv4 "$HOST"                                  || die "dst_host must be an IPv4 address"
        is_port "$PORT"                                  || die "dst_port must be a valid TCP/UDP port"

        # Replace-style apply: wipe any prior structure first so
        # repeated apply calls with different knobs are idempotent.
        "$TC" qdisc del dev "$IFACE" root 2>/dev/null || true

        # Root HTB with a default class that matches "everything
        # we didn't filter into the shaped class" — so browser UI,
        # SSH, DHCP, and mDNS are NOT shaped. default=20.
        "$TC" qdisc add dev "$IFACE" root handle "$ROOT_HANDLE" htb default 20

        # Shaped class: the rate cap lives here. `ceil == rate` so
        # the class cannot borrow from siblings. Burst is small so
        # the cap is actually enforced on short flows.
        "$TC" class add dev "$IFACE" parent "$ROOT_HANDLE" \
            classid "$SHAPED_CLASS" htb rate "${RATE}kbit" ceil "${RATE}kbit" \
            burst 15k cburst 15k

        # Default (un-shaped) class: gigabit ceiling so it is
        # effectively "no limit" on any sensible link.
        "$TC" class add dev "$IFACE" parent "$ROOT_HANDLE" \
            classid "$DEFAULT_CLASS" htb rate 1gbit ceil 1gbit

        # Netem under the shaped class: adds the LTE-style delay /
        # jitter / loss on top of the rate cap. Only the shaped
        # class is degraded; the default class is clean.
        NETEM_ARGS=()
        if [ "$LATENCY" -gt 0 ] || [ "$JITTER" -gt 0 ]; then
            if [ "$JITTER" -gt 0 ]; then
                NETEM_ARGS+=(delay "${LATENCY}ms" "${JITTER}ms" distribution normal)
            else
                NETEM_ARGS+=(delay "${LATENCY}ms")
            fi
        fi
        if [ "$LOSS" -gt 0 ]; then
            NETEM_ARGS+=(loss "${LOSS}%")
        fi
        if [ ${#NETEM_ARGS[@]} -gt 0 ]; then
            "$TC" qdisc add dev "$IFACE" parent "$SHAPED_CLASS" \
                handle "$NETEM_HANDLE" netem "${NETEM_ARGS[@]}"
        fi

        # u32 filter: match outbound packets whose destination IP is
        # $HOST AND destination port is $PORT. Both UDP (video) and
        # TCP (if the operator picked a TCP receiver) are handled by
        # the generic L4 port match at offset 22 in the IP header.
        # `match ip dport $PORT 0xffff` matches any L4 protocol —
        # good enough for this use case, and avoids needing two
        # parallel filters for UDP + TCP.
        "$TC" filter add dev "$IFACE" parent "$ROOT_HANDLE" protocol ip prio 1 u32 \
            match ip dst "$HOST/32" \
            match ip dport "$PORT" 0xffff \
            flowid "$SHAPED_CLASS"

        printf '%s: applied — %s:%s at %skbit, +%sms (±%sms), %s%% loss\n' \
            "$PROG" "$HOST" "$PORT" "$RATE" "$LATENCY" "$JITTER" "$LOSS"
        exit 0
        ;;
    *)
        die "unknown subcommand: ${subcmd}"
        ;;
esac
