// firewall.bpf.c
#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <linux/icmp.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_endian.h>

// === 配置常量 ===
#define ETH_ALEN 6
#define MAX_PORT_RULES    256
#define MAX_BLACKLIST     65536
#define MAX_WHITELIST     4096

// === 动作定义 ===
enum firewall_action {
    FW_PASS = 0,
    FW_DROP,
    FW_LOG,
};

// === 端口规则（5 元组）===
struct port_rule {
    u32 src_ip;        // 0 = 任意
    u32 dst_ip;        // 0 = 任意
    u16 src_port;      // 0 = 任意
    u16 dst_port;      // 0 = 任意
    u8  protocol;      // 0 = 任意
    u8  action;        // 0=PASS, 1=DROP
    u8  pad[2];
} __attribute__((packed));

// === 全局黑名单（hash，O(1) 查找）===
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, MAX_BLACKLIST);
    __type(key, u32);   // 黑名单 IP
    __type(value, u8);  // 1 = 启用
} blacklist SEC(".maps");

// === 全局白名单 ===
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, MAX_WHITELIST);
    __type(key, u32);
    __type(value, u8);
} whitelist SEC(".maps");

// === 端口规则（5 元组精确匹配）===
struct {
    __uint(type, BPF_MAP_TYPE_LRU_HASH);
    __uint(max_entries, 1024);
    __type(key, struct port_rule_key);
    __type(value, struct port_rule);
} port_rules SEC(".maps");

// 端口规则 key：源 IP + 目的 IP + 源端口 + 目的端口 + 协议
struct port_rule_key {
    u32 src_ip;
    u32 dst_ip;
    u16 src_port;
    u16 dst_port;
    u8  protocol;
} __attribute__((packed));

// === IP 段过滤（LPM trie）===
struct {
    __uint(type, BPF_MAP_TYPE_LPM_TRIE);
    __uint(max_entries, 1024);
    __type(key, struct lpm_key);
    __type(value, u8);  // 1 = DROP
} ip_range_deny SEC(".maps");

struct lpm_key {
    u32 prefixlen;   // 必须在第一位
    u32 addr;
} __attribute__((packed));

// === 统计（per-CPU，避免锁）===
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
    __uint(max_entries, 8);
    __type(key, u32);
    __type(value, u64);
} stats SEC(".maps");

enum {
    STAT_TOTAL = 0,
    STAT_PASS = 1,
    STAT_DROP = 2,
    STAT_BLACKLIST = 3,
    STAT_WHITELIST = 4,
    STAT_PORT = 5,
    STAT_IP_RANGE = 6,
    STAT_INVALID = 7,
};

// === 事件输出（ring buffer）===
struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 256 * 1024);
} events SEC(".maps");

struct firewall_event {
    u64 timestamp;
    u32 src_ip;
    u32 dst_ip;
    u16 src_port;
    u16 dst_port;
    u8  protocol;
    u8  action;
    u8  reason;        // 拦截原因
    u8  pad;
};

// 拦截原因
enum {
    REASON_NONE = 0,
    REASON_BLACKLIST,
    REASON_PORT,
    REASON_IP_RANGE,
};

// === 时间戳辅助 ===
static __always_inline u64 now_ns(void) {
    return bpf_ktime_get_ns();
}

// === 发送事件 ===
static __always_inline void emit_event(
    struct xdp_md *ctx,
    u32 src_ip, u32 dst_ip,
    u16 src_port, u16 dst_port,
    u8 protocol, u8 action, u8 reason) {

    struct firewall_event *e;
    e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
    if (!e) return;

    e->timestamp = bpf_ktime_get_ns();
    e->src_ip = src_ip;
    e->dst_ip = dst_ip;
    e->src_port = src_port;
    e->dst_port = dst_port;
    e->protocol = protocol;
    e->action = action;
    e->reason = reason;

    bpf_ringbuf_submit(e, 0);
}

