#!/bin/bash

# Script para configurar NAT BPF com TC

# Parâmetros
INTERFACE=${1:-"enp0s6"}
TARGET_IP=${2:-"10.15.0.45"}
GATEWAY_IP=${3:-"10.0.0.87"}

echo "=== Configurando NAT BPF com TC ==="
echo "Interface: $INTERFACE"
echo "IP alvo (container): $TARGET_IP"
echo "Gateway IP: $GATEWAY_IP"

# Compilar programa TC BPF
echo "Compilando programa TC BPF..."
clang -O2 -Wall -target bpf -g -D__TARGET_ARCH_arm64 -c tc_nat.c -o tc_nat.o -I/usr/include/bpf -I/usr/include/

if [ $? -ne 0 ]; then
    echo "Erro na compilação do programa TC!"
    exit 1
fi

# Verificar seções no arquivo objeto BPF
echo "Verificando seções no arquivo objeto..."
llvm-objdump -h tc_nat.o

# Remover programas TC existentes na interface
echo "Removendo programas TC existentes na interface $INTERFACE..."
tc filter del dev $INTERFACE ingress 2>/dev/null
tc filter del dev $INTERFACE egress 2>/dev/null
tc qdisc del dev $INTERFACE clsact 2>/dev/null

# Configurar qdisc
echo "Configurando qdisc no $INTERFACE..."
tc qdisc add dev $INTERFACE clsact

# Carregar programas TC
echo "Carregando programa TC ingress..."
tc filter add dev $INTERFACE ingress prio 1 handle 1 bpf da obj tc_nat.o sec tc_ingress
if [ $? -ne 0 ]; then
    echo "Erro ao carregar o programa ingress!"
    exit 1
fi

echo "Carregando programa TC egress..."
tc filter add dev $INTERFACE egress prio 1 handle 1 bpf da obj tc_nat.o sec tc_egress
if [ $? -ne 0 ]; then
    echo "Erro ao carregar o programa egress!"
    exit 1
fi

# Encontrar IDs dos mapas
echo "Encontrando IDs dos mapas BPF..."
DNAT_MAP_ID=$(bpftool map list | grep -m1 "name dnat_port_map" | awk '{print $1}' | tr -d ':')
SNAT_MAP_ID=$(bpftool map list | grep -m1 "name snat_networks" | awk '{print $1}' | tr -d ':')

if [ -z "$DNAT_MAP_ID" ] || [ -z "$SNAT_MAP_ID" ]; then
    echo "Não foi possível encontrar todos os mapas BPF necessários."
    exit 1
fi

echo "IDs dos mapas: DNAT=$DNAT_MAP_ID, SNAT=$SNAT_MAP_ID"

# Limpar mapas existentes
echo "Limpando mapas existentes..."
echo "Limpando DNAT..."
bpftool map dump id $DNAT_MAP_ID -j | grep -o '"key": [0-9]*' | awk '{print $2}' | \
while read key; do
    [ ! -z "$key" ] && bpftool map delete id $DNAT_MAP_ID key hex $(printf "%04x" $key | sed 's/\(..\)\(..\)/\2 \1/')
done

echo "Limpando SNAT..."
bpftool map dump id $SNAT_MAP_ID -j | grep -o '"key": [0-9]*' | awk '{print $2}' | \
while read key; do
    [ ! -z "$key" ] && bpftool map delete id $SNAT_MAP_ID key hex $(printf "%08x" $key | sed 's/\(..\)\(..\)\(..\)\(..\)/\4 \3 \2 \1/')
done

# Configurar regras DNAT
echo "Configurando regras DNAT..."

# Função auxiliar para converter IP para valores em little-endian
ip_to_hex_le() {
    local ip=$1
    IFS='.' read -r -a octets <<< "$ip"
    printf "%02x %02x %02x %02x" "${octets[3]}" "${octets[2]}" "${octets[1]}" "${octets[0]}"
}

# Função auxiliar para converter porta para little-endian
port_to_hex_le() {
    local port=$1
    printf "%02x %02x" $((port & 0xFF)) $(((port >> 8) & 0xFF))
}

# Obter valores em hex
TARGET_IP_HEX=$(ip_to_hex_le $TARGET_IP)

# Configurar redirecionamento HTTP
echo "Configurando redirecionamento HTTP (porta 80)..."
bpftool map update id $DNAT_MAP_ID key hex "$(port_to_hex_le 80)" value hex "$TARGET_IP_HEX $(port_to_hex_le 80) 00 00"

# Configurar redirecionamento HTTPS
echo "Configurando redirecionamento HTTPS (porta 443)..."
bpftool map update id $DNAT_MAP_ID key hex "$(port_to_hex_le 443)" value hex "$TARGET_IP_HEX $(port_to_hex_le 443) 00 00"

# Configurar redes para SNAT
echo "Configurando redes para SNAT..."
bpftool map update id $SNAT_MAP_ID key hex "00 00 0f 0a" value hex "01"  # 10.15.0.0/16
bpftool map update id $SNAT_MAP_ID key hex "00 00 11 ac" value hex "01"  # 172.17.0.0/16

# Habilitar IP forwarding
echo "Habilitando ip_forward..."
echo 1 > /proc/sys/net/ipv4/ip_forward

echo "NAT BPF configurado com sucesso. Utilize 'tc filter show dev $INTERFACE ingress' para verificar."
echo "Para ver logs: cat /sys/kernel/debug/tracing/trace_pipe"