#include "cap.h"
#include "common.h"
#include "libipset.h"
#include "libcurl.h"

pcap_if_t   *alldevs, *device;
pcap_t      *handle;            // 会话句柄
struct bpf_program fp;          // 编译后的过滤器

pid_t pid = -1;                 // 子进程全局PID

#define SHM_SIZE 1024           // 共享内存大小
#define SHM_KEY 0124            // 共享内存键值
int     shmid = -1;
int     RULE_NAME_NUMBER = 0;   // ipset 集合集合数
char    *RULE_NAME = NULL;      // 共享内存
char    *ip2region_area = NULL; // ip2region 解析结果
char    *command_result = NULL; // 执行命令的结果

#define CACHE_TTL 180           // 设定缓存的存活时间为 600 秒 (10 分钟)
#define MAX_CACHE_SIZE 100      // 缓存最多存储 100 个 IP 地址
struct ip_cache_node *ip_cache_head = NULL; //  缓存链表的头节点
int cache_size = 0;             // 当前缓存中的 IP 数量

// 定义链表结构，用于缓存 IP 地址
struct ip_cache_node {
    char ip[INET_ADDRSTRLEN];   // 存储 IP 地址
    time_t timestamp;           // 记录缓存时间
    struct ip_cache_node *next; // 指向下一个节点
};

#define MAXIPSET_ 100
#define MAXIPLEN 32
char cn_ip[MAXIPSET_][MAXIPLEN] = { 0 };

// 添加一个 IP 到集合（如果已存在则不添加）
int add_cn_ip(char cn_ip[MAXIPSET_][MAXIPLEN], char *ip)
{
    if (ip == NULL || strlen(ip) >= MAXIPLEN) {
        return -1;              // 错误：无效的 IP 地址或过长
    }
    // 检查是否已存在
    for (int i = 0; i < MAXIPSET_; i++) {
        if (cn_ip[i][0] != '\0' && strcmp(cn_ip[i], ip) == 0) {
            return 1;           // IP 已存在，返回特殊代码
        }
    }
    // 查找空位并添加
    for (int i = 0; i < MAXIPSET_; i++) {
        if (cn_ip[i][0] == '\0') { // 检查是否为空
            strcpy(cn_ip[i], ip);
            return 0;           // 成功添加
        }
    }
    return -1;                  // 错误：集合已满
}

// 判断给定的IP地址是否在集合中
int is_ip_in_set(char cn_ip[MAXIPSET_][MAXIPLEN], const char *ip)
{
    if (ip == NULL) {
        return 0;               // 错误：无效的IP地址指针
    }
    for (int i = 0; i < MAXIPSET_; i++) {
        if (cn_ip[i][0] != '\0' && strcmp(cn_ip[i], ip) == 0) {
            return 1;           // 找到匹配的IP地址
        }
    }
    return 0;                   // 没有找到匹配的IP地址
}

// 计算集合中非空IP的数量
int cn_ip_len(char cn_ip[MAXIPSET_][MAXIPLEN])
{
    int count = 0;              // 用于计数非空IP
    for (int i = 0; i < MAXIPSET_; i++) {
        if (cn_ip[i][0] != '\0') {
            count++;            // 非空IP计数
        }
    }
    return count;
}

// 安全的清理IP地址集合，将所有位置置为空字符串
void clear_ip_set(char cn_ip[MAXIPSET_][MAXIPLEN])
{
    if (cn_ip == NULL) {
        return;                 // 如果指针无效，则直接返回
    }
    for (int i = 0; i < MAXIPSET_; i++) {
        memset(cn_ip[i], '\0', MAXIPLEN);
    }
}

// 检查 IP 是否已在缓存中并是否过期
int is_ip_in_cache(const char *ip)
{
    time_t now = time(NULL);    // 获取当前时间
    struct ip_cache_node *current = ip_cache_head;
    struct ip_cache_node *prev = NULL;

    while (current != NULL) {
        // 如果 IP 匹配并且未过期
        if (strcmp(current->ip, ip) == 0) {
            if (now - current->timestamp <= CACHE_TTL) {
                return 1;       // IP 在缓存中，且未过期
            } else {
                // 如果过期，从链表中移除这个节点
                if (prev == NULL) {
                    ip_cache_head = current->next;
                } else {
                    prev->next = current->next;
                }

                free(current);
                cache_size--;
                return 0;       // IP 过期，不再缓存
            }
        }

        prev = current;
        current = current->next;
    }
    return 0;                   // IP 不在缓存中
}

