#!/usr/bin/env python3

import os
import sys
import socket
import struct
import json
import subprocess
import argparse
import ipaddress
import time
import re

def run_cmd(cmd, show_output=False):
    """Executa um comando shell e retorna sua saída"""
    if show_output:
        print(f"Executando: {cmd}")
    
    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if result.returncode != 0:
        print(f"Erro ao executar: {cmd}")
        print(f"Saída de erro: {result.stderr}")
        return None
    return result.stdout.strip()

def find_map_ids_by_name(map_name):
    """Encontra todos os IDs de mapas BPF com um determinado nome"""
    output = run_cmd("bpftool map list")
    if not output:
        return []
    
    map_ids = []
    for line in output.split('\n'):
        if f"name {map_name}" in line:
            match = re.search(r'^(\d+):', line)
            if match:
                map_ids.append(match.group(1))
    
    return map_ids

def find_map_id(map_name):
    """Encontra o primeiro ID de um mapa BPF pelo nome"""
    map_ids = find_map_ids_by_name(map_name)
    return map_ids[0] if map_ids else None

def clean_old_maps():
    """Limpa todos os mapas antigos com nomes relacionados a NAT"""
    map_names = ["dnat_port_map", "snat_networks", "hairpin_map"]
    count = 0
    
    for map_name in map_names:
        map_ids = find_map_ids_by_name(map_name)
        for map_id in map_ids:
            print(f"Limpando mapa antigo: {map_name} (ID: {map_id})")
            run_cmd(f"bpftool map dump id {map_id} -j | grep -o '\"key\": \"[^\"]*\"' | "
                    f"sed 's/\"key\": \"\\([^\"]*\\)\"/\\1/g' | while read key; do "
                    f"[ ! -z \"$key\" ] && bpftool map delete id {map_id} key hex $key; done")
            count += 1
    
    return count

def load_tc_program(interface, object_file="tc_nat.o"):
    """Carrega programas TC para interface especificada"""
    print(f"Carregando programas TC para interface {interface}...")
    
    # Compilar programa BPF
    print("Compilando programa TC BPF...")
    cmd = f"clang -O2 -Wall -target bpf -g -D__TARGET_ARCH_arm64 -c tc_nat.c -o {object_file} -I/usr/include/bpf -I/usr/include/"
    if run_cmd(cmd) is None:
        return False
    
    # Remover programas existentes
    run_cmd(f"tc filter del dev {interface} ingress 2>/dev/null")
    run_cmd(f"tc filter del dev {interface} egress 2>/dev/null")
    run_cmd(f"tc qdisc del dev {interface} clsact 2>/dev/null")
    
    # Configurar qdisc
    run_cmd(f"tc qdisc add dev {interface} clsact")
    
    # Carregar programas TC
    ingress_cmd = f"tc filter add dev {interface} ingress prio 1 handle 1 bpf da obj {object_file} sec tc_ingress"
    egress_cmd = f"tc filter add dev {interface} egress prio 1 handle 1 bpf da obj {object_file} sec tc_egress"
    
    if run_cmd(ingress_cmd, show_output=True) is None or run_cmd(egress_cmd, show_output=True) is None:
        return False
    
    print("Programas TC carregados com sucesso!")
    print("Aguardando criação dos mapas...")
    time.sleep(1)
    return True

def setup_dnat_rule_hex(map_id, external_port, internal_ip, internal_port):
    """Configura uma regra DNAT usando strings hexadecimais diretamente"""
    # Converter porta para hex (little-endian)
    port_hex = f"{external_port & 0xff:02x} {(external_port >> 8) & 0xff:02x}"
    
    # Converter IP para hex (formato de rede)
    octets = list(map(int, internal_ip.split('.')))
    ip_hex = f"{octets[3]:02x} {octets[2]:02x} {octets[1]:02x} {octets[0]:02x}"
    
    # Converter porta interna para hex (little-endian)
    int_port_hex = f"{internal_port & 0xff:02x} {(internal_port >> 8) & 0xff:02x}"
    
    # Valor: IP (4 bytes) + porta (2 bytes) + padding (2 bytes)
    value_hex = f"{ip_hex} {int_port_hex} 00 00"
    
    # Update map usando bpftool com valores hex
    cmd = f"bpftool map update id {map_id} key hex {port_hex} value hex {value_hex}"
    success = run_cmd(cmd, show_output=True) is not None
    
    if success:
        print(f"Regra DNAT configurada: {external_port} -> {internal_ip}:{internal_port}")
    else:
        print(f"Falha ao configurar regra DNAT: {external_port} -> {internal_ip}:{internal_port}")
    
    return success

def setup_snat_network_hex(map_id, network):
    """Configura uma rede para SNAT usando strings hexadecimais diretamente"""
    # Parse da rede (exemplo: "10.15.0.0/16")
    ip, prefix = network.split('/')
    prefix = int(prefix)
    
    if prefix != 16:
        print(f"AVISO: Rede {network} não é /16. Apenas redes /16 são totalmente suportadas.")
    
    # Extrair octetos do IP
    octets = list(map(int, ip.split('.')))
    
    # Para /16, montamos o prefixo
    network_hex = f"00 00 {octets[1]:02x} {octets[0]:02x}"
    
    # Valor é 1 byte: habilitado
    value_hex = "01"  # 1 = habilitado
    
    # Update map usando bpftool com valores hex
    cmd = f"bpftool map update id {map_id} key hex {network_hex} value hex {value_hex}"
    success = run_cmd(cmd, show_output=True) is not None
    
    if success:
        print(f"Rede SNAT configurada: {network}")
    else:
        print(f"Falha ao configurar rede SNAT: {network}")
    
    return success

