#include "cap.h"
#include "common.h"
#include "libipset.h"
#include "libcurl.h"
#include "cache.h"
#include "leak_detector_c/leak_detector_c.h"

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

pid_t pids[MAX_CHILDREN] = {-1};    // 子进程全局PID

char *ip2region_area = NULL;        // ip2region 解析结果

// 共享内存
#define SHM_SIZE (MAXIPSET_ * MAXIPLEN)
int shmid[3] = {-1};

char (*cn_ip)[MAXIPLEN] = NULL; // 用于存储国内 IP 地址
char (*whitelist_ip)[MAXIPLEN] = NULL;  // 最终要这样定义
char *RULE_NAME = NULL;         // 共享内存
int RULE_NAME_NUMBER = 0;       // ipset 集合集合数


void Processing_IP_addresses(char *src_ip)
{
    // 地域白名单
    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 };
    // 定义 Response 结构体
    Response response;


    if (1 == is_ip_in_set(whitelist_ip, src_ip)) {
        _printf(RED "IP:%s 白名单IP, 跳过！\n" REDEND, src_ip);
        //printf("%s %d\n", whitelist_ip[0], cn_ip_len(whitelist_ip));

        return;
    }
    

    // 如果 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);

        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 = CurlGetIpArea(src_ip);
        if (parse_json_to_struct(p, &response) == 0) { // 解析 JSON 到结构体            

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

                if (append_string_to_file("cn.txt", src_ip) != 0) {
                    _printf("append_string_to_file() Error!!!\n");
                }
            }
        } else {
            fprintf(stderr, "Failed to parse JSON.\n");
        }

        if (p != NULL) {
            free(p);
        }
    }

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

    // 实时检测内存泄漏
    dump_mem_leak();
    return;
}