// 将新 IP 添加到缓存，若缓存过大则移除最早的 IP
void add_ip_to_cache(const char *ip)
{
    // 如果缓存大小超过限制，移除最早的 IP
    if (cache_size >= MAX_CACHE_SIZE) {
        struct ip_cache_node *current = ip_cache_head;
        struct ip_cache_node *prev = NULL;

        // 找到链表的最后一个节点
        while (current->next != NULL) {
            prev = current;
            current = current->next;
        }

        // 移除最后一个节点（最早的 IP）
        if (prev != NULL) {
            prev->next = NULL;
        } else {
            ip_cache_head = NULL;
        }
        free(current);
        cache_size--;
    }

    // 创建新的缓存节点并添加到链表头部
    struct ip_cache_node *new_node = (struct ip_cache_node *)malloc(sizeof(struct ip_cache_node));
    if (new_node == NULL) {
        perror("malloc");
        return;
    }
    strncpy(new_node->ip, ip, INET_ADDRSTRLEN);
    new_node->timestamp = time(NULL); // 记录当前时间
    new_node->next = ip_cache_head;
    ip_cache_head = new_node;
    cache_size++;
}

// 清理缓存链表，释放所有节点的内存
void free_ip_cache()
{
    struct ip_cache_node *current = ip_cache_head;
    while (current != NULL) {
        struct ip_cache_node *next = current->next;
        free(current);
        current = next;
    }

    ip_cache_head = NULL;
    cache_size = 0;
}

