#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <sys/wait.h>
#include <arpa/inet.h>
#include <sys/resource.h>
#include <signal.h>
#include <sys/prctl.h>

#include "ip2region.h"
#include "qqwry.h"
#include "common.h"

#define RED "\033[31m"
#define RESET "\033[0m"
#define WHITELIST_IP_NUM 1024
#define MAXIPSET 65534

char *xdb_path = "ip2region.xdb";
pid_t pid1, pid2;               // 保存子进程的 PID


int is_valid_ip(const char *ip)
{
    struct sockaddr_in sa;
    // 尝试将字符串转换为IPv4地址
    int result = inet_pton(AF_INET, ip, &(sa.sin_addr));

    return result != 0;
}

int nice_(int increment)
{
    int oldprio = getpriority(PRIO_PROCESS, getpid());
    printf("%d\n", oldprio);

    return setpriority(PRIO_PROCESS, getpid(), oldprio + increment);
}

// 判断命令是否存在
int command_exists(const char *command)
{
    char buffer[BUFFER];
    snprintf(buffer, sizeof(buffer), "%s > /dev/null 2>&1", command);
    int status = system(buffer);

    return (status == 0);
}

void denyip_help()
{
    puts("      DenyIp");
    puts("Linux system firewall, reject non-Chinese IP");
    puts("Email: aixiao@aixiao.me");
    puts("Version: 0.1");
    puts("Usage: ./denyip [-i eth0|-h|-?] [start|stop] ");
    puts("Options:");
    puts("      stop    Enable firewall rules");
    puts("      start   Disable firewall rules");

    puts("Parameters:");
    puts("      -h|?    Help info ");
    puts("      -i      interface name");
    puts("");
    puts("");

    exit(0);
}

/* 处理僵尸进程 */
void sigchld_handler(int signal)
{
    while (waitpid(-1, NULL, WNOHANG) > 0) ;
}

void kill_tcpdump_processes()
{
    int result = system("pkill tcpdump"); // 用 pkill 命令终止所有 tcpdump 进程
    if (result == -1) {
        perror("Failed to kill tcpdump processes");
    }

    return ;
}

// 进程重启功能
void restart_process(pid_t pid1, pid_t pid2, char *argv[])
{
    // 发送 SIGTERM 信号终止两个子进程
    kill(pid1, SIGTERM);
    kill(pid2, SIGTERM);

    // 等待子进程完全退出
    waitpid(pid1, NULL, 0);
    waitpid(pid2, NULL, 0);

    // 终止 tcpdump 进程
    kill_tcpdump_processes();

    // 使用 execvp 重新启动程序自身
    printf("重启进程...\n");
    execvp(argv[0], argv);      // 重新启动程序
    perror("execvp failed");    // 如果 execvp 出错
    exit(EXIT_FAILURE);

    return ;
}

void cleanup_(int signum)
{
    int r;
    printf("Received signal %d, cleaning up...\n", signum);

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

    // 终止所有 tcpdump 进程
    r = system("pkill tcpdump");
    if (r == 0) {
        puts("pkill tcpdump");
    }

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

    return ;
}