def clear_map(map_id):
    """Limpa todas as entradas de um mapa BPF usando comandos diretos"""
    print(f"Limpando mapa ID {map_id}...")
    
    # Verificar se o mapa existe
    if run_cmd(f"bpftool map show id {map_id}") is None:
        print(f"  Mapa ID {map_id} não encontrado")
        return False
    
    # Obter conteúdo do mapa
    map_content = run_cmd(f"bpftool map dump id {map_id}")
    if not map_content or "key:" not in map_content:
        print("  Mapa já está vazio")
        return True
    
    # Usar abordagem direta para limpar
    run_cmd(f"echo '{map_content}' | grep -o 'key: [0-9a-fx ]*' | sed 's/key: //g' | while read -r key; do [ ! -z \"$key\" ] && bpftool map delete id {map_id} key hex $key; done")
    
    # Verificar se o mapa está vazio
    map_content = run_cmd(f"bpftool map dump id {map_id}")
    if not map_content or "key:" not in map_content:
        print("  Mapa limpo com sucesso")
        return True
    else:
        print("  Não foi possível limpar completamente o mapa")
        return False

def main():
    parser = argparse.ArgumentParser(description="Configuração de NAT BPF para Docker")
    parser.add_argument("interface", help="Interface de rede externa (ex: enp0s6)")
    parser.add_argument("--target", default="10.15.0.45", help="IP do container alvo (ex: 10.15.0.45)")
    parser.add_argument("--networks", default="10.15.0.0/16,172.17.0.0/16",
                       help="Redes para SNAT, separadas por vírgulas (ex: 10.15.0.0/16,172.17.0.0/16)")
    parser.add_argument("--http-port", type=int, default=80, help="Porta HTTP para redirecionar")
    parser.add_argument("--https-port", type=int, default=443, help="Porta HTTPS para redirecionar")
    parser.add_argument("--reload", action="store_true", help="Recarregar o programa TC")
    parser.add_argument("--clean", action="store_true", help="Limpar mapas antigos antes de configurar")
    
    args = parser.parse_args()
    
    print("=== Configuração de NAT BPF para Docker ===")
    print(f"Interface: {args.interface}")
    print(f"Container alvo: {args.target}")
    print(f"Redes para SNAT: {args.networks}")
    print(f"Portas: HTTP={args.http_port}, HTTPS={args.https_port}")
    print("")
    
    # Limpar mapas antigos se solicitado
    if args.clean:
        print("Limpando mapas antigos...")
        count = clean_old_maps()
        print(f"Total de {count} mapas antigos foram limpos.")
    
    # Carregar programa TC se necessário
    if args.reload:
        if not load_tc_program(args.interface):
            print("Falha ao carregar programas TC. Abortando.")
            return 1
    
    # Encontrar IDs dos mapas
    dnat_map_id = find_map_id("dnat_port_map")
    snat_map_id = find_map_id("snat_networks")
    
    if not dnat_map_id or not snat_map_id:
        print("Mapas BPF não encontrados. Carregando programas TC...")
        if not load_tc_program(args.interface):
            print("Falha ao carregar programas TC. Abortando.")
            return 1
        
        # Tentar novamente após carregar os programas
        dnat_map_id = find_map_id("dnat_port_map")
        snat_map_id = find_map_id("snat_networks")
        
        if not dnat_map_id or not snat_map_id:
            print("Ainda não foi possível encontrar os mapas. Verifique o carregamento dos programas TC.")
            return 1
    
    print(f"IDs dos mapas: DNAT={dnat_map_id}, SNAT={snat_map_id}")
    
    # Limpar mapas existentes
    clear_map(dnat_map_id)
    clear_map(snat_map_id)
    
    # Configurar regras DNAT usando método hexadecimal direto
    setup_dnat_rule_hex(dnat_map_id, args.http_port, args.target, args.http_port)
    setup_dnat_rule_hex(dnat_map_id, args.https_port, args.target, args.https_port)
    
    # Configurar redes SNAT
    for network in args.networks.split(','):
        setup_snat_network_hex(snat_map_id, network.strip())
    
    # Habilitar IP forwarding
    print("Habilitando IP forwarding...")
    with open("/proc/sys/net/ipv4/ip_forward", "w") as f:
        f.write("1")
    
    # Verificar configuração final
    print("\nConfiguração final:")
    print("Regras DNAT:")
    os.system(f"bpftool map dump id {dnat_map_id}")
    print("\nRedes SNAT:")
    os.system(f"bpftool map dump id {snat_map_id}")
    
    print("\nNAT BPF configurado com sucesso!")
    print("Para ver logs em tempo real: cat /sys/kernel/debug/tracing/trace_pipe")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())