#!/usr/bin/env python3

import os
import sys
import time
import socket
import struct
import argparse
import subprocess
import ipaddress
from bcc import BPF

# Constantes
DEFAULT_INTERFACE = "enp0s6"
DEFAULT_GATEWAY_IP = "10.0.0.87"
DNAT_RULES_FILE = "dnat_rules.txt"
SNAT_NETWORKS_FILE = "snat_networks.txt"

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 ip_to_int(ip_str):
    """Converte um endereço IP em inteiro."""
    return struct.unpack("!I", socket.inet_aton(ip_str))[0]

def int_to_ip(ip_int):
    """Converte um inteiro em endereço IP."""
    return socket.inet_ntoa(struct.pack("!I", ip_int))

def compile_xdp_program():
    """Compila o programa XDP."""
    print("Compilando o programa XDP...")
    cmd = "clang -O2 -target bpf -c xdp_accelerated_nat.c -o xdp_accelerated_nat.o"
    return run_cmd(cmd, True) is not None

def attach_xdp_program(interface, program_path):
    """Anexa o programa XDP à interface."""
    print(f"Anexando programa XDP à interface {interface}...")
    cmd = f"ip link set dev {interface} xdp obj {program_path} section xdp"
    return run_cmd(cmd, True) is not None

def detach_xdp_program(interface):
    """Remove o programa XDP da interface."""
    print(f"Removendo programa XDP da interface {interface}...")
    cmd = f"ip link set dev {interface} xdp off"
    return run_cmd(cmd, True) is not None

def load_dnat_rules(file_path):
    """Carrega regras DNAT de um arquivo."""
    rules = []
    
    try:
        with open(file_path, 'r') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                    
                # Formato: proto dst_ip:dst_port new_ip:new_port
                # Exemplo: tcp 10.0.0.87:80 10.15.0.45:80
                parts = line.split()
                if len(parts) != 3:
                    print(f"Formato inválido na linha: {line}")
                    continue
                    
                proto_name, dst, new = parts
                
                # Parse proto
                if proto_name.lower() == "tcp":
                    proto = socket.IPPROTO_TCP
                elif proto_name.lower() == "udp":
                    proto = socket.IPPROTO_UDP
                else:
                    print(f"Protocolo inválido: {proto_name}")
                    continue
                
                # Parse IP e portas
                try:
                    dst_ip, dst_port = dst.split(':')
                    new_ip, new_port = new.split(':')
                    
                    rule = {
                        'proto': proto,
                        'dst_ip': dst_ip,
                        'dst_port': int(dst_port),
                        'new_ip': new_ip,
                        'new_port': int(new_port)
                    }
                    rules.append(rule)
                except ValueError:
                    print(f"Erro ao analisar IP/porta na linha: {line}")
                    continue
                    
        print(f"Carregadas {len(rules)} regras DNAT")
        return rules
    except FileNotFoundError:
        print(f"Arquivo de regras DNAT não encontrado: {file_path}")
        return []

def load_snat_networks(file_path):
    """Carrega redes SNAT de um arquivo."""
    networks = []
    
    try:
        with open(file_path, 'r') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                    
                # Formato: rede/prefixo
                # Exemplo: 10.15.0.0/16
                try:
                    network = ipaddress.ip_network(line)
                    networks.append(str(network))
                except ValueError:
                    print(f"Rede inválida na linha: {line}")
                    continue
                    
        print(f"Carregadas {len(networks)} redes SNAT")
        return networks
    except FileNotFoundError:
        print(f"Arquivo de redes SNAT não encontrado: {file_path}")
        return []

def update_gateway_ip(bpf_obj, gateway_ip):
    """Atualiza o IP do gateway no mapa BPF."""
    try:
        ip_int = ip_to_int(gateway_ip)
        key = ctypes.c_uint32(0)
        value = ctypes.c_uint32(ip_int)
        bpf_obj["gateway_ip_map"][key] = value
        print(f"IP do gateway configurado: {gateway_ip}")
    except Exception as e:
        print(f"Erro ao configurar IP do gateway: {e}")

def update_dnat_rules(bpf_obj, rules):
    """Atualiza regras DNAT no mapa BPF."""
    try:
        dnat_map = bpf_obj["dnat_rules"]
        
        # Limpar mapa existente
        for k in list(dnat_map.keys()):
            del dnat_map[k]
        
        # Inserir novas regras
        for rule in rules:
            key = dnat_map.Key()
            key.dst_ip = ip_to_int(rule['dst_ip'])
            key.dst_port = rule['dst_port']
            key.proto = rule['proto']
            
            value = dnat_map.Value()
            value.new_dst_ip = ip_to_int(rule['new_ip'])
            value.new_dst_port = rule['new_port']
            
            dnat_map[key] = value
            
        print(f"{len(rules)} regras DNAT atualizadas")
    except Exception as e:
        print(f"Erro ao atualizar regras DNAT: {e}")

