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

/* Definições para contornar erros de compilação */
#ifndef __section
# define __section(NAME)                  \
   __attribute__((section(NAME), used))
#endif

/* Estrutura para conexões NAT */
struct nat_conn {
    __u32 orig_src_ip;    // IP de origem original
    __u16 orig_src_port;  // Porta de origem original
    __u32 orig_dst_ip;    // IP de destino original
    __u16 orig_dst_port;  // Porta de destino original
    __u32 new_src_ip;     // Novo IP de origem (para SNAT)
    __u16 new_src_port;   // Nova porta de origem (para SNAT)
    __u32 new_dst_ip;     // Novo IP de destino (para DNAT)
    __u16 new_dst_port;   // Nova porta de destino (para DNAT)
    __u8 proto;           // Protocolo (TCP/UDP)
    __u8 state;           // Estado da conexão (1=ativo)
    __u64 last_update;    // Timestamp da última atualização
};

/* Chave para consulta rápida de DNAT em pacotes de entrada */
struct dnat_key {
    __u32 dst_ip;
    __u16 dst_port;
    __u8 proto;
};

/* Chave para busca de conexões de saída */
struct conn_key {
    __u32 src_ip;
    __u16 src_port;
    __u32 dst_ip;
    __u16 dst_port;
    __u8 proto;
};

/* Mapas BPF para NAT */
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 65536);
    __type(key, struct conn_key);
    __type(value, struct nat_conn);
} conn_track_map SEC(".maps");

/* Mapa para regras DNAT estáticas */
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 1024);
    __type(key, struct dnat_key);
    __type(value, struct nat_conn);
} dnat_rules SEC(".maps");

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

/* Mapa para estatísticas */
struct nat_stats {
    __u64 processed;   // Pacotes processados
    __u64 snat_hits;   // Acertos de SNAT
    __u64 dnat_hits;   // Acertos de DNAT
    __u64 new_conns;   // Novas conexões
    __u64 errors;      // Erros
};

struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, struct nat_stats);
} stats_map SEC(".maps");

/* Configuração para IP do gateway (usado para SNAT) */
struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, __u32);
} gateway_ip_map SEC(".maps");

/* Função para calcular IPv4 checksum */
static __always_inline __u16 ip_checksum(struct iphdr *ip) {
    __u32 csum = 0;
    __u16 *next_ip_u16 = (__u16 *)ip;
    __u16 pad_bytes = ip->ihl * 4;
    
    ip->check = 0;
    
    for (int i = 0; i < pad_bytes / 2; i++) {
        csum += *next_ip_u16++;
    }
    
    csum = (csum & 0xffff) + (csum >> 16);
    csum = (csum & 0xffff) + (csum >> 16);
    
    return ~csum;
}

/* Função para atualizar o checksum TCP/UDP após alteração de IP/porta */
static __always_inline void update_l4_checksum(__u16 *csum, __u32 old_addr, __u32 new_addr, __u16 old_port, __u16 new_port) {
    __u32 check = ~bpf_ntohs(*csum);
    
    // Ajustar para mudança de IP
    check += ~((old_addr >> 16) & 0xffff) + ((new_addr >> 16) & 0xffff);
    check += ~(old_addr & 0xffff) + (new_addr & 0xffff);
    
    // Ajustar para mudança de porta
    check += ~old_port & 0xffff;
    check += new_port;
    
    // Finalizar
    check = (check & 0xffff) + (check >> 16);
    check = (check & 0xffff) + (check >> 16);
    *csum = bpf_htons(~check);
}

/* Função para atualizar estatísticas */
static __always_inline void update_stats(__u32 stat_type) {
    __u32 key = 0;
    struct nat_stats *stats;
    
    stats = bpf_map_lookup_elem(&stats_map, &key);
    if (stats) {
        switch (stat_type) {
            case 0: stats->processed++; break;
            case 1: stats->snat_hits++; break;
            case 2: stats->dnat_hits++; break;
            case 3: stats->new_conns++; break;
            case 4: stats->errors++; break;
        }
    }
}