int main(int argc, char *argv[])
{
    signal(SIGCHLD, sigchld_handler); // 防止子进程变成僵尸进程

    // 主进程设置
    //prctl(PR_SET_PDEATHSIG, SIGTERM);

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

    char interface[BUFFER] = { 0 };
    strcpy(interface, "eth0");

    int r;

    // 参数处理
    if (argc == 2) {
        if (0 == strcmp(argv[1], "start")) {
            if ((r = system("iptables -A INPUT -p tcp -m set --match-set root src -j DROP")) == -1) {
                puts("\"iptables -A INPUT -p tcp -m set --match-set root src -j DROP\" Error!");
            }
            exit(0);
        } else if (0 == strcmp(argv[1], "stop")) {
            if ((r = system("iptables -D INPUT -p tcp -m set --match-set root src -j DROP")) == -1) {
                puts("\"iptables -D INPUT -p tcp -m set --match-set root src -j DROP\" Error!");
            }
            exit(0);
        } else if (0 == strcmp(argv[1], "-h")) {
            denyip_help();
        } else if (0 == strcmp(argv[1], "-?")) {
            denyip_help();
        }
    }
    if (argc == 3) {
        if (0 == strcmp(argv[1], "-i")) {
            strcpy(interface, argv[2]);
        }
    }
    // 判断运行用户禁止非root用户运行
    if (geteuid() == 0) {
        ;
    } else {
        printf("This process is not running as root.\n");
        printf("\n");
        exit(-1);
    }

    // 判断网卡是否存在
    char command_ifconfig[BUFFER + 20] = { 0 };
    snprintf(command_ifconfig, BUFFER + 20, "ifconfig %s", interface);
    // 判断必要命令是否存在
    if (command_exists(command_ifconfig)) {
        ;
    } else {
        puts("The network card does not exist!");
        exit(-1);
    }

    // 后台运行
    if (daemon(1, 1)) {
        perror("daemon");
        return -1;
    }

    // 进程优先级
    if (-1 == (nice_(-20)))
        perror("nice_");

    // 哈希集合
    if ((r = system("ipset create root hash:ip > /dev/null 2>&1")) != -1) {
        ;
    }

    // 判断必要命令是否存在
    if (command_exists("which tcpdump")) {
        ;
    } else {
        r = system("yum -y install tcpdump > /dev/null 2>&1");
        r = system("apt -y install tcpdump > /dev/null 2>&1");
    }


    // 子进程
    pid1 = fork();              // 创建子进程
    if (pid1 == 0) {
        while (1) {
            FILE *fp = popen("ipset list root | grep \"Number of entries\" | cut -d : -f 2 | xargs", "r");
            char line[BUFFER] = { 0 };
            while (fgets(line, sizeof(line), fp) != NULL) {
                line[strcspn(line, "\n")] = '\0';
            }

            if (atoi(line) >= MAXIPSET) {
                r = system("ipset flush root");
            }

            printf("%s\n", line);
            pclose(fp);
            sleep(3);
        }
    }

    // 子进程
    pid2 = fork();              // 创建子进程
    if (pid2 == 0) {
        // 缓冲区用于存储每行的输出
        char line[BUFFER];
        // 要执行的命令
        char command_tcpdump[BUFFER + 256] = { 0 };
        snprintf(command_tcpdump, BUFFER + 256, "tcpdump -i %s -n 'tcp' | awk '{print $3}' | cut -d '.' -f 1-4", interface);

        // 地域白名单
        char _region_list[WHITELIST_IP_NUM][WHITELIST_IP_NUM] = { { 0 }, { 0 } };
        char qqwry_region_list[WHITELIST_IP_NUM][WHITELIST_IP_NUM] = { { 0 }, { 0 } };
        char _REGION_LIST_COPY[BUFFER] = { 0 };
        char QQWRY_REGION_LIST_COPY[BUFFER] = { 0 };

        if (access(xdb_path, F_OK) == -1) { // 判断 ip2region 地址定位库是否存在
            xdb_path = "ip2region/ip2region.xdb";
            if (access(xdb_path, F_OK) == -1) {
                printf("ip2region.xdb DOESN'T EXIST!\n");
                return -1;
            }
        }

        // 打开管道来执行命令
        FILE *fp = popen(command_tcpdump, "r");
        if (fp == NULL) {
            perror("popen failed");
            return 1;
        }

        // 逐行读取命令输出
        while (fgets(line, sizeof(line), fp) != NULL) {
            line[strcspn(line, "\n")] = '\0';

            if (is_valid_ip(line)) {

                char *qqwry_region = qqwry_(line);
                if (qqwry_region == NULL) {
                    printf("qqwry 解析地域错误\n");
                    continue;
                }
                char *area = ip2region(xdb_path, line);
                if (area == NULL) {
                    printf("ip2region 解析地域错误\n");
                    continue;
                }
                // 取环境变量
                const char *REGION_ENV = getenv("REGION");
                if (REGION_ENV != NULL) {
                    printf("REGION: %s\n", REGION_ENV);
                    strcpy(_REGION_LIST_COPY, REGION_ENV);
                    strcpy(QQWRY_REGION_LIST_COPY, REGION_ENV);
                } else {
                    strcpy(_REGION_LIST_COPY, "局域网 内网 中国 ");
                    strcpy(QQWRY_REGION_LIST_COPY, "局域网 内网 中国 ");
                }
                //printf("REGION_LIST : %s\n", _REGION_LIST_COPY);

                split_string(QQWRY_REGION_LIST_COPY, " ", qqwry_region_list); // 分割后存储在 qqwry_region_list
                if (isregion(qqwry_region, qqwry_region_list) == 1) { // 返回1表示在白名单列表
                    ;
                } else {
                    split_string(_REGION_LIST_COPY, " ", _region_list);
                    if (isregion(area, _region_list) == 1) { // 返回1表示在白名单列表
                        ;
                    } else {
                        char ipquery_command[BUFFER + 100] = { 0 };
                        snprintf(ipquery_command, BUFFER + 100, "./IP_region_query/ipquery %s", line);
                        FILE *fp = popen(ipquery_command, "r");
                        if (fp == NULL) {
                            perror("popen failed");
                            return 1;
                        }

                        // 创建足够大的缓冲区来存储命令输出
                        char buffer[1024 * 2]; // 2KB 缓冲区
                        size_t bytesRead = fread(buffer, 1, sizeof(buffer) - 1, fp);
                        buffer[bytesRead] = '\0';

                        pclose(fp);
                        sleep(1);

                        char *p = strstr(buffer, "中国");
                        if (p == NULL) {
                            printf("%s %s", line, buffer);

                            char command_ipset[BUFFER + 256] = { 0 };
                            snprintf(command_ipset, sizeof(command_ipset), "ipset add root %s > /dev/null 2>&1", line);
                            int r = system(command_ipset);
                            if (r == -1) {
                                perror("system command failed");
                            }
                        }
                    }
                }
                free(qqwry_region);
                free(area);

            } else {            // 是正确IP
                printf("%s is not a valid IPv4 address.\n", line);
            }
        }                       // while

        // 关闭管道
        pclose(fp);
    }


    // 父进程
    int iteration = 0;
    // 主进程循环，检查子进程运行情况
    while (1) {
        iteration++;

        if (iteration >= 3600) {
            iteration = 0;
            printf("准备重启进程...\n");
            restart_process(pid1, pid2, argv);
        }

        sleep(1);               // 每次检查间隔1秒
    }

    return 0;
}