#!/bin/bash
# ─────────────────────────────────────────────────────────
# loader.sh — Carrega e anexa os programas eBPF sockmap
#
# USO:
#   sudo ./scripts/loader.sh load     # Carrega e ativa o sockmap
#   sudo ./scripts/loader.sh unload   # Remove o sockmap
#   sudo ./scripts/loader.sh status   # Mostra status e estatísticas
# ─────────────────────────────────────────────────────────
set -euo pipefail

# Cores para output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

BPF_OBJ="build/sockmap.bpf.o"
BPF_PIN_DIR="/sys/fs/bpf/sockmap"
CGROUP_PATH=""

log_info()  { echo -e "${BLUE}[INFO]${NC} $1"; }
log_ok()    { echo -e "${GREEN}[OK]${NC} $1"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERRO]${NC} $1"; }

# ─────────────────────────────────────────────────────────
# Detecta o caminho do cgroup v2
# Docker coloca os containers dentro de um cgroup
# ─────────────────────────────────────────────────────────
detect_cgroup() {
    # Verifica se cgroup v2 está montado (unified hierarchy)
    if mount | grep -q "cgroup2 on /sys/fs/cgroup"; then
        CGROUP_PATH="/sys/fs/cgroup"
        log_info "Cgroup v2 (unified) detectado em $CGROUP_PATH"
    elif [ -d "/sys/fs/cgroup/unified" ]; then
        CGROUP_PATH="/sys/fs/cgroup/unified"
        log_info "Cgroup v2 (unified) detectado em $CGROUP_PATH"
    else
        log_error "Cgroup v2 não encontrado! eBPF sockmap requer cgroup v2."
        log_error "Adicione 'systemd.unified_cgroup_hierarchy=1' aos parâmetros do kernel."
        exit 1
    fi
}

# ─────────────────────────────────────────────────────────
# Verifica pré-requisitos
# ─────────────────────────────────────────────────────────
check_prereqs() {
    if [ "$(id -u)" -ne 0 ]; then
        log_error "Este script precisa ser executado como root (sudo)!"
        exit 1
    fi

    if ! command -v bpftool &>/dev/null; then
        log_error "'bpftool' não encontrado. Instale com: apt install linux-tools-common linux-tools-$(uname -r)"
        exit 1
    fi

    if [ ! -f "$BPF_OBJ" ]; then
        log_error "Objeto BPF não encontrado em '$BPF_OBJ'. Execute 'make' primeiro!"
        exit 1
    fi
}

# ─────────────────────────────────────────────────────────
# LOAD: Carrega os programas BPF e os anexa
# ─────────────────────────────────────────────────────────
do_load() {
    check_prereqs
    detect_cgroup

    # Limpa carregamento anterior se existir
    if [ -d "$BPF_PIN_DIR" ]; then
        log_warn "Sockmap já carregado. Removendo anterior..."
        do_unload
    fi

    log_info "Criando diretório de pins BPF..."
    mkdir -p "$BPF_PIN_DIR"

    # ── 1. Carregar TODOS os programas do objeto de uma vez ──
    # 'loadall' respeita as anotações SEC() de cada programa,
    # carregando sock_ops e sk_msg com seus tipos corretos.
    # Os programas são fixados (pinned) com o nome das funções C:
    #   bpf_sockops_handler  → tipo sockops
    #   bpf_msg_redirect     → tipo sk_msg
    log_info "Carregando programas BPF (loadall)..."
    bpftool prog loadall "$BPF_OBJ" "$BPF_PIN_DIR" \
        pinmaps "$BPF_PIN_DIR"
    log_ok "Programas sock_ops e sk_msg carregados"

    # Confirma que os pins foram criados
    if [ ! -f "$BPF_PIN_DIR/bpf_sockops_handler" ]; then
        log_error "Pin bpf_sockops_handler não encontrado em $BPF_PIN_DIR!"
        log_error "Programas no objeto:"
        bpftool prog list 2>/dev/null | tail -10
        exit 1
    fi
    if [ ! -f "$BPF_PIN_DIR/bpf_msg_redirect" ]; then
        log_error "Pin bpf_msg_redirect não encontrado em $BPF_PIN_DIR!"
        exit 1
    fi

    # ── 2. Anexar sock_ops ao cgroup ──
    # Flag 'multi' permite coexistir com outros programas sock_ops já
    # anexados (ex: systemd ou outros agentes), sem conflito.
    log_info "Anexando sock_ops ao cgroup ($CGROUP_PATH)..."
    bpftool cgroup attach "$CGROUP_PATH" sock_ops \
        pinned "$BPF_PIN_DIR/bpf_sockops_handler" \
        multi
    log_ok "sock_ops anexado ao cgroup"

    # ── 3. Anexar sk_msg ao mapa sockhash ──
    log_info "Anexando sk_msg ao mapa sockhash..."
    bpftool prog attach \
        pinned "$BPF_PIN_DIR/bpf_msg_redirect" \
        msg_verdict \
        pinned "$BPF_PIN_DIR/sock_hash"
    log_ok "sk_msg anexado ao sockhash"

    echo ""
    log_ok "═══════════════════════════════════════════════"
    log_ok "  eBPF Sockmap ATIVADO com sucesso!"
    log_ok "  Comunicação entre containers está acelerada."
    log_ok "═══════════════════════════════════════════════"
    echo ""
    log_info "Use 'sudo ./scripts/loader.sh status' para ver estatísticas"
    log_info "Use 'sudo ./scripts/loader.sh unload' para desativar"
}