/* Função para tentar DNAT em pacotes de entrada */
static __always_inline int try_dnat(struct iphdr *ip, void *data_end, __u8 proto, __u16 dport, struct nat_conn *conn) {
    struct dnat_key key = {0};
    struct nat_conn *rule;
    
    key.dst_ip = ip->daddr;
    key.dst_port = dport;
    key.proto = proto;
    
    // Buscar regra DNAT
    rule = bpf_map_lookup_elem(&dnat_rules, &key);
    if (!rule) {
        return 0; // Nenhuma regra DNAT para este destino
    }
    
    // Preencher a conexão com os valores da regra DNAT
    conn->orig_dst_ip = ip->daddr;
    conn->orig_dst_port = dport;
    conn->new_dst_ip = rule->new_dst_ip;
    conn->new_dst_port = rule->new_dst_port;
    
    // Aplicar DNAT
    ip->daddr = conn->new_dst_ip;
    
    // Modificar a porta de destino dependendo do protocolo
    if (proto == IPPROTO_TCP) {
        struct tcphdr *tcp = (void *)(ip + 1);
        if ((void *)(tcp + 1) > data_end)
            return 0;
            
        tcp->dest = bpf_htons(conn->new_dst_port);
        
        // Atualizar checksum TCP
        if (tcp->check) {
            update_l4_checksum(&tcp->check, conn->orig_dst_ip, conn->new_dst_ip, 
                             bpf_htons(conn->orig_dst_port), bpf_htons(conn->new_dst_port));
        }
    } else if (proto == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if ((void *)(udp + 1) > data_end)
            return 0;
            
        udp->dest = bpf_htons(conn->new_dst_port);
        
        // Atualizar checksum UDP se não for zero
        if (udp->check) {
            update_l4_checksum(&udp->check, conn->orig_dst_ip, conn->new_dst_ip, 
                             bpf_htons(conn->orig_dst_port), bpf_htons(conn->new_dst_port));
        }
    }
    
    // Recalcular IP checksum
    ip->check = ip_checksum(ip);
    
    update_stats(2); // DNAT hit
    return 1;
}

/* Função para verificar e aplicar SNAT para redes internas */
static __always_inline int should_apply_snat(__u32 src_ip) {
    // Extrair o prefixo /16
    __u32 prefix = src_ip & 0xFFFF0000;
    __u8 *flag;
    
    flag = bpf_map_lookup_elem(&snat_networks, &prefix);
    if (flag && *flag == 1) {
        return 1;
    }
    
    return 0;
}