// 回调函数，在捕获到每个数据包时调用
void packet_handler(u_char *args, const struct pcap_pkthdr *header, const u_char *packet)
{
    // 定义以太网帧头部的长度 (目标 MAC 6字节 + 源 MAC 6字节 + 类型 2字节 = 14字节)
    int ethernet_header_len = 14;

    // 假设数据包是 Ethernet II 帧封装的 IP 包
    // 计算 IP 头部的起始位置：跳过以太网头部
    // 将该位置的指针强制转换为 struct ip 指针
    struct ip *ip_header = (struct ip *)(packet + ethernet_header_len);
    // 检查 IP 协议版本号
    // ip_v 是 struct ip 结构中的版本号字段 (通常占 4 bits)
    if (ip_header->ip_v != 4)
        return;                 // 只处理 IPv4
    // 定义一个字符数组用于存储点分十进制格式的源 IP 地址字符串
    // INET_ADDRSTRLEN 是 <arpa/inet.h> 中定义的足够存储 IPv4 地址字符串的最大长度 (包括结尾的 '\0')
    char src_ip[INET_ADDRSTRLEN] = { 0 };

    // 将网络字节序 (二进制) 的 IPv4 源地址转换为点分十进制表示的字符串
    // AF_INET: 指定地址族为 IPv4
    // &(ip_header->ip_src): 指向 struct ip 结构中源地址字段 (struct in_addr 类型) 的指针
    // src_ip: 存储转换后字符串的缓冲区
    // INET_ADDRSTRLEN: 缓冲区的长度
    inet_ntop(AF_INET, &(ip_header->ip_src), src_ip, INET_ADDRSTRLEN);



    Processing_IP_addresses(src_ip);

    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("      -f --protocol       过滤器 [\"tcp\" | \"udp\" | \"tcp or udp\"] (default \"tcp\")");
    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 detach_and_delete_shm(void *addr, int shmid, int delete_flag) {
    if (addr && addr != (void *)-1) {
        shmdt(addr);
    }
    if (delete_flag && shmid >= 0) {
        shmctl(shmid, IPC_RMID, NULL);
    }
}

void cleanup_(int signum)
{

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

    // 终止子进程
    for (int i = 0; i < MAX_CHILDREN; i++) {
        if (pids[i] > 0) {
            printf("Terminating child process %d...\n", pids[i]);
            kill(pids[i], SIGTERM);     // 优雅终止
            sleep(1);                   // 等待其响应
            kill(pids[i], SIGKILL);     // 若仍未结束，强制终止
            waitpid(pids[i], NULL, 0);  // 回收资源
        }
    }

    // 释放共享内存
    detach_and_delete_shm(RULE_NAME, shmid[0], 1);
    detach_and_delete_shm(whitelist_ip, shmid[1], 1);
    detach_and_delete_shm(cn_ip, shmid[2], 1);

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

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


    // 清理
    pcap_freecode(&fp);
    pcap_freealldevs(alldevs);  // 释放设备列表
    if (handle)
        pcap_close(handle);

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

    return;
}

void *create_and_attach_shm(const char *path, int proj_id, size_t size, int *shmid_out) {
    key_t key = ftok(path, proj_id);
    if (key == -1) {
        perror("ftok");
        exit(1);
    }

    int shmid = shmget(key, size, IPC_CREAT | 0666);
    if (shmid < 0) {
        perror("shmget");
        exit(1);
    }

    void *addr = shmat(shmid, NULL, 0);
    if (addr == (void *)-1) {
        perror("shmat");
        exit(1);
    }

    if (shmid_out) *shmid_out = shmid;
    return addr;
}


int main(int argc, char **argv)
{
    atexit(report_mem_leak);

    // 注册 SIGTERM 信号处理函数
    signal(SIGINT, cleanup_);
    signal(SIGTERM, cleanup_);

    int opt;
    char errbuf[PCAP_ERRBUF_SIZE] = { 0 }; // 错误缓冲区
    char protocol[256] = "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:f:s:lh?";
    static struct option longopts[] = {
        { "interface", required_argument, 0, 'i' },
        { "protocol", required_argument, 0, 'f' },
        { "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 'f':
            strcpy(protocol, 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;
        }
    }

    // 创建共享内存
    RULE_NAME = (char *)create_and_attach_shm("/tmp", 'A', 1024, &shmid[0]);
    whitelist_ip = (char (*)[MAXIPLEN])create_and_attach_shm("/tmp", 'B', SHM_SIZE, &shmid[1]);
    cn_ip = (char (*)[MAXIPLEN])create_and_attach_shm("/tmp", 'C', SHM_SIZE, &shmid[2]);



    // 读取白名单
    if (1 == file_exists_access("whitelist.txt")) {
        int line_count = 0;
        int result = read_file_to_array("whitelist.txt", whitelist_ip, &line_count);
        if (result != 0) {
            fprintf(stderr, "Failed to read file with error code: %d\n", result);
            return 1;
        }

        for (int i = 0; i < line_count; i++) {
            printf("Line %d: %s\n", i + 1, whitelist_ip[i]);
        }
    }

    // 读取缓存
    if (1 == file_exists_access("cn.txt")) {
        int line_count = 0;
        int result = read_file_to_array("cn.txt", cn_ip, &line_count);
        if (result != 0) {
            fprintf(stderr, "Failed to read file with error code: %d\n", result);
            return 1;
        }

        char *local_addr = GetLocalAddr("https://inet-ip.aixiao.me/");
        if (local_addr == NULL) {
            perror("GetLocalAddr()");
            return 1;
        }
        remove_char(local_addr, '\n');
        printf("Local Address: %s\n", local_addr);

        strcpy(cn_ip[line_count], local_addr);
        line_count++;

        printf("Read %d lines from file:\n", line_count);
        for (int i = 0; i < line_count; i++) {
            printf("Line %d: %s\n", i + 1, cn_ip[i]);
        }

        free(local_addr);
    }


    pids[0] = fork();
    if ( 0 == pids[0] )  // 子进程
    {
        while(1) {

            if (1 == file_exists_access("whitelist.txt")) {
                clear_ip_set(whitelist_ip); // 清空白名单IP集合
                int line_count = 0;
                int result = read_file_to_array("whitelist.txt", whitelist_ip, &line_count);
                if (result != 0) {
                    fprintf(stderr, "Failed to read file with error code: %d\n", result);
                    return 1;
                }

                for (int i = 0; i < line_count; i++) {
                    printf("Line %d: %s\n", i + 1, whitelist_ip[i]);
                }
            }
            
            for (int i = 0; i < MAXIPSET_; i++) {
                if (cn_ip[i][0] != '\0') {
                    printf("%s ", cn_ip[i]);
                }
            }
            printf("cn_ip_len(cn_ip): %d\n", cn_ip_len(cn_ip));

            if (cn_ip_len(cn_ip) >= 10240) {    // 清理集合
                clear_ip_set(cn_ip);
                truncate_file("cn.txt"); // 清空文件
            }

            sleep(10);           // 每 60 秒检查一次
        }
    }

    pids[1] = fork();
    if ( 0 == pids[1] )               // 子进程
    {
        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)
        {
            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);
                         */
                        add_iptables_rule(RULE_NAME);
                    }
                }

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

            sleep(7);           // 每 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;
}