DenyIP/libipset.c

#include "libipset.h"

// 自定义输出处理函数
int custom_output_handler(struct ipset_session *session, void *p, const char *msg, ...)
{
    int *ip_count = (int *)p;
    char buffer[BUFFER];

    // 格式化输出消息
    va_list args;
    va_start(args, msg);
    vsnprintf(buffer, sizeof(buffer), msg, args);
    va_end(args);

    // 输出调试信息
    //printf("Buffer: %s\n", buffer); // 调试输出

    if (strlen(buffer) > 1) {
        char temp[BUFFER];
        char *p1 = strstr(buffer, "Number of entries:");
        if (p1 != NULL) {
            char *p2 = strstr(p1, "\n");
            if (p2 != NULL) {
                size_t len = p2 - p1; // 计算长度
                if (len < BUFFER) {
                    strncpy(temp, p1, len);
                    temp[len] = '\0'; // 确保字符串结束
                    //puts(temp);

                    // 查找冒号并提取数量
                    char *p3 = strstr(temp, ":");
                    if (p3 != NULL) {
                        // 提取数字
                        int count = atoi(p3 + 1); // 从冒号后面开始转换
                        *ip_count += count; // 更新计数
                    }
                }
            }
        }
    }

    return 0;
}

int custom_output_handler_(struct ipset_session *session, void *p, const char *msg, ...)
{
    char buffer[BUFFER];

    // 格式化输出消息
    va_list args;
    va_start(args, msg);
    vsnprintf(buffer, sizeof(buffer), msg, args);
    va_end(args);

    // 输出调试信息
    //printf("Buffer: %s\n", buffer); // 调试输出

    return 0;
}

// 自定义错误处理函数
int custom_error_handler(struct ipset *ipset, void *p, int errnum, const char *msg, ...)
{
    va_list args;
    va_start(args, msg);
    fprintf(stderr, "自定义错误: ");
    vfprintf(stderr, msg, args);
    va_end(args);
    return 0;                   // 返回0表示处理成功
}

// 创建 IPSet，如果它不存在
int create_ipset(char *set_name)
{
    struct ipset *ipset = ipset_init();
    if (!ipset) {
        fprintf(stderr, "Failed to initialize IPSet\n");
        return -1;
    }

    ipset_load_types();

    struct ipset_session *session = ipset_session(ipset);
    if (!session) {
        fprintf(stderr, "Failed to create IPSet session\n");
        ipset_fini(ipset);
        return -1;
    }
    // 设置自定义错误和输出处理函数
    if (ipset_custom_printf(ipset, custom_error_handler, NULL, custom_output_handler_, NULL) != 0) {
        fprintf(stderr, "设置自定义打印函数失败。\n");
        ipset_fini(ipset);
        return -1;
    }
    // 创建集合
    char *args[] = { "ipset", "create", set_name, "hash:ip", NULL };
    if (ipset_parse_argv(ipset, 4, args) != 0) {
        ipset_fini(ipset);
        return -1;
    }

    ipset_fini(ipset);

    return 0;
}

// 向指定的 ipset 添加 IP
int add_ip_to_ipset(char *set_name, char *ip)
{
    struct ipset *ipset = NULL;
    struct ipset_session *session = NULL;

    // 初始化 ipset 和 session
    ipset = ipset_init();
    if (!ipset) {
        fprintf(stderr, "初始化 IPSet 失败\n");
        return -1;              // 返回 -1，但不退出
    }

    ipset_load_types();

    session = ipset_session(ipset);
    if (!session) {
        fprintf(stderr, "创建 IPSet 会话失败\n");
        ipset_fini(ipset);
        return -1;              // 返回 -1，但不退出
    }
    // 设置 IPSet 名称
    if (ipset_session_data_set(session, IPSET_SETNAME, set_name) != 0) {
        fprintf(stderr, "设置 IPSet 名称失败\n");
        ipset_session_fini(session);
        ipset_fini(ipset);
        return -1;              // 返回 -1，但不退出
    }
    // 设置自定义错误和输出处理函数
    if (ipset_custom_printf(ipset, custom_error_handler, NULL, custom_output_handler_, NULL) != 0) {
        fprintf(stderr, "设置自定义打印函数失败。\n");
        ipset_fini(ipset);
        return -1;
    }
    // 将ip添加到集合
    char *args[] = { "ipset", "add", set_name, ip, NULL };
    if (ipset_parse_argv(ipset, 4, args) != 0) {
        ipset_fini(ipset);
        return -1;
    }

    ipset_fini(ipset);

    return 0;                   // 始终返回 0，表示执行成功
}

// 定义一个函数来清空指定名称的 ipset 集合
int flush_ipset(char *set_name)
{
    struct ipset *ipset = ipset_init();
    if (!ipset) {
        fprintf(stderr, "Failed to initialize IPSet\n");
        return -1;
    }

    ipset_load_types();

    // 设置自定义错误和输出处理函数
    if (ipset_custom_printf(ipset, custom_error_handler, NULL, custom_output_handler_, NULL) != 0) {
        fprintf(stderr, "设置自定义打印函数失败。\n");
        ipset_fini(ipset);
        return -1;
    }
    // 清空集合
    char *args[] = { "ipset", "flush", set_name, NULL };
    if (ipset_parse_argv(ipset, 3, args) != 0) {
        ipset_fini(ipset);
        return -1;
    }

    printf("IPSet %s flushed successfully\n", set_name);
    ipset_fini(ipset);
    return 0;
}

// 获取指定 IPSet 中的 IP 数量
int get_ip_count_in_ipset(char *set_name)
{
    int ip_count = 0;
    struct ipset *ipset = ipset_init();
    if (!ipset) {
        fprintf(stderr, "初始化 ipset 失败。\n");
        return -1;
    }

    ipset_load_types();

    // 设置自定义错误和输出处理函数
    if (ipset_custom_printf(ipset, custom_error_handler, NULL, custom_output_handler, &ip_count) != 0) {
        fprintf(stderr, "设置自定义打印函数失败。\n");
        ipset_fini(ipset);
        return -1;
    }
    // 列出集合
    char *args[] = { "ipset", "list", set_name, NULL };
    if (ipset_parse_argv(ipset, 3, args) != 0) {
        ipset_fini(ipset);
        return -1;
    }
    // 释放资源
    ipset_fini(ipset);
    return ip_count;
}


void add_iptables_rule(const char *rule_name)
{
    char iptables_command[256];

    // 使用 snprintf 避免溢出
    int written = snprintf(iptables_command, sizeof(iptables_command),
                           "iptables -I INPUT -m set --match-set %s src -j DROP", rule_name);

    if (written < 0 || written >= sizeof(iptables_command)) {
        fprintf(stderr, "Error: iptables command is too long.\n");
        return;
    }

    // 执行命令并检查返回值
    int ret = system(iptables_command);
    if (ret != 0) {
        fprintf(stderr, "Failed to execute iptables command: %s\n", iptables_command);
    }
}

/*
int main()
{
    int r;
    char *set_name = "root";
    char *ip = "1.1.1.2";

    // 确保 IPSet 已存在
    r = create_ipset(set_name);
    printf("create_ipset %d\n", r);

    r = flush_ipset(set_name);
    printf("flush_ipset %d\n", r);

    // 尝试添加 IP
    r = add_ip_to_ipset(set_name, ip);
    printf("add_ip_to_ipset %d\n", r);

    // 获取并打印 IPSet 中的 IP 数量
    int count = get_ip_count_in_ipset(set_name);
    if (count >= 0) {
        printf("IPSet %s 中的 IP 数量: %d\n", set_name, count);
    }

    return 0;
}
*/