#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 <linux/icmp.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

// Otimização: constantes ajustadas para seu ambiente específico
#define MAX_HIGH_PRIORITY_PACKETS_PER_SECOND 3000  // Aumentado para serviços críticos (HTTP/HTTPS)
#define MAX_LOW_PRIORITY_PACKETS_PER_SECOND 500
#define MAX_PACKETS_PER_IP 60000
#define MAX_ICMP_PACKETS_PER_IP 30                 // Reduzido conforme política ICMP restritiva
#define MAX_SYN_PACKETS_PER_IP 80                  // Reduzido para proteção preventiva
#define TIMEOUT_SEC 1
#define PACKET_COST 1
#define HIGH_PRIO_TOKEN_RATE 3000                  // Aumentado para combinar com QoS HTB
#define LOW_PRIO_TOKEN_RATE 500
#define MAX_TOKENS 3000
#define DOCKER_NETWORK_PREFIX 0x0A0F0000           // 10.15.0.0
#define DOCKER_NETWORK_MASK   0xFFFF0000           // /16
#define WG_NETWORK_PREFIX     0x0A630100           // 10.99.1.0
#define WG_NETWORK_MASK       0xFFFFFF00           // /24
#define NGINX_SERVER_IP       0x0A0F002D           // 10.15.0.45

// Estrutura de dados otimizada
struct ip_stats {
    __u32 pkt_count;
    __u32 last_seen;
    __u16 icmp_count;
    __u16 syn_count;
    __u16 high_prio_tokens;
    __u16 low_prio_tokens;
    __u8 flags;                                    // Flags para identificar comportamentos suspeitos
} __attribute__((packed));

// Adicionar flags para comportamento de tráfego
#define FLAG_POTENTIAL_SCANNER 0x01
#define FLAG_SYN_FLOOD         0x02
#define FLAG_ICMP_FLOOD        0x04
#define FLAG_MARKED_SUSPICIOUS 0x08

struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
    __uint(max_entries, 100000);
    __type(key, __u32);
    __type(value, struct ip_stats);
} ddos_map SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 100000);
    __type(key, __u32);
    __type(value, __u32);
} blacklist_map SEC(".maps");

// Mapa para estatísticas de tráfego para integração com userspace
struct traffic_stats {
    __u64 total_packets;
    __u64 dropped_packets;
    __u64 passed_packets;
    __u64 blacklisted_count;
    __u64 syn_flood_count;
    __u64 icmp_flood_count;
};

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

// Fast-path para serviços críticos - pré-verifica IPs confiáveis
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 1000);
    __type(key, __u32);
    __type(value, __u8);     // 1 = confiável
} trusted_hosts_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");

// Lista de permissões fixa
static const __u32 whitelist[] = {
    0x0A000001,  // 10.0.0.1
    0x0A630101,  // 10.99.1.1
    0x0a630102,  // 10.99.1.2
    0x0a630104   // 10.99.1.4
};

// Função otimizada para timestamp
static __always_inline __u32 get_time_sec(void) {
    return bpf_ktime_get_ns() / 1000000000;
}

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: verificar se é rede interna para bypass
static __always_inline int is_internal_network(__u32 ip) {
    // Verificar se é rede Docker
    if ((ip & DOCKER_NETWORK_MASK) == DOCKER_NETWORK_PREFIX)
        return 1;
        
    // Verificar se é rede WireGuard
    if ((ip & WG_NETWORK_MASK) == WG_NETWORK_PREFIX)
        return 1;
        
    return 0;
}

// Função otimizada para classificação de portas
static __always_inline int get_port_priority(__u16 port) {
    // Otimização: usar switch para fast lookup
    switch (port) {
        case 80:    // HTTP
        case 443:   // HTTPS
        case 32100: // WireGuard
        case 32400: // Outro serviço crítico
            return 2;  // Alta prioridade
        case 8117:  // SSH personalizado
        case 8085:  // Serviço interno 1
        case 8086:  // Serviço interno 2
            return 1;  // Média prioridade
        default:
            return 0;  // Baixa prioridade
    }
}

// Função otimizada para detecção de serviços expostos
static __always_inline int is_filtered_port(__u16 port) {
    // Portas críticas verificadas primeiro
    switch (port) {
        case 80:     // HTTP
        case 443:    // HTTPS
        case 8117:   // SSH personalizado
        case 8085:   // Serviço interno 1
        case 8086:   // Serviço interno 2
        case 32100:  // WireGuard
        case 32400:  // Outro serviço crítico
            return 1;
    }
    
    // Demais portas privilegiadas
    return (port <= 1024);
}

// Função otimizada para verificar se é porta WireGuard
static __always_inline int is_wireguard_traffic(__u8 protocol, __u16 port) {
    return (protocol == IPPROTO_UDP && port == 32100);
}