// === 更新统计 ===
static __always_inline void update_stat(u32 key) {
    u64 *count = bpf_map_lookup_elem(&stats, &key);
    if (count) (*count)++;
}

// === 检查 IP 段 ===
static __always_inline int check_ip_range(u32 ip, int is_src) {
    struct lpm_key key = {
        .prefixlen = 32,
        .addr = ip,
    };

    // LPM trie 自动处理前缀匹配
    u8 *action = bpf_map_lookup_elem(&ip_range_deny, &key);
    if (action && *action == 1) {
        return 1;  // 命中黑名单段
    }
    return 0;
}

// === 主程序 ===
SEC("xdp")
int xdp_firewall(struct xdp_md *ctx) {
    void *data = (void *)(long)ctx->data;
    void *data_end = (void *)(long)ctx->data_end;

    update_stat(STAT_TOTAL);

    // === 1. 以太网头 ===
    struct ethhdr *eth = data;
    if ((void *)(eth + 1) > data_end) {
        update_stat(STAT_INVALID);
        return XDP_PASS;
    }

    __u16 eth_proto = eth->h_proto;
    u32 src_ip = 0, dst_ip = 0;
    u8  protocol = 0;
    u16 src_port = 0, dst_port = 0;
    int is_ipv6 = 0;

    // === 2. 解析网络层 ===
    if (eth_proto == bpf_htons(ETH_P_IP)) {
        struct iphdr *ip = (void *)(eth + 1);
        if ((void *)(ip + 1) > data_end) {
            update_stat(STAT_INVALID);
            return XDP_PASS;
        }

        src_ip = ip->saddr;
        dst_ip = ip->daddr;
        protocol = ip->protocol;

        // === 3. 解析传输层 ===
        if (protocol == IPPROTO_TCP) {
            struct tcphdr *tcp = (void *)ip + (ip->ihl * 4);
            if ((void *)(tcp + 1) > data_end) {
                update_stat(STAT_INVALID);
                return XDP_PASS;
            }
            src_port = tcp->source;
            dst_port = tcp->dest;
        } else if (protocol == IPPROTO_UDP) {
            struct udphdr *udp = (void *)ip + (ip->ihl * 4);
            if ((void *)(udp + 1) > data_end) {
                update_stat(STAT_INVALID);
                return XDP_PASS;
            }
            src_port = udp->source;
            dst_port = udp->dest;
        }
    } else if (eth_proto == bpf_htons(ETH_P_IPV6)) {
        is_ipv6 = 1;
        struct ipv6hdr *ip6 = (void *)(eth + 1);
        if ((void *)(ip6 + 1) > data_end) {
            update_stat(STAT_INVALID);
            return XDP_PASS;
        }
        // 提取 IPv6 末尾 32 位作为简化 hash
        src_ip = ip6->saddr.in6_u.u6_addr32[3];
        dst_ip = ip6->daddr.in6_u.u6_addr32[3];
        protocol = ip6->nexthdr;
        // (类似解析 transport 层)
    } else {
        // 非 IP 协议，直接放行
        update_stat(STAT_PASS);
        return XDP_PASS;
    }

    // === 4. 规则检查 ===

    // 4.1 白名单（最高优先级）
    u8 *whitelisted = bpf_map_lookup_elem(&whitelist, &src_ip);
    if (whitelisted) {
        update_stat(STAT_WHITELIST);
        update_stat(STAT_PASS);
        return XDP_PASS;
    }

    // 4.2 黑名单
    u8 *blacklisted = bpf_map_lookup_elem(&blacklist, &src_ip);
    if (blacklisted) {
        update_stat(STAT_BLACKLIST);
        update_stat(STAT_DROP);
        emit_event(ctx, src_ip, dst_ip, src_port, dst_port,
                   protocol, FW_DROP, REASON_BLACKLIST);
        return XDP_DROP;
    }

    // 4.3 IP 段黑名单
    if (check_ip_range(src_ip, 1)) {
        update_stat(STAT_IP_RANGE);
        update_stat(STAT_DROP);
        emit_event(ctx, src_ip, dst_ip, src_port, dst_port,
                   protocol, FW_DROP, REASON_IP_RANGE);
        return XDP_DROP;
    }

    // 4.4 端口规则
    struct port_rule_key pk = {
        .src_ip = src_ip,
        .dst_ip = dst_ip,
        .src_port = src_port,
        .dst_port = dst_port,
        .protocol = protocol,
    };

    struct port_rule *rule = bpf_map_lookup_elem(&port_rules, &pk);
    if (rule) {
        if (rule->action == FW_DROP) {
            update_stat(STAT_PORT);
            update_stat(STAT_DROP);
            emit_event(ctx, src_ip, dst_ip, src_port, dst_port,
                       protocol, FW_DROP, REASON_PORT);
            return XDP_DROP;
        }
    }

    // 4.5 默认通过
    update_stat(STAT_PASS);
    return XDP_PASS;
}

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