def update_snat_networks(bpf_obj, networks):
    """Atualiza redes SNAT no mapa BPF."""
    try:
        snat_map = bpf_obj["snat_networks"]
        
        # Limpar mapa existente
        for k in list(snat_map.keys()):
            del snat_map[k]
        
        # Inserir novas redes
        for net_str in networks:
            network = ipaddress.ip_network(net_str)
            if network.prefixlen != 16:
                print(f"Aviso: Rede {net_str} não é /16. Pode não funcionar corretamente.")
            
            # Extrair prefixo /16
            prefix = int(network.network_address) & 0xFFFF0000
            key = ctypes.c_uint32(prefix)
            value = ctypes.c_uint8(1)
            
            snat_map[key] = value
            
        print(f"{len(networks)} redes SNAT atualizadas")
    except Exception as e:
        print(f"Erro ao atualizar redes SNAT: {e}")

def print_stats(bpf_obj):
    """Imprime estatísticas de NAT."""
    stats = bpf_obj["stats_map"][0]
    print("\nEstatísticas de NAT:")
    print(f"  Pacotes processados: {stats.processed}")
    print(f"  Acertos de SNAT: {stats.snat_hits}")
    print(f"  Acertos de DNAT: {stats.dnat_hits}")
    print(f"  Novas conexões: {stats.new_conns}")
    print(f"  Erros: {stats.errors}")

def print_connections(bpf_obj):
    """Imprime conexões ativas."""
    conn_map = bpf_obj["conn_track_map"]
    print("\nConexões ativas:")
    
    count = 0
    for k, v in conn_map.items():
        proto_name = "TCP" if k.proto == socket.IPPROTO_TCP else "UDP"
        src_ip = int_to_ip(k.src_ip)
        dst_ip = int_to_ip(k.dst_ip)
        
        # Mostrar apenas as 20 primeiras conexões
        if count < 20:
            print(f"  {proto_name}: {src_ip}:{k.src_port} -> {dst_ip}:{k.dst_port}")
        count += 1
        
    if count > 20:
        print(f"  ... e mais {count - 20} conexões")
    print(f"  Total: {count} conexões ativas")

def main():
    parser = argparse.ArgumentParser(description='Ferramenta para configuração de aceleração NAT com XDP')
    parser.add_argument('-i', '--interface', default=DEFAULT_INTERFACE, help=f'Interface de rede (padrão: {DEFAULT_INTERFACE})')
    parser.add_argument('-g', '--gateway', default=DEFAULT_GATEWAY_IP, help=f'IP do gateway para SNAT (padrão: {DEFAULT_GATEWAY_IP})')
    parser.add_argument('--compile', action='store_true', help='Compilar e carregar o programa XDP')
    parser.add_argument('--unload', action='store_true', help='Descarregar o programa XDP')
    parser.add_argument('--stats', action='store_true', help='Mostrar estatísticas')
    parser.add_argument('--monitor', action='store_true', help='Monitorar estatísticas continuamente')
    
    args = parser.parse_args()
    
    # Descarregar programa
    if args.unload:
        detach_xdp_program(args.interface)
        sys.exit(0)
    
    # Compilar e carregar programa
    if args.compile:
        if not compile_xdp_program():
            print("Falha ao compilar o programa XDP")
            sys.exit(1)
            
        # Desanexar qualquer programa existente
        detach_xdp_program(args.interface)
        
        # Anexar o novo programa
        if not attach_xdp_program(args.interface, "xdp_accelerated_nat.o"):
            print("Falha ao anexar o programa XDP à interface")
            sys.exit(1)
    
    # Carregar o objeto BPF para interagir com o programa carregado
    try:
        b = BPF(obj="xdp_accelerated_nat.o")
    except Exception as e:
        print(f"Erro ao carregar o objeto BPF: {e}")
        sys.exit(1)
    
    # Carregar regras DNAT e redes SNAT
    dnat_rules = load_dnat_rules(DNAT_RULES_FILE)
    snat_networks = load_snat_networks(SNAT_NETWORKS_FILE)
    
    # Atualizar configurações
    update_gateway_ip(b, args.gateway)
    if dnat_rules:
        update_dnat_rules(b, dnat_rules)
    if snat_networks:
        update_snat_networks(b, snat_networks)
    
    # Mostrar estatísticas
    if args.stats:
        print_stats(b)
        print_connections(b)
    
    # Monitorar estatísticas
    if args.monitor:
        try:
            print("Monitorando estatísticas (Ctrl+C para sair)...")
            while True:
                os.system('clear')
                print_stats(b)
                print_connections(b)
                time.sleep(1)
        except KeyboardInterrupt:
            print("\nMonitoramento interrompido.")
    
    print("\nPrograma XDP NAT acelerado em execução.")
    print("Use '--stats' para visualizar estatísticas.")
    print("Use '--unload' para descarregar o programa.")

if __name__ == "__main__":
    main()