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

#define MAX_HIGH_PRIORITY_PACKETS_PER_SECOND 2000 // Limite de pacotes por segundo para alta prioridade
#define MAX_LOW_PRIORITY_PACKETS_PER_SECOND 500   // Limite de pacotes por segundo para baixa prioridade
#define MAX_PACKETS_PER_IP 60000                    // Limite de pacotes por IP antes de bloquear
#define MAX_ICMP_PACKETS_PER_IP 50                // Limite de pacotes ICMP por IP antes de bloquear
#define MAX_SYN_PACKETS_PER_IP 100                // Limite de pacotes SYN por IP antes de bloquear
#define TIMEOUT_SEC 1                             // Tempo para expiração do contador (em segundos)

// Estrutura para armazenar contagem de pacotes e o último timestamp visto para um IP
struct ip_stats {
    __u64 pkt_count;   // Contador de pacotes
    __u64 icmp_count;  // Contador de pacotes ICMP
    __u64 syn_count;   // Contador de pacotes SYN
    __u64 last_seen;   // Timestamp do último pacote
    __u64 high_priority_count; // Contador de pacotes de alta prioridade
    __u64 low_priority_count;  // Contador de pacotes de baixa prioridade
};

// Estruturas para NAT
struct nat_entry {
    __u32 new_ip;      // IP de destino após NAT
    __u16 new_port;    // Porta de destino após NAT
    __u16 proto;       // Protocolo (TCP/UDP)
} __attribute__((aligned(8)));

// Modificar mapa DNAT para usar porta como chave
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 10000);
    __type(key, __u16);           // Porta de destino
    __type(value, struct nat_entry); // Nova configuração
} dnat_map SEC(".maps");

struct masquerade_key {
    __u32 src_ip;      // IP de origem
    __u32 dst_ip;      // IP de destino
    __u16 dst_port;    // Porta de destino
    __u16 proto;       // Protocolo
} __attribute__((aligned(8)));

// Mapa de hash para armazenar estatísticas de IPs
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 100000);  // Número máximo de entradas no mapa
    __type(key, __u32);           // Tipo da chave (endereço IP)
    __type(value, struct ip_stats); // Tipo do valor (estatísticas do IP)
} ddos_map SEC(".maps");

// Mapa de hash para armazenar a blacklist
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 100000);  // Número máximo de entradas no mapa
    __type(key, __u32);           // Tipo da chave (endereço IP)
    __type(value, __u32);         // Tipo do valor (dummy value)
} blacklist_map SEC(".maps");

// Mapas adicionais para NAT
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 10000);
    __type(key, struct masquerade_key);
    __type(value, __u32);         // IP mascarado
} masquerade_map SEC(".maps");

// Adicionar estrutura para SNAT
struct snat_entry {
    __u32 new_ip;      // IP após NAT
    __u16 new_port;    // Porta após NAT
} __attribute__((aligned(8)));

// Adicionar mapa para SNAT
struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 10000);
    __type(key, struct masquerade_key);
    __type(value, struct snat_entry);
} snat_map SEC(".maps");

// Lista de permissões estática (whitelist) de IPs confiáveis
const __u32 whitelist[] = {
    0x0A000001, // 10.0.0.1
    0x0A630101,  // 192.168.0.1
    0x0a630102,
    0x0a630104
};

// Adicionar 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");

// Função auxiliar para obter o timestamp atual em nanossegundos
static __inline __u64 get_time_ns(void) {
    return bpf_ktime_get_ns();
}

// Função para verificar se o IP está na lista de permissões
static __inline int is_whitelisted(__u32 ip) {
    int i;
    for (i = 0; i < sizeof(whitelist) / sizeof(whitelist[0]); i++) {
        if (ip == whitelist[i]) {
            return 1;
        }
    }
    return 0;
}

// Função para verificar se o IP está na lista negra
static __inline int is_blacklisted(__u32 ip) {
    __u32 *value = bpf_map_lookup_elem(&blacklist_map, &ip);
    return value != NULL;
}

// Função para verificar se a porta é de alta prioridade
static __inline int is_high_priority_port(__u16 port) {
    return (port == 80 || port == 443 || port == 32100 || port == 32400);
}

// Função para verificar se a porta é de baixa prioridade
static __inline int is_low_priority_port(__u16 port) {
    return !is_high_priority_port(port);
}

// Função para verificar se a porta deve ser filtrada
static __inline int is_filtered_port(__u16 port) {
    // Intervalo de portas 1-1024
    if (port >= 1 && port <= 1024) {
        return 1;
    }
    // Portas específicas
    if (port == 8117 || port == 8085 || port == 8086 || port == 32100 || port == 32400) {
        return 1;
    }
    return 0;
}

// Função para verificar se o pacote é destinado ao servidor e deve ser filtrado
static __inline int is_destined_to_server(struct iphdr *ip, void *data_end) {
    if (ip->protocol == IPPROTO_TCP) {
        struct tcphdr *tcp = (void *)(ip + 1);
        if ((void *)(tcp + 1) > data_end)
            return 0;
        // Verifique se a porta está na lista de portas a serem filtradas
        __u16 dest_port = __constant_ntohs(tcp->dest);
        if (is_filtered_port(dest_port)) {
            return 1;
        }
    } else if (ip->protocol == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if ((void *)(udp + 1) > data_end)
            return 0;
        // Verifique se a porta está na lista de portas a serem filtradas
        __u16 dest_port = __constant_ntohs(udp->dest);
        if (is_filtered_port(dest_port)) {
            return 1;
        }
    } else if (ip->protocol == IPPROTO_ICMP) {
        return 1;
    }
    return 0;
}