// fwctl.c — 用户态管理工具
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <getopt.h>
#include <signal.h>
#include <bpf/bpf.h>
#include <bpf/libbpf.h>
#include <arpa/inet.h>
#include "firewall.skel.h"

static volatile bool exiting = false;
static void sig_handler(int sig) { (void)sig; exiting = true; }

static int libbpf_print_fn(enum libbpf_print_level level,
                            const char *format, va_list args) {
    if (level == LIBBPF_DEBUG) return 0;
    return vfprintf(stderr, format, args);
}

void print_usage(const char *prog) {
    fprintf(stderr, "Usage:\n");
    fprintf(stderr, "  %s attach <ifname>          # Attach XDP to interface\n", prog);
    fprintf(stderr, "  %s detach <ifname>          # Detach XDP\n", prog);
    fprintf(stderr, "  %s blacklist add <ip>       # Add to blacklist\n", prog);
    fprintf(stderr, "  %s blacklist del <ip>       # Remove from blacklist\n", prog);
    fprintf(stderr, "  %s whitelist add <ip>       # Add to whitelist\n", prog);
    fprintf(stderr, "  %s whitelist del <ip>       # Remove from whitelist\n", prog);
    fprintf(stderr, "  %s port <proto> <port> <action>  # Add port rule\n", prog);
    fprintf(stderr, "  %s range add <ip/cidr>      # Add IP range to deny\n", prog);
    fprintf(stderr, "  %s range del <ip/cidr>      # Remove IP range\n", prog);
    fprintf(stderr, "  %s stats                    # Show statistics\n", prog);
    fprintf(stderr, "  %s list                     # List all rules\n", prog);
    fprintf(stderr, "  %s log                      # Show recent events\n", prog);
}

// === 工具函数 ===
int parse_ip(const char *s, uint32_t *ip) {
    struct in_addr addr;
    if (inet_pton(AF_INET, s, &addr) != 1) return -1;
    *ip = addr.s_addr;
    return 0;
}

int parse_cidr(const char *s, uint32_t *ip, uint32_t *prefixlen) {
    char buf[64];
    strncpy(buf, s, sizeof(buf) - 1);
    char *slash = strchr(buf, '/');
    if (slash) {
        *slash = '\0';
        *prefixlen = atoi(slash + 1);
    } else {
        *prefixlen = 32;
    }
    return parse_ip(buf, ip);
}

