FeatureExtractor.h

#pragma once

#include "netStat.h"
#include <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <vector>

/**
 * @brief 从pcap文件中提取Kitsune特征
 *
 * 从给定的pcap文件中一次提取一个数据包的特征。
 * 如果安装了Wireshark(tshark)，会使用它来解析(速度更快)
 * 否则，使用其他方式解析(速度较慢)
 * 如果使用了Wireshark，会生成一个tsv文件(pcap的解析版本)，可以在下次作为输入使用
 */
class FeatureExtractor {
  private:
    std::string path;               // 文件路径
    double limit;                   // 数据包数量限制
    std::string parse_type;         // 解析类型
    uint64_t curPacketIndx;         // 当前数据包索引
    std::unique_ptr<netStat> nstat; // 网络统计
    FILE *tsvinf;                   // TSV文件

    // 获取tshark路径
    std::string _get_tshark_path();

    // 使用tshark将pcap转换为tsv
    void pcap2tsv_with_tshark();

    // 准备文件
    void _prep();

  public:
    /**
     * @brief 构造函数
     * @param file_path 要解析的文件路径
     * @param limit 要处理的最大数据包数量，默认为无限
     */
    FeatureExtractor(const std::string &file_path,
                     double limit = std::numeric_limits<double>::infinity());

    /**
     * @brief 析构函数
     */
    ~FeatureExtractor();

    /**
     * @brief 获取下一个特征向量
     * @return 数据包的特征向量
     */
    std::vector<double> get_next_vector();

    /**
     * @brief 获取特征数量
     * @return 特征向量的大小
     */
    size_t get_num_features();
};