example.cpp

#include "Kitsune.h"
#include <algorithm> // 用于排序和计算统计量
#include <chrono>
#include <cmath>
#include <filesystem> // 用于检查文件是否存在
#include <fstream>
#include <iostream>
#include <limits>
#include <memory>
#include <numeric> // 用于求和
#include <string>
#include <vector>

// 用于解压缩文件的函数声明
void extractZipFile(const std::string &zipFileName);

// 检查文件是否存在
bool fileExists(const std::string &filename) {
    return std::filesystem::exists(filename);
}

// 检查文件是否为支持的捕获文件格式
bool isCaptureFile(const std::string &filename) {
    std::string extension = filename.substr(filename.find_last_of(".") + 1);
    std::transform(extension.begin(), extension.end(), extension.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return (extension == "pcap" || extension == "pcapng" || extension == "tsv");
}

// 获取文件扩展名（小写形式）
std::string getFileExtension(const std::string &filename) {
    std::string extension = filename.substr(filename.find_last_of(".") + 1);
    std::transform(extension.begin(), extension.end(), extension.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return extension;
}

// 检查是否安装了tshark
bool isTsharkInstalled() {
#ifdef _WIN32
    std::string tshark_path = "C:\\Program Files\\Wireshark\\tshark.exe";
    return std::filesystem::exists(tshark_path);
#else
    std::string command = "which tshark > /dev/null 2>&1";
    return system(command.c_str()) == 0;
#endif
}

// 计算均值
double mean(const std::vector<double> &data) {
    double sum = std::accumulate(data.begin(), data.end(), 0.0);
    return sum / data.size();
}

// 计算中位数
double median(std::vector<double> data) {
    size_t n = data.size();
    if (n == 0) return 0;
    
    std::sort(data.begin(), data.end());
    if (n % 2 == 0) {
        return (data[n/2 - 1] + data[n/2]) / 2.0;
    } else {
        return data[n/2];
    }
}

// 计算标准差
double stdDev(const std::vector<double> &data, double mean_val) {
    double sq_sum =
        std::inner_product(data.begin(), data.end(), data.begin(), 0.0);
    return std::sqrt(sq_sum / data.size() - mean_val * mean_val);
}

// 计算MAD（中位数绝对偏差）
double mad(const std::vector<double> &data, double med) {
    std::vector<double> deviations;
    for (double val : data) {
        deviations.push_back(std::abs(val - med));
    }
    return median(deviations);
}

// 使用MAD方法进行异常检测
std::vector<size_t> detectAnomaliesMAD(
    const std::vector<double> &rmse_values, 
    size_t start_idx,
    double mad_factor = 3.0) {
    
    // 仅使用训练之后的数据进行异常检测
    std::vector<double> test_data(rmse_values.begin() + start_idx, rmse_values.end());
    
    // 使用中位数绝对偏差 (MAD) 方法
    double med = median(test_data);
    double mad_value = mad(test_data, med);
    
    std::vector<size_t> mad_anomalies;
    for (size_t i = 0; i < test_data.size(); i++) {
        if (mad_value > 0 && std::abs(test_data[i] - med) > mad_factor * mad_value) {
            mad_anomalies.push_back(start_idx + i);
        }
    }
    
    return mad_anomalies;
}

// 计算对数正态分布的参数
struct LogNormalParams {
    double mu;      // 对数均值
    double sigma;   // 对数标准差
};

LogNormalParams fitLogNormal(const std::vector<double>& data) {
    std::vector<double> log_data;
    log_data.reserve(data.size());
    
    // 计算数据的自然对数
    for (double x : data) {
        if (x > 0) {  // 确保数据为正
            log_data.push_back(std::log(x));
        }
    }
    
    // 计算对数数据的均值和标准差
    double log_mean = mean(log_data);
    double log_std = stdDev(log_data, log_mean);
    
    return {log_mean, log_std};
}

// 计算对数正态分布的分位数
double logNormalQuantile(double mu, double sigma, double p) {
    // 使用标准正态分布的分位数
    double z = std::sqrt(2) * std::erf(2 * p - 1);
    return std::exp(mu + sigma * z);
}

// 使用统计概率方法进行异常检测
std::vector<size_t> detectAnomaliesStatistical(
    const std::vector<double>& rmse_values,
    size_t start_idx,
    double confidence_level = 0.99) {
    
    // 使用训练数据拟合对数正态分布
    std::vector<double> training_data(rmse_values.begin() + start_idx, rmse_values.begin() + start_idx * 2);
    LogNormalParams params = fitLogNormal(training_data);
    
    // 计算阈值（基于置信水平）
    double threshold = logNormalQuantile(params.mu, params.sigma, confidence_level);
    
    std::vector<size_t> anomalies;
    for (size_t i = start_idx; i < rmse_values.size(); i++) {
        if (rmse_values[i] > threshold) {
            anomalies.push_back(i);
        }
    }
    
    return anomalies;
}


int main() {

    // 文件位置 - 支持pcap和pcapng格式
    std::string captureFile = "mirai.pcap"; // 默认文件
    std::string zipFile = "mirai.zip";      // 压缩文件
    std::vector<std::string> supportedFiles;

    // 查找可能的捕获文件
    for (const auto &entry : std::filesystem::directory_iterator(".")) {
        std::string filename = entry.path().filename().string();
        if (isCaptureFile(filename)) {
            supportedFiles.push_back(filename);
        }
    }

    // 如果找到支持的文件，提示用户选择
    if (!supportedFiles.empty()) {
        std::cout << "找到以下捕获文件:" << std::endl;
        for (size_t i = 0; i < supportedFiles.size(); i++) {
            std::cout << i + 1 << ". " << supportedFiles[i] << " ("
                      << getFileExtension(supportedFiles[i]) << "格式)"
                      << std::endl;
        }

        std::cout << "请选择要分析的文件 (1-" << supportedFiles.size()
                  << "), 或按Enter使用默认文件 (" << captureFile << "): ";

        std::string input;
        std::getline(std::cin, input);
        if (!input.empty()) {
            try {
                int choice = std::stoi(input);
                if (choice >= 1 &&
                    choice <= static_cast<int>(supportedFiles.size())) {
                    captureFile = supportedFiles[choice - 1];
                }
            } catch (...) {
                // 使用默认选择
            }
        }
    }

    // 检查选择的文件是否存在
    if (!fileExists(captureFile)) {
        std::cout << "所选文件不存在，尝试检查默认文件或解压..." << std::endl;

        // 检查默认pcap文件是否存在
        if (!fileExists("mirai.pcap")) {
            std::cout << "未找到默认数据文件，正在解压样本捕获文件..."
                      << std::endl;
            extractZipFile(zipFile);

            // 解压后重新检查支持的文件
            supportedFiles.clear();
            for (const auto &entry : std::filesystem::directory_iterator(".")) {
                std::string filename = entry.path().filename().string();
                if (isCaptureFile(filename)) {
                    supportedFiles.push_back(filename);
                    if (captureFile.empty()) {
                        captureFile = filename; // 使用第一个找到的文件
                    }
                }
            }

            if (captureFile.empty()) {
                std::cerr << "解压后未找到可用的捕获文件，退出。" << std::endl;
                return 1;
            }
        } else {
            captureFile = "mirai.pcap"; // 使用默认文件
        }
    }

    // 检查tshark是否可用（处理pcap/pcapng需要）
    std::string extension = getFileExtension(captureFile);
    if ((extension == "pcap" || extension == "pcapng") &&
        !isTsharkInstalled()) {
        std::cerr << "警告: 未检测到tshark（Wireshark命令行工具）。"
                  << std::endl;
        std::cerr << "处理" << extension << "文件需要安装Wireshark。"
                  << std::endl;
        std::cerr << "请安装Wireshark后重试，或使用已转换的.tsv文件。"
                  << std::endl;
        return 1;
    }

    std::cout << "将使用捕获文件: " << captureFile << std::endl;

    double packet_limit =
        std::numeric_limits<double>::infinity(); // 要处理的数据包数量

    // KitNET参数:
    size_t maxAE = 10;      // 集成层中任何自动编码器的最大大小
    size_t FMgrace = 5000;  // 用于学习特征映射的实例数（集成的架构）
    size_t ADgrace = 50000; // 用于训练异常检测器的实例数（集成本身）

    // 异常检测参数
    double confidence_level = 0.99;  // 置信水平（99%）
    
    // 让用户设置置信水平
    std::cout << "\n统计概率方法参数设置:" << std::endl;
    std::cout << "请输入置信水平 (0.90-0.99, 默认" << confidence_level << "): ";
    std::string input;
    std::getline(std::cin, input);
    if (!input.empty()) {
        try {
            double val = std::stod(input);
            if (val >= 0.90 && val <= 0.99) {
                confidence_level = val;
            }
        } catch (...) {
            // 使用默认值
        }
    }
    
    std::cout << "使用检测方法: 统计概率方法 (对数正态分布)" << std::endl;
    std::cout << "置信水平: " << confidence_level << std::endl;

    // 构建Kitsune
    std::cout << "初始化Kitsune并解析" << extension << "文件..." << std::endl;
    Kitsune K(captureFile, packet_limit, maxAE, FMgrace, ADgrace);

    std::cout << "运行Kitsune进行特征提取和异常检测:" << std::endl;
    std::vector<double> RMSEs;
    int i = 0;

    // 计时开始
    auto start = std::chrono::high_resolution_clock::now();

    // 这里我们处理（训练/执行）每个单独的数据包。
    // 通过这种方式，每个观测值在执行process()方法后就会被丢弃。
    while (true) {
        i++;
        if (i % 1000 == 0) {
            std::cout << i << std::endl;
        }
        double rmse = K.proc_next_packet();
        if (rmse == -1) {
            break;
        }
        RMSEs.push_back(rmse);
    }

    // 计时结束
    auto stop = std::chrono::high_resolution_clock::now();
    auto duration =
        std::chrono::duration_cast<std::chrono::milliseconds>(stop - start);
    std::cout << "完成。耗时: " << duration.count() / 1000.0 << " 秒"
              << std::endl;

    // 指定起始检测点（训练完成后）
    size_t start_idx = FMgrace + ADgrace + 1;
    
    if (start_idx >= RMSEs.size()) {
        std::cerr << "警告: 训练期过长，没有足够的数据进行异常检测。" << std::endl;
        start_idx = RMSEs.size() > 100 ? RMSEs.size() - 100 : 0;
    }
    
    std::cout << "开始执行统计概率异常检测..." << std::endl;
    // 执行统计概率异常检测
    std::vector<size_t> anomalies = detectAnomaliesStatistical(RMSEs, start_idx, confidence_level);
    
    double anomaly_percent = 100.0 * anomalies.size() / (RMSEs.size() - start_idx);
    std::cout << "\n统计概率方法检测到 " << anomalies.size() 
              << " 个异常数据包，占测试数据的 "
              << anomaly_percent << "%"
              << std::endl;

    // 创建结果目录
    std::string resultsDir = "kitsune_results";
    std::filesystem::create_directory(resultsDir);

    // 将结果保存到文件中，以便后续分析或绘图
    std::string resultsFile = resultsDir + "/kitsune_results.csv";
    std::ofstream outFile(resultsFile);
    
    if (!outFile.is_open()) {
        std::cerr << "无法创建结果文件: " << resultsFile << std::endl;
        return 1;
    }
    
    // 设置较大的缓冲区以加速文件写入
    outFile.rdbuf()->pubsetbuf(nullptr, 0); // 禁用缓冲以提高性能
    
    // 写入标题行
    outFile << "index,rmse,is_anomaly" << std::endl;

    const size_t BUFFER_THRESHOLD = 1 << 20; // 例如：1 MB

    std::string buffer;
    buffer.reserve(BUFFER_THRESHOLD);  // 预留一定容量，减少动态扩容
    
    for (size_t idx = 0; idx < RMSEs.size(); idx++) {
        bool is_anomaly = std::find(anomalies.begin(), anomalies.end(), idx) != anomalies.end();

        // 构建CSV行
        std::stringstream ss;
        ss << idx << "," << RMSEs[idx] << "," << (is_anomaly ? 1 : 0) << std::endl;
        
        buffer.append(ss.str());

        if (buffer.size() >= BUFFER_THRESHOLD) {
            outFile.write(buffer.data(), buffer.size());
            outFile.flush();  // 可选: 确保数据写入磁盘
            buffer.clear();   // 清空 buffer，继续复用内存
        }
    }
    
    // 写入剩余部分
    if (!buffer.empty()) {
        outFile.write(buffer.data(), buffer.size());
        outFile.flush();
    }
    
    outFile.close();

    // 保存异常数据包索引
    std::string anomalyFile = resultsDir + "/anomaly_indices_statistical.csv";
    std::ofstream anomalyOutFile(anomalyFile);
    
    if (!anomalyOutFile.is_open()) {
        std::cerr << "无法创建异常索引文件: " << anomalyFile << std::endl;
        return 1;
    }
    
    anomalyOutFile << "index,rmse" << std::endl;
    
    // 使用缓冲区写入异常数据
    buffer.clear();
    for (size_t anomaly_idx : anomalies) {
        std::stringstream ss;
        ss << anomaly_idx << "," << RMSEs[anomaly_idx] << std::endl;
        buffer.append(ss.str());

        if (buffer.size() >= BUFFER_THRESHOLD) {
            anomalyOutFile.write(buffer.data(), buffer.size());
            anomalyOutFile.flush();  // 可选: 确保数据写入磁盘
            buffer.clear();   // 清空 buffer，继续复用内存
        }
    }
    

    // 写入剩余异常数据
    if (!buffer.empty()) {
        anomalyOutFile.write(buffer.data(), buffer.size());
        anomalyOutFile.flush();
    }
    
    anomalyOutFile.close();

    // 保存元数据
    std::string metadataFile = resultsDir + "/metadata.txt";
    std::ofstream metadataOutFile(metadataFile);
    metadataOutFile << "捕获文件: " << captureFile << std::endl;
    metadataOutFile << "文件格式: " << extension << std::endl;
    metadataOutFile << "处理的数据包数量: " << RMSEs.size() << std::endl;
    metadataOutFile << "特征映射宽限期: " << FMgrace << std::endl;
    metadataOutFile << "异常检测宽限期: " << ADgrace << std::endl;
    metadataOutFile << "最大自动编码器大小: " << maxAE << std::endl;
    metadataOutFile << "检测方法: 统计概率方法 (对数正态分布)" << std::endl;
    metadataOutFile << "置信水平: " << confidence_level << std::endl;
    metadataOutFile << "异常数量: " << anomalies.size() 
                    << " (" << anomaly_percent << "%)" << std::endl;
    metadataOutFile << "处理时间: " << duration.count() / 1000.0 << " 秒" << std::endl;
    metadataOutFile.close();

    std::cout << "结果已保存到 " << resultsDir << " 目录:" << std::endl;
    std::cout << "- " << resultsFile << " (所有结果)" << std::endl;
    std::cout << "- " << anomalyFile << " (统计概率方法的异常)" << std::endl;
    std::cout << "- " << metadataFile << " (元数据)" << std::endl;
    
    std::cout << "\n请使用Python或其他工具进行后续分析和绘图" << std::endl;
    std::cout << "可使用evaluate_detection程序评估检测效果" << std::endl;

    return 0;
}

// 解压缩文件的简单实现
// 注意：在实际应用中，你需要使用一个C++的ZIP库（如libzip、zlib等）
void extractZipFile(const std::string &zipFileName) {
// 在实际实现中，这里应该使用ZIP库
// 这里我们使用系统命令作为简单实现
#ifdef _WIN32
    std::string command = "powershell -command \"Expand-Archive -Path " +
                          zipFileName + " -DestinationPath .\"";
#else
    std::string command = "unzip -o " + zipFileName;
#endif

    int result = system(command.c_str());
    if (result != 0) {
        std::cerr << "解压失败，请确保安装了相应的解压工具" << std::endl;
    }
}