// 回调函数，在捕获到每个数据包时调用
void packet_handler(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{
    int ethernet_header_len = 14;
    struct ip *ip_header = (struct ip *)(packet + ethernet_header_len);
    char src_ip[INET_ADDRSTRLEN] = { 0 };

    // 地域白名单
    char _region_list[WHITELIST_IP_NUM][WHITELIST_IP_NUM] = { { 0 }, { 0 } };
    char _REGION_LIST[BUFFER] = { 0 };
    const char *REGION_ENV = NULL;

    char ipset_query_command[256] = { 0 };
    //char ip_query_command[256] = { 0 };

    // 定义 Response 结构体
    Response response;

    inet_ntop(AF_INET, &(ip_header->ip_src), src_ip, INET_ADDRSTRLEN);

    // 如果 IP 地址已在缓存中且未过期，则跳过查询
    if (is_ip_in_cache(src_ip)) {
        _printf(RED "IP:%s 已在缓存中，跳过查询\n" REDEND, src_ip);
        return;
    }

    // 如果ipset规则已经存在，则跳过查询
    snprintf(ipset_query_command, sizeof(ipset_query_command), "ipset test %s %s 2>/dev/null", RULE_NAME, src_ip);
    if (system(ipset_query_command) == 0) {
        // _printf(RED "Ipset 规则内已经存在 %s\n" REDEND, src_ip);
        return;
    }

    if (1 == is_ip_in_set(cn_ip, src_ip)) {
        _printf(RED "IP:%s 已经标记为国内，跳过！！！\n" REDEND, src_ip);

        for (int i = 0; i < MAXIPSET_; i++) {
            if (cn_ip[i][0] != '\0') {
                printf("%s ", cn_ip[i]);
            }
        }

        if (cn_ip_len(cn_ip) >= 100) {      // 清理集合
            clear_ip_set(cn_ip);
        }

        printf("cn_ip_len(cn_ip): %d\n", cn_ip_len(cn_ip));


        return;
    }

    // 执行查询并添加到缓存
    ip2region_area = ip2region("ip2region/ip2region.xdb", src_ip);
    if (ip2region_area == NULL) {
        _printf(RED "ip2region 解析地域错误\n" REDEND);
        return;
    }

    // 取环境变量
    REGION_ENV = getenv("REGION");
    if (REGION_ENV != NULL) {
        _printf("REGION: %s\n", REGION_ENV);
        strcpy(_REGION_LIST, REGION_ENV);
    } else {
        strcpy(_REGION_LIST, "局域网 内网 中国 ");
    }

    split_string(_REGION_LIST, " ", _region_list);
    if (isregion(ip2region_area, _region_list) == 1) { // 返回1表示在白名单列表
        ;
    } else {

        if (cache_size < MAX_CACHE_SIZE) // 缓存IP数不够预备设定值
        {
            sleep(1);
            _printf("缓存 IP 数 %d\n", cache_size);
        }

        add_ip_to_cache(src_ip); // 添加到缓存

        char *p = curl_get_area(src_ip);
        if (parse_json_to_struct(p, &response) == 0) { // 解析 JSON 到结构体            

            if (NULL == strstr(response.continent_country, "中国")) { // 这时是国外IP
                _printf(RED "%s %s\n" REDEND, src_ip, response.continent_country);
                add_ip_to_ipset(RULE_NAME, src_ip);
            } else {            // 这时是国内IP
                add_cn_ip(cn_ip, src_ip); // 添加国内IP到缓存
                _printf("IP: %s 离线库为国外, API 判断为国内, 标记为已处理!!!\n", src_ip);
            }

        } else {
            fprintf(stderr, "Failed to parse JSON.\n");
        }

        free(p);

    }

    if (ip2region_area != NULL) {
        free(ip2region_area);
        ip2region_area = NULL;
    }
    return;
}

void usage()
{
    printf("DenyIP version %s\n", _VERSION);
    puts("拒绝Linux服务器非大陆IP工具");
    puts("MAIL: aixiao@aixiao.me");
    puts("Date: 20241024");
    puts("  Usage: denyip [-d] [-i <interface>] [-s <start|stop>] [-h|-?]");
    puts("      -d --daemon         Daemon mode");
    puts("      -i --interface      interface (default eth0)");
    puts("      -l                  print iptables rule");
    puts("      -s --signal         regular signal (default start|stop)                        ");
    puts("                              start Enable Iptables rule");
    puts("                              stop  Disable Iptables rule");
    puts("      -h|-?               Help Information");
    puts("");

    exit(0);
}

void cleanup_(int signum)
{

    _printf("Received signal %d, cleaning up...\n", signum);

    // 终止子进程
    if (pid > 0) {
        kill(pid, SIGTERM);
    }

    // 释放共享内存
    if (RULE_NAME != NULL) {
        shmdt(RULE_NAME);
        shmctl(shmid, IPC_RMID, NULL);
    }

    // 在程序结束时，清理缓存链表
    free_ip_cache();

    if (ip2region_area != NULL) {
        free(ip2region_area);
        ip2region_area = NULL;
    }
    if (command_result != NULL) {
        free(command_result);
        command_result = NULL;
    }

    // 清理
    pcap_freecode(&fp);
    pcap_freealldevs(alldevs);  // 释放设备列表
    //pcap_close(handle);         // 关闭会话句柄

    // 退出主进程 
    exit(0);

    return;
}

int main(int argc, char **argv)
{
    // 注册 SIGTERM 信号处理函数
    signal(SIGTERM, cleanup_);

    int opt;
    char errbuf[PCAP_ERRBUF_SIZE] = { 0 }; // 错误缓冲区
    char protocol[] = "tcp";
    char interface[256] = "{ 0 }";
    char Ipset_Command[BUFFER] = { 0 };

    strcpy(interface, "eth0");
    memset(&alldevs, 0, sizeof(alldevs));
    memset(&device, 0, sizeof(device));
    memset(errbuf, 0, PCAP_ERRBUF_SIZE);

    int longindex = 0;
    char optstring[] = "di:s:lh?";
    static struct option longopts[] = {
        { "interface", required_argument, 0, 'i' },
        { "signal", required_argument, 0, 's' },
        { "daemon", no_argument, 0, 'd' },
        { "l", no_argument, 0, 'l' },
        { "help", no_argument, 0, 'h' },
        { "?", no_argument, 0, '?' },
        { 0, 0, 0, 0 }
    };

    while (-1 != (opt = getopt_long(argc, argv, optstring, longopts, &longindex))) {
        switch (opt) {
        case 'd':
            if (daemon(1, 1)) {
                perror("daemon");
            }
            break;
        case 'i':
            strcpy(interface, optarg);
            break;
        case 'l':
            system("iptables -L -v -n --line-numbers");
            exit(0);
            break;
        case 's':
            if (strcmp(optarg, "start") == 0) {
                memset(Ipset_Command, 0, BUFFER);
                // 将 MAXIPSET_RULT_NAME_NUM 替换为实际值
                snprintf(Ipset_Command, sizeof(Ipset_Command), "for n in $(seq 0 %d); do iptables -A INPUT -p tcp -m set --match-set root$n src -j DROP 2> /dev/null; done", MAXIPSET_RULT_NAME_NUM);
                system(Ipset_Command);
                exit(0);
            } else if (strcmp(optarg, "stop") == 0) {
                memset(Ipset_Command, 0, BUFFER);
                // 将 MAXIPSET_RULT_NAME_NUM 替换为实际值
                snprintf(Ipset_Command, sizeof(Ipset_Command), "for n in $(seq 0 %d); do iptables -D INPUT -p tcp -m set --match-set root$n src -j DROP 2> /dev/null; done", MAXIPSET_RULT_NAME_NUM);
                system(Ipset_Command);
                exit(0);
            } else {
                usage();
                exit(0);
            }

            break;
        case 'h':
        case '?':
            usage();
            exit(0);
            break;
        default:
            usage();
            exit(0);
            break;
        }
    }

    // 创建共享内存
    shmid = shmget(SHM_KEY, SHM_SIZE, IPC_CREAT | 0666);
    if (shmid < 0) {
        perror("shmget");
        exit(1);
    }

    // 连接共享内存到进程地址空间
    RULE_NAME = (char *)shmat(shmid, NULL, 0);
    if (RULE_NAME == (char *)-1) {
        perror("shmat");
        exit(1);
    }

    pid = fork();               // 创建子进程
    if (pid == 0)               // 子进程
    {

        int count = 0;
        snprintf(RULE_NAME, BUFFER, "root%d", RULE_NAME_NUMBER);

        if (create_ipset(RULE_NAME) != 0) {
            _printf("创建 IPSet %s 失败\n", RULE_NAME);
        }

        while (1) {
            //_printf("子进程当前 Ipset Rule 名 %s\n", RULE_NAME);

            count = get_ip_count_in_ipset(RULE_NAME);
            if (count >= 0) {
                _printf("子进程当前 Ipset Rule 名 %s, 数量: %d \n", RULE_NAME, count);
                if (count >= MAXIPSET && RULE_NAME_NUMBER <= MAXIPSET_RULT_NAME_NUM) // RULE_中的IP数量不超过MAXIPSET，并且集合不能超过 MAXIPSET_RULT_NAME_NUM 个
                {
                    RULE_NAME_NUMBER++;

                    snprintf(RULE_NAME, BUFFER, "root%d", RULE_NAME_NUMBER); // 更新规则名称
                    // 创建新的 IPSet
                    if (create_ipset(RULE_NAME) != 0) {
                        _printf("创建 IPSet %s 失败\n", RULE_NAME);
                    } else {
                        char iptables_command[256];
                        sprintf(iptables_command, "iptables -I INPUT -m set --match-set %s src -j DROP", RULE_NAME);
                        system(iptables_command);
                    }

                }

                if (RULE_NAME_NUMBER == MAXIPSET_RULT_NAME_NUM) {
                    _printf("已达到最大规则数限制，停止创建新规则。\n");
                    _printf("请手动清理Ipset规则\n");
                }
            }

            sleep(9);           // 每 3 秒检查一次
        }
    }

    // 查找可用的网络设备
    if (pcap_findalldevs(&alldevs, errbuf) == -1) {
        fprintf(stderr, "无法找到设备: %s\n", errbuf);
        return 1;
    }

    // 打印可用设备列表
    _printf("可用的设备:\n");
    for (device = alldevs; device != NULL; device = device->next) {
        _printf("设备: %s\n", device->name);
    }

    // 打开设备以进行数据包捕获
    handle = pcap_open_live(interface, BUFSIZ, 1, 1000, errbuf);
    if (handle == NULL) {
        fprintf(stderr, "无法打开设备 %s: %s\n", interface, errbuf);
        pcap_freealldevs(alldevs);
        return 1;
    }

    // 编译过滤器
    if (pcap_compile(handle, &fp, protocol, 0, PCAP_NETMASK_UNKNOWN) == -1) {
        fprintf(stderr, "无法编译过滤器 %s: %s\n", protocol, pcap_geterr(handle));
        pcap_close(handle);
        pcap_freealldevs(alldevs);
        return 1;
    }

    // 设置过滤器
    if (pcap_setfilter(handle, &fp) == -1) {
        fprintf(stderr, "无法设置过滤器 %s: %s\n", protocol, pcap_geterr(handle));
        pcap_freecode(&fp);
        pcap_close(handle);
        pcap_freealldevs(alldevs);
        return 1;
    }

    // 开始捕获数据包
    if (pcap_loop(handle, 0, packet_handler, NULL) < 0) {
        fprintf(stderr, "捕获数据包时出错: %s\n", pcap_geterr(handle));
        pcap_freecode(&fp);
        pcap_close(handle);
        pcap_freealldevs(alldevs);
        return 1;
    }

    // 在程序结束时，清理缓存链表
    free_ip_cache();

    // 清理
    pcap_freecode(&fp);
    pcap_freealldevs(alldevs);  // 释放设备列表
    pcap_close(handle);         // 关闭会话句柄

    return 0;
}