// === 主程序 ===
int main(int argc, char **argv) {
    if (argc < 2) {
        print_usage(argv[0]);
        return 1;
    }

    libbpf_set_print(libbpf_print_fn);
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    // 打开 BPF skeleton
    struct firewall_bpf *skel = firewall_bpf__open();
    if (!skel) { fprintf(stderr, "Failed to open BPF skeleton\n"); return 1; }

    if (firewall_bpf__load(skel)) {
        fprintf(stderr, "Failed to load BPF\n");
        firewall_bpf__destroy(skel);
        return 1;
    }

    const char *cmd = argv[1];

    // === attach ===
    if (strcmp(cmd, "attach") == 0) {
        if (argc < 3) { print_usage(argv[0]); return 1; }
        int ifindex = if_nametoindex(argv[2]);

        // 尝试 native 模式，失败回退到 generic
        if (bpf_xdp_attach(ifindex, bpf_program__fd(skel->progs.xdp_firewall),
                           XDP_FLAGS_DRV_MODE, NULL) < 0) {
            fprintf(stderr, "Native XDP failed, trying generic\n");
            if (bpf_xdp_attach(ifindex, bpf_program__fd(skel->progs.xdp_firewall),
                               XDP_FLAGS_SKB_MODE, NULL) < 0) {
                fprintf(stderr, "Failed to attach XDP\n");
                return 1;
            }
        }
        printf("XDP attached to %s\n", argv[2]);
    }

    // === detach ===
    else if (strcmp(cmd, "detach") == 0) {
        if (argc < 3) { print_usage(argv[0]); return 1; }
        int ifindex = if_nametoindex(argv[2]);
        bpf_xdp_detach(ifindex, XDP_FLAGS_DRV_MODE, NULL);
        bpf_xdp_detach(ifindex, XDP_FLAGS_SKB_MODE, NULL);
        printf("XDP detached from %s\n", argv[2]);
    }

    // === blacklist add/del ===
    else if (strcmp(cmd, "blacklist") == 0) {
        if (argc < 4) { print_usage(argv[0]); return 1; }
        const char *op = argv[2];
        uint32_t ip;
        if (parse_ip(argv[3], &ip) < 0) {
            fprintf(stderr, "Invalid IP\n");
            return 1;
        }

        struct bpf_map *map = skel->maps.blacklist;
        if (strcmp(op, "add") == 0) {
            uint8_t v = 1;
            bpf_map__update_elem(map, &ip, sizeof(ip), &v, sizeof(v), BPF_ANY);
            printf("Added %s to blacklist\n", argv[3]);
        } else if (strcmp(op, "del") == 0) {
            bpf_map__delete_elem(map, &ip, sizeof(ip), 0);
            printf("Removed %s from blacklist\n", argv[3]);
        }
    }

    // === whitelist add/del ===
    else if (strcmp(cmd, "whitelist") == 0) {
        if (argc < 4) { print_usage(argv[0]); return 1; }
        const char *op = argv[2];
        uint32_t ip;
        if (parse_ip(argv[3], &ip) < 0) {
            fprintf(stderr, "Invalid IP\n");
            return 1;
        }

        struct bpf_map *map = skel->maps.whitelist;
        if (strcmp(op, "add") == 0) {
            uint8_t v = 1;
            bpf_map__update_elem(map, &ip, sizeof(ip), &v, sizeof(v), BPF_ANY);
            printf("Added %s to whitelist\n", argv[3]);
        } else if (strcmp(op, "del") == 0) {
            bpf_map__delete_elem(map, &ip, sizeof(ip), 0);
            printf("Removed %s from whitelist\n", argv[3]);
        }
    }

    // === port ===
    else if (strcmp(cmd, "port") == 0) {
        if (argc < 5) { print_usage(argv[0]); return 1; }
        const char *proto_s = argv[2];
        int port = atoi(argv[3]);
        const char *action_s = argv[4];

        uint8_t protocol = strcmp(proto_s, "tcp") == 0 ? IPPROTO_TCP : IPPROTO_UDP;
        uint8_t action = strcmp(action_s, "drop") == 0 ? 1 : 0;

        struct port_rule_key key = {0};
        key.dst_port = bpf_htons(port);
        key.protocol = protocol;

        struct port_rule rule = {0};
        rule.dst_port = bpf_htons(port);
        rule.protocol = protocol;
        rule.action = action;

        struct bpf_map *map = skel->maps.port_rules;
        bpf_map__update_elem(map, &key, sizeof(key), &rule, sizeof(rule), BPF_ANY);
        printf("Added port rule: %s %d -> %s\n", proto_s, port, action_s);
    }

    // === range ===
    else if (strcmp(cmd, "range") == 0) {
        if (argc < 4) { print_usage(argv[0]); return 1; }
        const char *op = argv[2];
        uint32_t ip, prefixlen;
        if (parse_cidr(argv[3], &ip, &prefixlen) < 0) {
            fprintf(stderr, "Invalid CIDR\n");
            return 1;
        }

        struct lpm_key key = {
            .prefixlen = prefixlen,
            .addr = ip,
        };

        struct bpf_map *map = skel->maps.ip_range_deny;
        if (strcmp(op, "add") == 0) {
            uint8_t v = 1;
            bpf_map__update_elem(map, &key, sizeof(key), &v, sizeof(v), BPF_ANY);
            printf("Added IP range %s to deny\n", argv[3]);
        } else if (strcmp(op, "del") == 0) {
            bpf_map__delete_elem(map, &key, sizeof(key), 0);
            printf("Removed IP range %s\n", argv[3]);
        }
    }

    // === stats ===
    else if (strcmp(cmd, "stats") == 0) {
        int ncpu = libbpf_num_possible_cpus();
        uint64_t values[ncpu];
        const char *names[] = {
            "Total", "Pass", "Drop", "Blacklist",
            "Whitelist", "Port", "IP-Range", "Invalid"
        };

        printf("=== Firewall Statistics ===\n");
        for (int i = 0; i < 8; i++) {
            if (bpf_map__lookup_elem(skel->maps.stats, &i, sizeof(i),
                                      values, sizeof(values), 0) == 0) {
                uint64_t total = 0;
                for (int c = 0; c < ncpu; c++) total += values[c];
                printf("  %-12s: %lu\n", names[i], total);
            }
        }
    }

    // === list ===
    else if (strcmp(cmd, "list") == 0) {
        // 遍历所有 map 并打印内容
        struct bpf_map *maps[] = {
            skel->maps.blacklist,
            skel->maps.whitelist,
            skel->maps.port_rules,
            skel->maps.ip_range_deny,
        };
        const char *map_names[] = {"Blacklist", "Whitelist", "Port Rules", "IP Range"};

        for (int i = 0; i < 4; i++) {
            printf("\n=== %s ===\n", map_names[i]);
            // ... 遍历 map 并打印 ...
        }
    }

    // === log ===
    else if (strcmp(cmd, "log") == 0) {
        printf("Streaming firewall events (Ctrl-C to stop)...\n");

        struct ring_buffer *rb = ring_buffer__new(
            bpf_map__fd(skel->maps.events),
            handle_event, NULL, NULL);

        while (!exiting) {
            ring_buffer__poll(rb, 1000);
        }
    }

    else {
        print_usage(argv[0]);
    }

    firewall_bpf__destroy(skel);
    return 0;
}

// 事件回调
static int handle_event(void *ctx, void *data, size_t data_sz) {
    const struct firewall_event *e = data;
    char src[INET_ADDRSTRLEN], dst[INET_ADDRSTRLEN];
    struct in_addr s = { .s_addr = e->src_ip };
    struct in_addr d = { .s_addr = e->dst_ip };

    inet_ntop(AF_INET, &s, src, sizeof(src));
    inet_ntop(AF_INET, &d, dst, sizeof(dst));

    const char *reason = e->reason == 1 ? "blacklist" :
                         e->reason == 2 ? "port" :
                         e->reason == 3 ? "ip-range" : "unknown";

    printf("[%llu] %s:%d -> %s:%d proto=%d ACTION=%s reason=%s\n",
           e->timestamp,
           src, ntohs(e->src_port),
           dst, ntohs(e->dst_port),
           e->protocol,
           e->action == 1 ? "DROP" : "PASS",
           reason);
    return 0;
}