// 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;

    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end)
        return XDP_PASS;

    // Verifica se é um pacote IP
    if (eth->h_proto != __constant_htons(ETH_P_IP))
        return XDP_PASS;

    struct iphdr *ip = (void *)(eth + 1);
    if ((void *)(ip + 1) > data_end)
        return XDP_PASS;

    // Verifica se o cabeçalho IP é válido
    if (ip->ihl != 5 || ip->version != 4)
        return XDP_PASS;

    // Verifica se o IP está na lista de permissões
    __u32 src_ip = ip->saddr;
    if (is_whitelisted(src_ip)) {
        return XDP_PASS;
    }

    // Verifica se o IP está na lista negra
    if (is_blacklisted(src_ip)) {
        return XDP_DROP;
    }

    // Verifica se o pacote é destinado ao servidor e deve ser filtrado
    if (!is_destined_to_server(ip, data_end))
        return XDP_PASS;

    __u16 dest_port = 0;
    struct tcphdr *tcp = NULL;

    // Verifica se é um pacote TCP ou UDP
    if (ip->protocol == IPPROTO_TCP) {
        tcp = (void *)(ip + 1);
        if ((void *)(tcp + 1) > data_end)
            return XDP_PASS;
        dest_port = __constant_ntohs(tcp->dest);

        // Verifica se é um pacote SYN
        if (tcp->syn && !tcp->ack) {
            // Consulta o mapa para verificar se o IP já está registrado
            struct ip_stats *entry = bpf_map_lookup_elem(&ddos_map, &src_ip);

            if (entry) {
                __u64 now = get_time_ns();

                // Reinicia o contador se o tempo decorrido for maior que um segundo
                if ((now - entry->last_seen) > TIMEOUT_SEC * 1000000000ULL) {
                    entry->syn_count = 0;
                    entry->last_seen = now;
                }

                // Incrementa o contador de pacotes SYN
                entry->syn_count++;
                entry->last_seen = now;

                // Bloqueia o tráfego se o limite de pacotes SYN for excedido
                if (entry->syn_count > MAX_SYN_PACKETS_PER_IP) {
                    return XDP_DROP;
                }
            } else {
                // Adiciona um novo IP ao mapa
                struct ip_stats new_entry = {
                    .syn_count = 1,
                    .last_seen = get_time_ns()
                };
                bpf_map_update_elem(&ddos_map, &src_ip, &new_entry, BPF_ANY);
            }
        }
    } else if (ip->protocol == IPPROTO_UDP) {
        struct udphdr *udp = (void *)(ip + 1);
        if ((void *)(udp + 1) > data_end)
            return XDP_PASS;
        dest_port = __constant_ntohs(udp->dest);
    }

    // Consulta o mapa para verificar se o IP já está registrado
    struct ip_stats *entry = bpf_map_lookup_elem(&ddos_map, &src_ip);

    if (entry) {
        __u64 now = get_time_ns();

        // Reinicia o contador se o tempo decorrido for maior que um segundo
        if ((now - entry->last_seen) > TIMEOUT_SEC * 1000000000ULL) {
            entry->pkt_count = 0;
            entry->icmp_count = 0;
            entry->syn_count = 0;
            entry->high_priority_count = 0;
            entry->low_priority_count = 0;
            entry->last_seen = now;
        }

        // Rate limiting para pacotes de alta prioridade
        if (is_high_priority_port(dest_port)) {
            entry->high_priority_count++;
            if (entry->high_priority_count > MAX_HIGH_PRIORITY_PACKETS_PER_SECOND) {
                return XDP_DROP;
            }
        }

        // Rate limiting para pacotes de baixa prioridade
        if (is_low_priority_port(dest_port)) {
            entry->low_priority_count++;
            if (entry->low_priority_count > MAX_LOW_PRIORITY_PACKETS_PER_SECOND) {
                return XDP_DROP;
            }
        }

        // Incrementa o contador de pacotes
        entry->pkt_count++;
        entry->last_seen = now;

        // Filtra pacotes ICMP
        if (ip->protocol == IPPROTO_ICMP) {
            entry->icmp_count++;
            if (entry->icmp_count > MAX_ICMP_PACKETS_PER_IP) {
                return XDP_DROP;
            }
        }

        // Bloqueia o tráfego se o limite de pacotes for excedido
        if (entry->pkt_count > MAX_PACKETS_PER_IP) {
            return XDP_DROP;
        }
    } else {
        // Adiciona um novo IP ao mapa
        struct ip_stats new_entry = {
            .pkt_count = 1,
            .icmp_count = (ip->protocol == IPPROTO_ICMP) ? 1 : 0,
            .syn_count = (ip->protocol == IPPROTO_TCP && tcp && tcp->syn && !tcp->ack) ? 1 : 0,
            .last_seen = get_time_ns(),
            .high_priority_count = is_high_priority_port(dest_port) ? 1 : 0,
            .low_priority_count = is_low_priority_port(dest_port) ? 1 : 0
        };
        bpf_map_update_elem(&ddos_map, &src_ip, &new_entry, BPF_ANY);
    }

    return XDP_PASS;
}

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