#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>

// Otimização: melhorar constantes para ajuste fino
#define MAX_HIGH_PRIORITY_PACKETS_PER_SECOND 2000
#define MAX_LOW_PRIORITY_PACKETS_PER_SECOND 500
#define MAX_PACKETS_PER_IP 60000                    
#define MAX_ICMP_PACKETS_PER_IP 50                
#define MAX_SYN_PACKETS_PER_IP 100                
#define TIMEOUT_SEC 1
#define PACKET_COST 1                // Custo padrão por pacote
#define HIGH_PRIO_TOKEN_RATE 2000    // Tokens/segundo para alta prioridade
#define LOW_PRIO_TOKEN_RATE 500      // Tokens/segundo para baixa prioridade
#define MAX_TOKENS 2000              // Máximo de tokens acumulados

// Otimização: estrutura de dados mais compacta para reduzir acessos de memória
struct ip_stats {
    __u32 pkt_count;          // Contador de pacotes total
    __u32 last_seen;          // Último timestamp em segundos (não nanossegundos)
    __u16 icmp_count;         // Contador ICMP (reduzido para 16 bits)
    __u16 syn_count;          // Contador SYN (reduzido para 16 bits)
    __u16 high_prio_tokens;   // Tokens disponíveis para alta prioridade
    __u16 low_prio_tokens;    // Tokens disponíveis para baixa prioridade
} __attribute__((packed));

// Otimização: usar PERCPU map para evitar contenção e locks
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
    __uint(max_entries, 100000);
    __type(key, __u32);           // IP address
    __type(value, struct ip_stats);
} ddos_map SEC(".maps");

// Otimização: usar LRU para blacklist para evitar crescimento infinito
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 100000);
    __type(key, __u32);
    __type(value, __u32);
} blacklist_map SEC(".maps");

// Mapa para interface de redirecionamento
struct {
    __uint(type, BPF_MAP_TYPE_XSKMAP);
    __uint(max_entries, 1);
    __type(key, int);
    __type(value, int);
} xsks_map SEC(".maps");

// Otimização: usar lista de permissões fixa e compacta
static const __u32 whitelist[] = {
    0x0A000001,  // 10.0.0.1
    0x0A630101,  // 10.99.1.1
    0x0a630102,  // 10.99.1.2
    0x0a630104   // 10.99.1.4
};

// Otimização: função inline para obter o timestamp em segundos em vez de nanossegundos
static __always_inline __u32 get_time_sec(void) {
    return bpf_ktime_get_ns() / 1000000000;
}

// Otimização: usar always_inline e __builtin_expect para funções críticas
static __always_inline int is_whitelisted(__u32 ip) {
    #pragma unroll
    for (int i = 0; i < sizeof(whitelist) / sizeof(whitelist[0]); i++) {
        if (ip == whitelist[i]) {
            return 1;
        }
    }
    return 0;
}

static __always_inline int is_blacklisted(__u32 ip) {
    return bpf_map_lookup_elem(&blacklist_map, &ip) != NULL;
}

// Otimização: usar switch-case em vez de verificações múltiplas para portas
static __always_inline int get_port_priority(__u16 port) {
    switch (port) {
        case 80:    // HTTP
        case 443:   // HTTPS
        case 32100: // Serviço de alta prioridade 1
        case 32400: // Serviço de alta prioridade 2
            return 2;  // Alta prioridade
        case 22:    // SSH
        case 25:    // SMTP
        case 53:    // DNS
            return 1;  // Média prioridade
        default:
            return 0;  // Baixa prioridade
    }
}

// Otimização: função de verificação de porta mais eficiente
static __always_inline int is_filtered_port(__u16 port) {
    // Portas importantes usam verificação direta
    switch (port) {
        case 8117:
        case 8085:
        case 8086: 
        case 32100:
        case 32400:
            return 1;
        default:
            // Verificação de intervalo apenas para outras portas
            return (port <= 1024);
    }
}

