#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/pkt_cls.h>
#include <linux/in.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

// Adicionar definição para ARM64
#ifndef __TARGET_ARCH_arm64
#define __TARGET_ARCH_arm64
#endif

// Estrutura para regras DNAT
struct dnat_rule {
    __u32 new_ip;      // IP destino após NAT
    __u16 new_port;    // Porta destino após NAT
};

// Mapa para regras DNAT baseadas na porta
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 64);
    __type(key, __u16);             // Porta destino
    __type(value, struct dnat_rule); // Regra DNAT
} dnat_port_map SEC(".maps");

// Mapa para definir redes para SNAT (masquerade)
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 16);
    __type(key, __u32);   // Prefixo de rede (/16)
    __type(value, __u8);  // Flag (1=ativo)
} snat_networks SEC(".maps");

// Mapa para armazenar configurações de hairpinning
struct hairpin_key {
    __u32 dst_ip;     // IP destino
    __u16 dst_port;   // Porta destino
};

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 64);
    __type(key, struct hairpin_key);  // IP e porta destino
    __type(value, __u8);              // Flag (1=ativo)
} hairpin_map SEC(".maps");

// Função para calcular checksum com bounds checking
static inline __u16 compute_ip_checksum(struct iphdr *iph) {
    __u32 csum = 0;
    __u16 *next_iph_u16 = (__u16 *)iph;
    
    // IP header sem o campo de checksum (5 words, 10 bytes)
#pragma unroll
    for (int i = 0; i < 10; i++) {
        csum += *next_iph_u16++;
    }
    
    // Fold 32-bit to 16-bit
    csum = (csum & 0xffff) + (csum >> 16);
    csum = (csum & 0xffff) + (csum >> 16);
    
    return ~csum;
}

// Gancho TC para INGRESS (DNAT)
SEC("tc_ingress")
int nat_ingress_func(struct __sk_buff *skb) {
    // Cabeçalhos
    void *data = (void *)(long)skb->data;
    void *data_end = (void *)(long)skb->data_end;
    
    // Validação de ethernet
    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end)
        return TC_ACT_OK;
    
    // Verificar se é IPv4
    if (eth->h_proto != bpf_htons(ETH_P_IP))
        return TC_ACT_OK;
    
    // Validação de cabeçalho IP
    struct iphdr *ip = (void *)(eth + 1);
    if ((void *)(ip + 1) > data_end)
        return TC_ACT_OK;
    
    // Verificar tamanho do cabeçalho IP
    if (ip->ihl != 5) // Apenas cabeçalho padrão sem opções
        return TC_ACT_OK;
    
    // Depuração - imprimir informações do pacote
    bpf_printk("INGRESS: IP Source %pI4, Dest %pI4", &ip->saddr, &ip->daddr);
    
    // DNAT para TCP
    if (ip->protocol == IPPROTO_TCP) {
        struct tcphdr *tcp = (void *)(ip + 1);
        if ((void *)(tcp + 1) > data_end)
            return TC_ACT_OK;
        
        __u16 dport = bpf_ntohs(tcp->dest);
        bpf_printk("INGRESS: TCP port %d", dport);
        
        // Procurar regra para esta porta
        struct dnat_rule *rule = bpf_map_lookup_elem(&dnat_port_map, &dport);
        if (rule) {
            bpf_printk("Regra encontrada: %d -> %pI4:%d", 
                      dport, &rule->new_ip, rule->new_port);
            
            // Guardar o IP e porta originais
            __u32 old_ip = ip->daddr;
            __u16 old_port = tcp->dest;
            
            // Aplicar a regra DNAT
            ip->daddr = rule->new_ip;
            tcp->dest = bpf_htons(rule->new_port);
            
            // Recalcular checksum IP
            ip->check = 0;
            ip->check = compute_ip_checksum(ip);
            
            // Recalcular checksum TCP (parcialmente)
            if (tcp->check) {
                // Primeiro ajustar para o IP de destino alterado
                __u32 check = ~bpf_ntohs(tcp->check);
                check += ~bpf_ntohs(old_port) & 0xffff;
                check += bpf_ntohs(tcp->dest);
                check += ~((old_ip >> 16) & 0xffff) + ~(old_ip & 0xffff);
                check += ((ip->daddr >> 16) & 0xffff) + (ip->daddr & 0xffff);
                
                // Dobrar para obter o checksum final
                check = (check & 0xffff) + (check >> 16);
                check = (check & 0xffff) + (check >> 16);
                tcp->check = bpf_htons(~check);
            }
            
            bpf_printk("DNAT aplicado: %pI4:%d -> %pI4:%d",
                     &old_ip, bpf_ntohs(old_port), 
                     &ip->daddr, bpf_ntohs(tcp->dest));
        }
    }
    // DNAT para UDP
    else if (ip->protocol == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if ((void *)(udp + 1) > data_end)
            return TC_ACT_OK;
        
        __u16 dport = bpf_ntohs(udp->dest);
        bpf_printk("INGRESS: UDP port %d", dport);
        
        // Procurar regra para esta porta
        struct dnat_rule *rule = bpf_map_lookup_elem(&dnat_port_map, &dport);
        if (rule) {
            bpf_printk("Regra encontrada: %d -> %pI4:%d", 
                      dport, &rule->new_ip, rule->new_port);
            
            // Guardar o IP e porta originais
            __u32 old_ip = ip->daddr;
            __u16 old_port = udp->dest;
            
            // Aplicar a regra DNAT
            ip->daddr = rule->new_ip;
            udp->dest = bpf_htons(rule->new_port);
            
            // Recalcular checksum IP
            ip->check = 0;
            ip->check = compute_ip_checksum(ip);
            
            // Recalcular checksum UDP (caso exista)
            if (udp->check) {
                // Ajuste similar ao do TCP
                __u32 check = ~bpf_ntohs(udp->check);
                check += ~bpf_ntohs(old_port) & 0xffff;
                check += bpf_ntohs(udp->dest);
                check += ~((old_ip >> 16) & 0xffff) + ~(old_ip & 0xffff);
                check += ((ip->daddr >> 16) & 0xffff) + (ip->daddr & 0xffff);
                
                // Dobrar para obter o checksum final
                check = (check & 0xffff) + (check >> 16);
                check = (check & 0xffff) + (check >> 16);
                udp->check = bpf_htons(~check);
            }
            
            bpf_printk("DNAT aplicado: %pI4:%d -> %pI4:%d",
                     &old_ip, bpf_ntohs(old_port), 
                     &ip->daddr, bpf_ntohs(udp->dest));
        }
    }
    
    return TC_ACT_OK;
}