# ─────────────────────────────────────────────────────────
# UNLOAD: Remove os programas BPF
# ─────────────────────────────────────────────────────────
do_unload() {
    detect_cgroup

    log_info "Desanexando programas BPF..."

    # Desanexa sock_ops do cgroup
    if [ -f "$BPF_PIN_DIR/bpf_sockops_handler" ]; then
        bpftool cgroup detach "$CGROUP_PATH" sock_ops \
            pinned "$BPF_PIN_DIR/bpf_sockops_handler" \
            multi 2>/dev/null || true
        log_ok "sock_ops desanexado do cgroup"
    fi

    # Desanexa sk_msg do mapa
    if [ -f "$BPF_PIN_DIR/bpf_msg_redirect" ]; then
        bpftool prog detach \
            pinned "$BPF_PIN_DIR/bpf_msg_redirect" \
            msg_verdict \
            pinned "$BPF_PIN_DIR/sock_hash" 2>/dev/null || true
        log_ok "sk_msg desanexado do sockhash"
    fi

    # Remove pins
    log_info "Removendo pins BPF..."
    rm -rf "$BPF_PIN_DIR"
    log_ok "Pins removidos"

    echo ""
    log_ok "eBPF Sockmap DESATIVADO com sucesso."
}

# ─────────────────────────────────────────────────────────
# STATUS: Mostra status e estatísticas
# ─────────────────────────────────────────────────────────
do_status() {
    echo ""
    echo -e "${BLUE}═══ eBPF Sockmap Status ═══${NC}"
    echo ""

    if [ ! -d "$BPF_PIN_DIR" ]; then
        log_warn "Sockmap NÃO está carregado."
        return
    fi

    log_ok "Sockmap está ATIVO"
    echo ""

    # Mostra programas carregados
    echo -e "${YELLOW}Programas BPF carregados:${NC}"
    bpftool prog show pinned "$BPF_PIN_DIR/bpf_sockops_handler" 2>/dev/null || echo "  sock_ops: não encontrado"
    bpftool prog show pinned "$BPF_PIN_DIR/bpf_msg_redirect" 2>/dev/null || echo "  sk_msg: não encontrado"
    echo ""

    # Mostra estatísticas
    # O valor no mapa é u64 little-endian (8 bytes).
    # bpftool --json pode retornar cada byte como int (10) ou string hex ("0x0a")
    # dependendo da versão. O helper parse_bpftool_bytes trata ambos.
    echo -e "${YELLOW}Estatísticas:${NC}"
    if [ -f "$BPF_PIN_DIR/stats" ]; then
        read_stat() {
            local idx=$1
            bpftool map lookup pinned "$BPF_PIN_DIR/stats" \
                key $(printf '0x%02x 0x00 0x00 0x00' "$idx") \
                --json 2>/dev/null \
            | python3 -c "
import sys, json
def to_bytes(arr):
    # bpftool >= 7.x retorna strings hex ('0x0a'), antigas retornam ints (10)
    return bytes(int(x, 16) if isinstance(x, str) else int(x) for x in arr)
try:
    d = json.load(sys.stdin)
    print(int.from_bytes(to_bytes(d['value']), 'little'))
except:
    print(0)
"
        }
        CONNS=$(read_stat 0)
        REDIR=$(read_stat 1)
        ERRS=$(read_stat 2)
        SKIP=$(read_stat 3)
        echo "  Conexões registradas no sockhash : $CONNS"
        echo "  Mensagens redirecionadas (accel.) : $REDIR"
        echo "  Fallback (par não encontrado)     : $ERRS"
        echo "  Ignoradas (IPs externos)          : $SKIP"
    else
        echo "  Mapa de estatísticas não encontrado (recompile e recarregue)"
    fi
    echo ""

    # Mostra sockets no mapa
    # NOTA: 'No space left on device' nos values é NORMAL.
    # SOCKHASH armazena ponteiros de socket do kernel, não legíveis de userspace.
    echo -e "${YELLOW}Conexões ativas no sockhash (IPs/portas):${NC}"
    echo -e "  ${BLUE}(Os 'value: No space left on device' são normais para SOCKHASH)${NC}"
    if [ -f "$BPF_PIN_DIR/sock_hash" ]; then
        SOCKMAP_OUT=$(bpftool map dump pinned "$BPF_PIN_DIR/sock_hash" --json 2>/dev/null \
            | python3 -c "
import sys, json, socket, struct

def to_int(x):
    # Aceita tanto int (10) quanto string hex ('0x0a') — depende da versão do bpftool
    return int(x, 16) if isinstance(x, str) else int(x)

try:
    entries = json.load(sys.stdin)
    if not entries:
        print('  (nenhuma conexão registrada)')
    else:
        print(f'  Total: {len(entries)} conexão(oes):')
        for e in entries[:25]:
            k = e.get('key', [])
            if len(k) >= 20:
                # struct sock_key: sip4 dip4 sport dport family (5x u32)
                # local_ip4/remote_ip4 estão em big-endian (network order)
                # portas em host order (little-endian)
                raw = [to_int(b) for b in k]
                sip   = socket.inet_ntoa(bytes(raw[0:4]))
                dip   = socket.inet_ntoa(bytes(raw[4:8]))
                sport = struct.unpack_from('<I', bytes(raw[8:12]))[0]
                dport = struct.unpack_from('<I', bytes(raw[12:16]))[0]
                print(f'  {sip}:{sport}  →  {dip}:{dport}')
except Exception as ex:
    print(f'  (erro ao decodificar: {ex})')
" 2>/dev/null || echo "  Nenhum socket registrado")
        SOCK_COUNT=$(echo "$SOCKMAP_OUT" | grep -oP 'Total: \K[0-9]+' || echo 0)
        echo "$SOCKMAP_OUT"
        # Aviso inteligente: usa contagem real do sockhash quando stats falham
        if [ "${CONNS:-0}" -eq 0 ] 2>/dev/null && [ "${SOCK_COUNT:-0}" -gt 0 ] 2>/dev/null; then
            echo ""
            echo -e "  ${GREEN}✔ Sockhash tem ${SOCK_COUNT} conexão(oes) ativas!${NC}"
            echo -e "  ${YELLOW}(Contadores de stats zerados podem indicar que os containers"
            echo -e "  foram reiniciados antes do sockmap ser carregado nesta sessão.)${NC}"
        elif [ "${CONNS:-0}" -eq 0 ] 2>/dev/null && [ "${SOCK_COUNT:-0}" -eq 0 ] 2>/dev/null; then
            echo ""
            echo -e "  ${YELLOW}⚠ Nenhuma conexão no sockhash.${NC}"
            echo "  → Reinicie os containers: cd /home/docker && docker compose restart"
            echo "  → Ou rode: sudo ./scripts/loader.sh diagnose"
        fi
    fi
    echo ""

    # Mostra logs do BPF (trace_pipe)
    echo -e "${YELLOW}Últimas mensagens do BPF (trace_pipe):${NC}"
    if [ -f /sys/kernel/debug/tracing/trace_pipe ]; then
        timeout 1 cat /sys/kernel/debug/tracing/trace_pipe 2>/dev/null | tail -10 || echo "  Nenhuma mensagem recente"
    fi
}

# ─────────────────────────────────────────────────────────
# DIAGNOSE: Força conexões TCP internas para verificar se o sockmap
# está interceptando corretamente.
# ─────────────────────────────────────────────────────────
do_diagnose() {
    if [ ! -d "$BPF_PIN_DIR" ]; then
        log_error "Sockmap não está carregado. Execute 'load' primeiro."
        exit 1
    fi

    echo ""
    echo -e "${BLUE}═══ Diagnóstico do sockmap ═══${NC}"
    echo ""

    # Lê stats antes
    CONNS_BEFORE=$(bpftool map lookup pinned "$BPF_PIN_DIR/stats" \
        key 0x00 0x00 0x00 0x00 --json 2>/dev/null \
        | python3 -c "import sys,json; d=json.load(sys.stdin); print(int.from_bytes(bytes(d['value']),'little'))" 2>/dev/null || echo 0)

    log_info "Detectando containers Docker em execução..."
    CONTAINERS=$(docker ps --format '{{.Names}} {{.NetworkSettings.Networks}}' 2>/dev/null | head -20 || echo "")

    # Testa conexões TCP para IPs conhecidos da rede 10.15.0.0/16
    log_info "Tentando conexões TCP internas para ativar o sock_ops..."
    # Pega IPs de containers ativos
    INTERNAL_IPS=$(docker inspect $(docker ps -q 2>/dev/null) 2>/dev/null \
        | python3 -c "
import sys, json
try:
    data = json.load(sys.stdin)
    ips = []
    for c in data:
        nets = c.get('NetworkSettings', {}).get('Networks', {})
        for net in nets.values():
            ip = net.get('IPAddress', '')
            if ip.startswith('10.'):
                name = c.get('Name', '?').strip('/')
                port_map = c.get('NetworkSettings', {}).get('Ports', {})
                port = 80
                for p in port_map:
                    try: port = int(p.split('/')[0]); break
                    except: pass
                ips.append(f'{ip}:{port}:{name}')
    print(' '.join(ips[:10]))
except Exception as e:
    print('')
" 2>/dev/null || echo "")

    if [ -z "$INTERNAL_IPS" ]; then
        log_warn "Nenhum container com IP 10.x detectado via Docker."
        log_info "Tentando IPs conhecidos da dockernet..."
        INTERNAL_IPS="10.15.0.45:80:nginx 10.15.0.42:3306:mariadb 10.15.0.54:5432:postgres 10.15.0.55:6379:redis"
    fi

    log_info "Testando conexões (timeout 1s cada):"
    for entry in $INTERNAL_IPS; do
        IP=$(echo $entry | cut -d: -f1)
        PORT=$(echo $entry | cut -d: -f2)
        NAME=$(echo $entry | cut -d: -f3)
        if timeout 1 bash -c "echo > /dev/tcp/$IP/$PORT" 2>/dev/null; then
            log_ok "  $NAME ($IP:$PORT) — conexão OK"
        else
            log_warn "  $NAME ($IP:$PORT) — sem resposta (normal se a porta estiver fechada)"
        fi
    done

    sleep 1

    # Lê stats depois
    CONNS_AFTER=$(bpftool map lookup pinned "$BPF_PIN_DIR/stats" \
        key 0x00 0x00 0x00 0x00 --json 2>/dev/null \
        | python3 -c "import sys,json; d=json.load(sys.stdin); print(int.from_bytes(bytes(d['value']),'little'))" 2>/dev/null || echo 0)

    NEW=$((CONNS_AFTER - CONNS_BEFORE))
    echo ""
    if [ "$NEW" -gt 0 ]; then
        log_ok "$NEW nova(s) conexão(oes) registradas no sockhash! ✅"
        log_ok "O sockmap está interceptando tráfego inter-container corretamente."
    else
        log_warn "Nenhuma nova conexão registrada."
        echo ""
        echo "Possíveis causas:"
        echo "  1. Containers não reiniciados após 'make load'"
        echo "     → Solucao: cd /home/docker && docker compose restart"
        echo "  2. A rede dockernet usa macvlan (não bridge)"
        echo "     → Macvlan não passa pelo cgroup host, mas sim pelo kernel de rede"
        echo "     → Neste caso o sockmap não intercepta (limitação de design)"
        echo "  3. Verificar o driver da rede:"
        echo "     docker network inspect dockernet | grep Driver"
    fi
    echo ""
}

# ─────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────
case "${1:-help}" in
    load)
        do_load
        ;;
    unload)
        do_unload
        ;;
    status)
        do_status
        ;;
    diagnose)
        do_diagnose
        ;;
    *)
        echo "eBPF Sockmap Loader"
        echo ""
        echo "Uso: sudo $0 {load|unload|status|diagnose}"
        echo ""
        echo "  load      - Carrega e ativa o sockmap"
        echo "  unload    - Remove o sockmap"
        echo "  status    - Mostra status e estatísticas"
        echo "  diagnose  - Testa se o sockmap está interceptando conexões"
        ;;
esac