// Programa principal XDP
SEC("xdp")
int xdp_ddos_filter(struct xdp_md *ctx) {
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;
    
    // Otimização: early return para pacotes muito pequenos
    if (__builtin_expect(data + sizeof(struct ethhdr) + sizeof(struct iphdr) > data_end, 0))
        return XDP_PASS;

    struct ethhdr *eth = data;

    // Otimização: early return para pacotes não-IP
    if (__builtin_expect(eth->h_proto != __constant_htons(ETH_P_IP), 0))
        return XDP_PASS;

    struct iphdr *ip = (void *)(eth + 1);
    
    // Otimização: verificar o cabeçalho IP de forma mais eficiente
    if (__builtin_expect((void *)(ip + 1) > data_end || 
                         ip->ihl != 5 || 
                         ip->version != 4, 0))
        return XDP_PASS;

    // Otimização: early return para IPs na whitelist
    __u32 src_ip = ip->saddr;
    if (__builtin_expect(is_whitelisted(src_ip), 0)) {
        return XDP_PASS;
    }

    // Otimização: early drop para IPs bloqueados
    if (__builtin_expect(is_blacklisted(src_ip), 0)) {
        return XDP_DROP;
    }

    // Identificar protocolo e porta para classificação
    __u16 dest_port = 0;
    int is_syn = 0;
    int is_icmp = 0;

    if (ip->protocol == IPPROTO_TCP) {
        struct tcphdr *tcp = (void *)(ip + 1);
        if (__builtin_expect((void *)(tcp + 1) > data_end, 0))
            return XDP_PASS;
        
        dest_port = __constant_ntohs(tcp->dest);
        is_syn = tcp->syn && !tcp->ack;
        
    } else if (ip->protocol == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if (__builtin_expect((void *)(udp + 1) > data_end, 0))
            return XDP_PASS;
        
        dest_port = __constant_ntohs(udp->dest);
        
    } else if (ip->protocol == IPPROTO_ICMP) {
        is_icmp = 1;
    } else {
        // Outros protocolos passam sem filtro adicional
        return XDP_PASS;
    }

    // Otimização: early return se não for para portas filtradas
    if (__builtin_expect(!is_filtered_port(dest_port) && !is_icmp, 1))
        return XDP_PASS;

    // Prioridade do pacote com base na porta
    int priority = get_port_priority(dest_port);
    
    // Obter estatísticas atuais do IP
    struct ip_stats *entry = bpf_map_lookup_elem(&ddos_map, &src_ip);
    __u32 now = get_time_sec();
    
    if (entry) {
        // Otimização: cálculo de tempo decorrido mais eficiente
        __u32 elapsed = 0;
        if (__builtin_expect(now > entry->last_seen, 1)) {
            elapsed = now - entry->last_seen;
        }

        // Implementação de token bucket para cada prioridade
        if (elapsed > 0) {
            // Adicionar tokens baseado no tempo decorrido, limitado ao máximo
            __u16 new_high_tokens = entry->high_prio_tokens + (HIGH_PRIO_TOKEN_RATE * elapsed);
            entry->high_prio_tokens = new_high_tokens > MAX_TOKENS ? MAX_TOKENS : new_high_tokens;
            
            __u16 new_low_tokens = entry->low_prio_tokens + (LOW_PRIO_TOKEN_RATE * elapsed);
            entry->low_prio_tokens = new_low_tokens > MAX_TOKENS ? MAX_TOKENS : new_low_tokens;
            
            // Resetar contadores específicos de protocolo após 1 segundo
            if (elapsed >= TIMEOUT_SEC) {
                entry->icmp_count = 0;
                entry->syn_count = 0;
            }
        }

        // Atualizar o último timestamp visto
        entry->last_seen = now;
        
        // Atualizar contadores específicos
        entry->pkt_count++;
        
        if (is_icmp) {
            entry->icmp_count++;
            if (__builtin_expect(entry->icmp_count > MAX_ICMP_PACKETS_PER_IP, 0)) {
                __u32 one = 1;
                bpf_map_update_elem(&blacklist_map, &src_ip, &one, BPF_ANY);
                return XDP_DROP;
            }
        }
        
        if (is_syn) {
            entry->syn_count++;
            if (__builtin_expect(entry->syn_count > MAX_SYN_PACKETS_PER_IP, 0)) {
                __u32 one = 1;
                bpf_map_update_elem(&blacklist_map, &src_ip, &one, BPF_ANY);
                return XDP_DROP;
            }
        }

        // Otimização: verificação total de pacotes
        if (__builtin_expect(entry->pkt_count > MAX_PACKETS_PER_IP, 0)) {
            __u32 one = 1;
            bpf_map_update_elem(&blacklist_map, &src_ip, &one, BPF_ANY);
            return XDP_DROP;
        }

        // Verificar tokens disponíveis com base na prioridade
        if (priority == 2) {  // Alta prioridade
            if (entry->high_prio_tokens >= PACKET_COST) {
                entry->high_prio_tokens -= PACKET_COST;
                return XDP_PASS;
            }
        } else {  // Média e baixa prioridade
            if (entry->low_prio_tokens >= PACKET_COST) {
                entry->low_prio_tokens -= PACKET_COST;
                return XDP_PASS;
            }
        }
        
        // Sem tokens disponíveis para esta prioridade, descartar
        return XDP_DROP;
        
    } else {
        // Primeira vez que vemos este IP, inicializar com valores generosos
        struct ip_stats new_entry = {
            .pkt_count = 1,
            .last_seen = now,
            .icmp_count = is_icmp ? 1 : 0,
            .syn_count = is_syn ? 1 : 0,
            .high_prio_tokens = MAX_TOKENS - PACKET_COST,
            .low_prio_tokens = MAX_TOKENS - PACKET_COST
        };
        
        bpf_map_update_elem(&ddos_map, &src_ip, &new_entry, BPF_ANY);
        return XDP_PASS;
    }
}

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