/* Programa principal XDP para aceleração de NAT */
SEC("xdp")
int xdp_accelerate_nat(struct xdp_md *ctx) {
    void *data = (void *)(long)ctx->data;
    void *data_end = (void *)(long)ctx->data_end;
    struct ethhdr *eth = data;
    
    // Verificar limites do frame
    if (data + sizeof(struct ethhdr) > data_end)
        return XDP_PASS;
    
    // Só processa IPv4
    if (eth->h_proto != bpf_htons(ETH_P_IP))
        return XDP_PASS;
    
    struct iphdr *ip = (void *)(eth + 1);
    if ((void *)(ip + 1) > data_end)
        return XDP_PASS;
    
    // Ignorar pacotes não-TCP/UDP por enquanto
    if (ip->protocol != IPPROTO_TCP && ip->protocol != IPPROTO_UDP)
        return XDP_PASS;
    
    update_stats(0); // Pacote processado
    
    __u16 src_port = 0;
    __u16 dst_port = 0;
    
    // Extrair portas base no protocolo
    if (ip->protocol == IPPROTO_TCP) {
        struct tcphdr *tcp = (void *)(ip + 1);
        if ((void *)(tcp + 1) > data_end)
            return XDP_PASS;
            
        src_port = bpf_ntohs(tcp->source);
        dst_port = bpf_ntohs(tcp->dest);
    } else if (ip->protocol == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if ((void *)(udp + 1) > data_end)
            return XDP_PASS;
            
        src_port = bpf_ntohs(udp->source);
        dst_port = bpf_ntohs(udp->dest);
    }
    
    struct conn_key key = {0};
    struct nat_conn conn = {0}, *conn_ptr;
    
    // Criar chave de conexão
    key.src_ip = ip->saddr;
    key.dst_ip = ip->daddr;
    key.src_port = src_port;
    key.dst_port = dst_port;
    key.proto = ip->protocol;
    
    // Procurar a conexão na tabela
    conn_ptr = bpf_map_lookup_elem(&conn_track_map, &key);
    
    if (conn_ptr) {
        // Conexão existente, atualizar timestamp
        conn_ptr->last_update = bpf_ktime_get_ns();
    } else {
        // Nova conexão potencial
        __u32 gateway_key = 0;
        __u32 *gateway_ip;
        
        // Preencher conexão original
        conn.orig_src_ip = ip->saddr;
        conn.orig_src_port = src_port;
        conn.orig_dst_ip = ip->daddr;
        conn.orig_dst_port = dst_port;
        conn.proto = ip->protocol;
        conn.state = 1;
        conn.last_update = bpf_ktime_get_ns();
        
        // Verificar se é um pacote de entrada que precisa de DNAT
        if (try_dnat(ip, data_end, ip->protocol, dst_port, &conn)) {
            // Atualizar campos de conexão após DNAT
            key.dst_ip = conn.new_dst_ip;
            key.dst_port = conn.new_dst_port;
        } 
        // Verificar se é um pacote de saída que precisa de SNAT
        else if (should_apply_snat(ip->saddr)) {
            gateway_ip = bpf_map_lookup_elem(&gateway_ip_map, &gateway_key);
            if (!gateway_ip) {
                update_stats(4); // Erro - gateway IP não configurado
                return XDP_PASS; // Sem IP gateway, não pode fazer SNAT
            }
            
            conn.new_src_ip = *gateway_ip;
            conn.new_src_port = src_port; // Usar mesma porta por enquanto
            
            // Aplicar SNAT
            __u32 old_ip = ip->saddr;
            ip->saddr = conn.new_src_ip;
            
            // Atualizar checksum de nível 4
            if (ip->protocol == IPPROTO_TCP) {
                struct tcphdr *tcp = (void *)(ip + 1);
                if ((void *)(tcp + 1) > data_end)
                    return XDP_PASS;
                    
                if (tcp->check) {
                    update_l4_checksum(&tcp->check, old_ip, conn.new_src_ip, 
                                    bpf_htons(src_port), bpf_htons(conn.new_src_port));
                }
            } else if (ip->protocol == IPPROTO_UDP) {
                struct udphdr *udp = (void *)(ip + 1);
                if ((void *)(udp + 1) > data_end)
                    return XDP_PASS;
                    
                if (udp->check) {
                    update_l4_checksum(&udp->check, old_ip, conn.new_src_ip, 
                                    bpf_htons(src_port), bpf_htons(conn.new_src_port));
                }
            }
            
            // Recalcular IP checksum
            ip->check = ip_checksum(ip);
            
            update_stats(1); // SNAT hit
        } else {
            // Pacote normal sem NAT, apenas passa
            return XDP_PASS;
        }
        
        // Criar conexão reversa para resposta
        struct conn_key reverse_key = {0};
        struct nat_conn reverse_conn = {0};
        
        // Configurar chave reversa
        reverse_key.src_ip = conn.new_dst_ip ? conn.new_dst_ip : ip->daddr;
        reverse_key.src_port = conn.new_dst_port ? conn.new_dst_port : dst_port;
        reverse_key.dst_ip = conn.new_src_ip ? conn.new_src_ip : ip->saddr;
        reverse_key.dst_port = conn.new_src_port ? conn.new_src_port : src_port;
        reverse_key.proto = ip->protocol;
        
        // Configurar conexão reversa
        reverse_conn.orig_src_ip = reverse_key.src_ip;
        reverse_conn.orig_src_port = reverse_key.src_port;
        reverse_conn.orig_dst_ip = reverse_key.dst_ip;
        reverse_conn.orig_dst_port = reverse_key.dst_port;
        reverse_conn.new_src_ip = conn.orig_dst_ip;
        reverse_conn.new_src_port = conn.orig_dst_port;
        reverse_conn.new_dst_ip = conn.orig_src_ip;
        reverse_conn.new_dst_port = conn.orig_src_port;
        reverse_conn.proto = ip->protocol;
        reverse_conn.state = 1;
        reverse_conn.last_update = conn.last_update;
        
        // Inserir ambas as conexões (original e reversa) no mapa
        bpf_map_update_elem(&conn_track_map, &key, &conn, BPF_ANY);
        bpf_map_update_elem(&conn_track_map, &reverse_key, &reverse_conn, BPF_ANY);
        
        update_stats(3); // Nova conexão
    }
    
    return XDP_PASS;
}

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