#!/usr/bin/env bash
#
# deerflow-firewall.sh — egress firewall for the hardened DeerFlow container.
#
# Restricts outbound traffic from the DeerFlow Docker bridge so the agent
# cannot reach LAN devices, only the Internet plus a small allowlist of
# Wireguard hosts (Searx + local model servers).
#
# Threat model: a prompt-injected agent inside the container should be
# unable to scan or attack other devices on the home LAN. Inbound LAN
# traffic to the container's published ports is unaffected because the
# rules are stateful (ESTABLISHED,RELATED returns first).
#
# Idempotent: running `up` repeatedly will only ever leave one copy of
# each rule in DOCKER-USER. Running `down` removes them all.
#
# Usage:
#   deerflow-firewall.sh up      # install rules
#   deerflow-firewall.sh down    # remove rules
#   deerflow-firewall.sh status  # show DOCKER-USER chain
#
# Requires: iptables, root privileges, the `DOCKER-USER` chain (created
# by the Docker daemon at startup).

set -euo pipefail

readonly BRIDGE="br-deerflow"
readonly CHAIN="DOCKER-USER"

# IPs that the container IS allowed to reach inside RFC1918:
readonly -a ALLOW_HOSTS=(
    "10.67.67.1"   # Searx (Wireguard)
    "10.67.67.2"   # XTTS / Whisper / Ollama-local (Wireguard)
)

# Subnets that the container is NOT allowed to reach:
readonly -a BLOCK_NETS=(
    "192.168.3.0/24"   # home LAN
    "10.0.0.0/8"       # rest of /8 except whitelisted /32s above
    "172.16.0.0/12"    # other Docker bridges + RFC1918 leftovers
)

if [[ $EUID -ne 0 ]]; then
    echo "deerflow-firewall: must run as root" >&2
    exit 1
fi

require_chain() {
    if ! iptables -w -nL "$CHAIN" >/dev/null 2>&1; then
        echo "deerflow-firewall: chain $CHAIN does not exist — is the Docker daemon running?" >&2
        exit 2
    fi
}

# Run iptables only if the same rule is not already present.
add_rule() {
    local args=("$@")
    if ! iptables -w -C "$CHAIN" "${args[@]}" 2>/dev/null; then
        iptables -w -I "$CHAIN" "${args[@]}"
    fi
}

# Run iptables -D until the rule is gone (in case it was added multiple times).
del_rule() {
    local args=("$@")
    while iptables -w -C "$CHAIN" "${args[@]}" 2>/dev/null; do
        iptables -w -D "$CHAIN" "${args[@]}"
    done
}

cmd_up() {
    require_chain

    # NOTE: -I prepends, so build the chain bottom-up.
    # Final order from top of DOCKER-USER:
    #   1. ESTABLISHED,RELATED -> RETURN          (let inbound responses out)
    #   2. -d 10.67.67.1       -> RETURN          (Searx)
    #   3. -d 10.67.67.2       -> RETURN          (XTTS/Whisper/Ollama)
    #   4. -d 192.168.3.0/24   -> REJECT          (home LAN)
    #   5. -d 10.0.0.0/8       -> REJECT
    #   6. -d 172.16.0.0/12    -> REJECT
    # Insert rules in reverse so the final order matches the spec.

    for net in "${BLOCK_NETS[@]}"; do
        add_rule -i "$BRIDGE" -d "$net" -j REJECT --reject-with icmp-net-prohibited
    done

    for host in "${ALLOW_HOSTS[@]}"; do
        add_rule -i "$BRIDGE" -d "$host" -j RETURN
    done

    add_rule -i "$BRIDGE" -m conntrack --ctstate ESTABLISHED,RELATED -j RETURN

    if ! ip link show "$BRIDGE" >/dev/null 2>&1; then
        echo "deerflow-firewall: WARNING: bridge $BRIDGE does not exist yet."
        echo "deerflow-firewall: rules are installed and will activate as soon as you run"
        echo "    docker compose -f deer-flow/docker/docker-compose.yaml \\"
        echo "                   -f docker/docker-compose.override.yaml up -d"
    fi

    echo "deerflow-firewall: up"
}

cmd_down() {
    require_chain

    for net in "${BLOCK_NETS[@]}"; do
        del_rule -i "$BRIDGE" -d "$net" -j REJECT --reject-with icmp-net-prohibited
    done

    for host in "${ALLOW_HOSTS[@]}"; do
        del_rule -i "$BRIDGE" -d "$host" -j RETURN
    done

    del_rule -i "$BRIDGE" -m conntrack --ctstate ESTABLISHED,RELATED -j RETURN

    echo "deerflow-firewall: down"
}

cmd_status() {
    require_chain
    echo "DOCKER-USER chain (rules matching $BRIDGE):"
    # -nvL prints the input interface column so we can grep for our bridge.
    iptables -w -nvL "$CHAIN" --line-numbers | awk -v b="$BRIDGE" 'NR<=2 || $0 ~ b'

    if ip link show "$BRIDGE" >/dev/null 2>&1; then
        echo
        echo "Bridge $BRIDGE: present"
    else
        echo
        echo "Bridge $BRIDGE: NOT present (DeerFlow container not started)"
    fi
}

case "${1:-}" in
    up)     cmd_up ;;
    down)   cmd_down ;;
    status) cmd_status ;;
    *)
        echo "Usage: $0 {up|down|status}" >&2
        exit 64
        ;;
esac