// Gancho TC para EGRESS (SNAT/Masquerade)
SEC("tc_egress")
int nat_egress_func(struct __sk_buff *skb) {
    // Cabeçalhos
    void *data = (void *)(long)skb->data;
    void *data_end = (void *)(long)skb->data_end;
    
    // Validação de ethernet
    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end)
        return TC_ACT_OK;
    
    // Verificar se é IPv4
    if (eth->h_proto != bpf_htons(ETH_P_IP))
        return TC_ACT_OK;
    
    // Validação de cabeçalho IP
    struct iphdr *ip = (void *)(eth + 1);
    if ((void *)(ip + 1) > data_end)
        return TC_ACT_OK;
    
    // Depuração - imprimir informações do pacote
    bpf_printk("EGRESS: IP Source %pI4, Dest %pI4", &ip->saddr, &ip->daddr);
    
    // Verificar se o IP de origem está em nossas redes para SNAT/MASQUERADE
    __u32 src_prefix = ip->saddr & 0xFFFF0000; // /16
    __u8 *do_snat = bpf_map_lookup_elem(&snat_networks, &src_prefix);
    
    // Processar SNAT/Masquerade apenas para pacotes de redes internas
    if (do_snat && *do_snat) {
        // Guardar o IP original
        __u32 old_ip = ip->saddr;
        
        // IP fixo para SNAT: 10.0.0.87
        __u32 new_ip = 0x5700000A; // 10.0.0.87 em ordem de rede
        
        bpf_printk("SNAT: Aplicando masquerade para %pI4 -> %pI4", 
                  &old_ip, &new_ip);
        
        // Aplicar SNAT
        ip->saddr = new_ip;
        
        // Recalcular checksum IP
        ip->check = 0;
        ip->check = compute_ip_checksum(ip);
        
        // Recalcular checksums L4 (TCP ou UDP)
        if (ip->protocol == IPPROTO_TCP) {
            struct tcphdr *tcp = (void *)(ip + 1);
            if ((void *)(tcp + 1) <= data_end && tcp->check) {
                // Ajustar checksum TCP para o IP de origem alterado
                __u32 check = ~bpf_ntohs(tcp->check);
                check += ~((old_ip >> 16) & 0xffff) + ~(old_ip & 0xffff);
                check += ((new_ip >> 16) & 0xffff) + (new_ip & 0xffff);
                
                // Dobrar para obter o checksum final
                check = (check & 0xffff) + (check >> 16);
                check = (check & 0xffff) + (check >> 16);
                tcp->check = bpf_htons(~check);
                
                bpf_printk("SNAT TCP: porta %d, checksum atualizado", 
                          bpf_ntohs(tcp->source));
            }
        } 
        else if (ip->protocol == IPPROTO_UDP) {
            struct udphdr *udp = (void *)(ip + 1);
            if ((void *)(udp + 1) <= data_end && udp->check) {
                // Ajustar checksum UDP para o IP de origem alterado
                __u32 check = ~bpf_ntohs(udp->check);
                check += ~((old_ip >> 16) & 0xffff) + ~(old_ip & 0xffff);
                check += ((new_ip >> 16) & 0xffff) + (new_ip & 0xffff);
                
                // Dobrar para obter o checksum final
                check = (check & 0xffff) + (check >> 16);
                check = (check & 0xffff) + (check >> 16);
                udp->check = bpf_htons(~check);
                
                bpf_printk("SNAT UDP: porta %d, checksum atualizado", 
                          bpf_ntohs(udp->source));
            }
        }
        
        bpf_printk("SNAT concluído: %pI4 -> %pI4", &old_ip, &ip->saddr);
    }
    
    return TC_ACT_OK;
}

char __license[] SEC("license") = "GPL";