// Função para atualizar estatísticas de tráfego
static __always_inline void update_stats(__u32 action, __u8 flags) {
    __u32 key = 0;
    struct traffic_stats *stats = bpf_map_lookup_elem(&stats_map, &key);
    if (stats) {
        stats->total_packets++;
        
        if (action == XDP_DROP) {
            stats->dropped_packets++;
            if (flags & FLAG_SYN_FLOOD)
                stats->syn_flood_count++;
            if (flags & FLAG_ICMP_FLOOD)
                stats->icmp_flood_count++;
        } else if (action == XDP_PASS) {
            stats->passed_packets++;
        }
        
        if (flags & FLAG_MARKED_SUSPICIOUS)
            stats->blacklisted_count++;
    }
}

// Versão otimizada do programa principal
SEC("xdp")
int xdp_ddos_filter(struct xdp_md *ctx) {
    void *data_end = (void *)(long)ctx->data_end;
    void *data = (void *)(long)ctx->data;
    __u8 flags = 0;
    
    // Otimização: verificação de tamanho mínimo do pacote mais eficiente
    if (__builtin_expect(data + sizeof(struct ethhdr) + sizeof(struct iphdr) > data_end, 0))
        return XDP_PASS;

    struct ethhdr *eth = data;
    if (__builtin_expect(eth->h_proto != __constant_htons(ETH_P_IP), 0))
        return XDP_PASS;

    struct iphdr *ip = (void *)(eth + 1);
    if (__builtin_expect((void *)(ip + 1) > data_end || ip->ihl != 5 || ip->version != 4, 0))
        return XDP_PASS;

    __u32 src_ip = ip->saddr;
    __u32 dst_ip = ip->daddr;
    
    // Fast path 1: Permitir tráfego de redes internas
    if (__builtin_expect(is_internal_network(src_ip), 0)) {
        return XDP_PASS;
    }
    
    // Fast path 2: Verificar whitelist estática
    if (__builtin_expect(is_whitelisted(src_ip), 0)) {
        return XDP_PASS;
    }
    
    // Fast path 3: WireGuard - verificação prioritária para UDP porta 32100
    if (ip->protocol == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if ((void *)(udp + 1) <= data_end) {
            __u16 dest_port = __constant_ntohs(udp->dest);
            if (dest_port == 32100) {  // WireGuard port
                // Permitir todo tráfego WireGuard sem verificação adicional
                update_stats(XDP_PASS, 0);
                return XDP_PASS;
            }
        }
    }
    
    // Fast path 4: Verificar IPs já confiáveis (hosts que já passaram por verificação)
    __u8 *trusted = bpf_map_lookup_elem(&trusted_hosts_map, &src_ip);
    if (__builtin_expect(trusted != NULL && *trusted == 1, 0)) {
        return XDP_PASS;
    }
    
    // Fast path 5: Rejeitar IPs na blacklist
    if (__builtin_expect(is_blacklisted(src_ip), 0)) {
        update_stats(XDP_DROP, FLAG_MARKED_SUSPICIOUS);
        return XDP_DROP;
    }

    // Ignorar pacotes que não são destinados ao nosso servidor NGINX ou serviços expostos
    if (dst_ip != NGINX_SERVER_IP) {
        // Verificar protocolo e porta para classificação de tráfego
        __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;
            
            // Detecção de port scan - múltiplas flags ativas (inválido)
            if ((tcp->syn && tcp->rst) || (tcp->fin && tcp->rst) || (tcp->fin && tcp->syn)) {
                flags |= FLAG_POTENTIAL_SCANNER;
            }
            
        } 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;
            // Verificar tipo de ICMP
            struct icmphdr *icmp = (void *)(ip + 1);
            if (__builtin_expect((void *)(icmp + 1) > data_end, 0))
                return XDP_PASS;
                
            // Permitir apenas echo-reply (0) e limitar echo-request (8)
            if (icmp->type != 0 && icmp->type != 8) {
                update_stats(XDP_DROP, 0);
                return XDP_DROP;
            }
        } else {
            // Outros protocolos passam sem filtro adicional
            return XDP_PASS;
        }

        // Otimização: Se não for para portas filtradas, deixar passar
        if (__builtin_expect(!is_filtered_port(dest_port) && !is_icmp, 1)) {
            // Se for um cliente fazendo várias conexões legítimas, marcá-lo como confiável
            struct ip_stats *entry = bpf_map_lookup_elem(&ddos_map, &src_ip);
            if (entry && entry->pkt_count > 5 && entry->flags == 0) {
                __u8 mark_trusted = 1;
                bpf_map_update_elem(&trusted_hosts_map, &src_ip, &mark_trusted, BPF_ANY);
            }
            return XDP_PASS;
        }

        // Prioridade do pacote
        int priority = get_port_priority(dest_port);
        
        // Rate limiting baseado em tokens
        struct ip_stats *entry = bpf_map_lookup_elem(&ddos_map, &src_ip);
        __u32 now = get_time_sec();
        
        if (entry) {
            // Cálculo de tokens disponíveis
            __u32 elapsed = 0;
            if (__builtin_expect(now > entry->last_seen, 1)) {
                elapsed = now - entry->last_seen;
            }

            // Implementação de token bucket
            if (elapsed > 0) {
                // Adicionar tokens com limite 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 se passou tempo suficiente
                if (elapsed >= TIMEOUT_SEC) {
                    entry->icmp_count = 0;
                    entry->syn_count = 0;
                    entry->flags &= ~(FLAG_SYN_FLOOD | FLAG_ICMP_FLOOD);
                }
            }

            // Atualizar flags e contadores
            entry->last_seen = now;
            entry->pkt_count++;
            
            if (flags & FLAG_POTENTIAL_SCANNER) {
                entry->flags |= FLAG_POTENTIAL_SCANNER;
            }
            
            if (is_icmp) {
                entry->icmp_count++;
                if (__builtin_expect(entry->icmp_count > MAX_ICMP_PACKETS_PER_IP, 0)) {
                    entry->flags |= FLAG_ICMP_FLOOD;
                    __u32 one = 1;
                    bpf_map_update_elem(&blacklist_map, &src_ip, &one, BPF_ANY);
                    update_stats(XDP_DROP, FLAG_ICMP_FLOOD);
                    return XDP_DROP;
                }
            }
            
            if (is_syn) {
                entry->syn_count++;
                if (__builtin_expect(entry->syn_count > MAX_SYN_PACKETS_PER_IP, 0)) {
                    entry->flags |= FLAG_SYN_FLOOD;
                    __u32 one = 1;
                    bpf_map_update_elem(&blacklist_map, &src_ip, &one, BPF_ANY);
                    update_stats(XDP_DROP, FLAG_SYN_FLOOD);
                    return XDP_DROP;
                }
            }

            // Blacklist se exceder total de pacotes
            if (__builtin_expect(entry->pkt_count > MAX_PACKETS_PER_IP, 0)) {
                __u32 one = 1;
                entry->flags |= FLAG_MARKED_SUSPICIOUS;
                bpf_map_update_elem(&blacklist_map, &src_ip, &one, BPF_ANY);
                update_stats(XDP_DROP, FLAG_MARKED_SUSPICIOUS);
                return XDP_DROP;
            }

            // Aplicar token bucket conforme prioridade
            if (priority == 2) {  // Alta prioridade
                if (entry->high_prio_tokens >= PACKET_COST) {
                    entry->high_prio_tokens -= PACKET_COST;
                    update_stats(XDP_PASS, 0);
                    return XDP_PASS;
                }
            } else if (priority == 1) {  // Média prioridade
                // Média prioridade usa metade de cada bucket
                if (entry->high_prio_tokens >= PACKET_COST/2 && 
                    entry->low_prio_tokens >= PACKET_COST/2) {
                    entry->high_prio_tokens -= PACKET_COST/2;
                    entry->low_prio_tokens -= PACKET_COST/2;
                    update_stats(XDP_PASS, 0);
                    return XDP_PASS;
                }
            } else {  // Baixa prioridade
                if (entry->low_prio_tokens >= PACKET_COST) {
                    entry->low_prio_tokens -= PACKET_COST;
                    update_stats(XDP_PASS, 0);
                    return XDP_PASS;
                }
            }
            
            update_stats(XDP_DROP, 0);
            return XDP_DROP;
            
        } else {
            // Primeira vez que vemos este IP
            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,
                .flags = flags
            };
            
            bpf_map_update_elem(&ddos_map, &src_ip, &new_entry, BPF_ANY);
            update_stats(XDP_PASS, 0);
            return XDP_PASS;
        }
    }
    
    // Processar pacotes destinados ao servidor NGINX
    if (dst_ip == NGINX_SERVER_IP) {
        struct ip_stats *entry = bpf_map_lookup_elem(&ddos_map, &src_ip);
        __u32 now = get_time_sec();
        
        // Verificar cabeçalhos TCP para tráfego HTTP(S)
        if (ip->protocol == IPPROTO_TCP) {
            struct tcphdr *tcp = (void *)(ip + 1);
            if (__builtin_expect((void *)(tcp + 1) > data_end, 0))
                return XDP_PASS;
                
            __u16 dest_port = __constant_ntohs(tcp->dest);
            if (dest_port == 80 || dest_port == 443) {
                // Lógica específica para proteção de HTTP/HTTPS
                
                // Se for SYN, fazer rate limiting específico
                if (tcp->syn && !tcp->ack) {
                    if (entry) {
                        entry->syn_count++;
                        if (entry->syn_count > MAX_SYN_PACKETS_PER_IP) {
                            update_stats(XDP_DROP, FLAG_SYN_FLOOD);
                            return XDP_DROP;
                        }
                    } else {
                        // Primeiro SYN deste IP
                        struct ip_stats new_entry = {
                            .pkt_count = 1,
                            .last_seen = now,
                            .syn_count = 1,
                            .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);
                    }
                }
                
                // Permitir o pacote (processamento de entrada HTTP é prioritário)
                update_stats(XDP_PASS, 0);
                return XDP_PASS;
            }
        }
    }
    
    // Se chegou até aqui, passa o tráfego
    update_stats(XDP_PASS, 0);
    return XDP_PASS;
}

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