packet_handler.h

Go to the documentation of this file.

 
#pragma once
 
#include <vector>
#include <queue>
#include <unordered_map>
#include <functional>
#include <sstream>
#include "TcpLayer.h"
#include <UdpLayer.h>
#include <IcmpLayer.h>
#include "IPv4Layer.h"
#include "PayloadLayer.h"
#include "PacketUtils.h"
#include "SystemUtils.h"
#include <ctime>
#include <fstream>
#include <curl/curl.h>
#include<string>
#include <clickhouse/client.h>
#include <headers/unifiedPacketAttr.h>
#include <headers/HttpStatsCollector.h>

#define MAX_QUEUE_LENGTH 20000

#define WRITE_STAT_LINE(description, counter, measurement) \
        file<< description \
            << ','\
            << counter \
            << ','\
            <<measurement<< std::endl;
 
using namespace clickhouse;

void StatsSummary2File(HttpStatsCollector& collector,std::string filename)
{
    std::ofstream file(filename);
    if (!file.is_open()) {
        std::cerr << "Failed to open file: " << filename << std::endl;
        exit ;
    }
    WRITE_STAT_LINE("Sample time", collector.getGeneralStats().sampleTime, "Seconds");
    WRITE_STAT_LINE("Number of HTTP packets", collector.getGeneralStats().numOfHttpPackets, "Packets");
    WRITE_STAT_LINE("Rate of HTTP packets", collector.getGeneralStats().httpPacketRate.totalRate, "Packets/sec");
    WRITE_STAT_LINE("Number of HTTP flows", collector.getGeneralStats().numOfHttpFlows, "Flows");
    WRITE_STAT_LINE("Rate of HTTP flows", collector.getGeneralStats().httpFlowRate.totalRate, "Flows/sec");
    WRITE_STAT_LINE("Number of HTTP pipelining flows", collector.getGeneralStats().numOfHttpPipeliningFlows, "Flows");
    WRITE_STAT_LINE("Number of HTTP transactions", collector.getGeneralStats().numOfHttpTransactions, "Transactions");
    WRITE_STAT_LINE("Rate of HTTP transactions", collector.getGeneralStats().httpTransactionsRate.totalRate, "Transactions/sec");
    WRITE_STAT_LINE("Total HTTP data", collector.getGeneralStats().amountOfHttpTraffic, "Bytes");
    WRITE_STAT_LINE("Rate of HTTP data", collector.getGeneralStats().httpTrafficRate.totalRate, "Bytes/sec");
    WRITE_STAT_LINE("Average packets per flow", collector.getGeneralStats().averageNumOfPacketsPerFlow, "Packets");
    WRITE_STAT_LINE("Average transactions per flow", collector.getGeneralStats().averageNumOfHttpTransactionsPerFlow, "Transactions");
    WRITE_STAT_LINE("Average data per flow", collector.getGeneralStats().averageAmountOfDataPerFlow, "Bytes");
 
    WRITE_STAT_LINE("Number of HTTP requests", collector.getRequestStats().numOfMessages, "Requests");
    WRITE_STAT_LINE("Rate of HTTP requests", collector.getRequestStats().messageRate.totalRate, "Requests/sec");
    WRITE_STAT_LINE("Total data in headers", collector.getRequestStats().totalMessageHeaderSize, "Bytes");
    WRITE_STAT_LINE("Average header size", collector.getRequestStats().averageMessageHeaderSize, "Bytes");
 
    WRITE_STAT_LINE("Number of HTTP responses", collector.getResponseStats().numOfMessages, "Responses");
    WRITE_STAT_LINE("Rate of HTTP responses", collector.getResponseStats().messageRate.totalRate, "Responses/sec");
    WRITE_STAT_LINE("Total data in headers", collector.getResponseStats().totalMessageHeaderSize, "Bytes");
    WRITE_STAT_LINE("Average header size", collector.getResponseStats().averageMessageHeaderSize, "Bytes");
    WRITE_STAT_LINE("Num of responses with content-length", collector.getResponseStats().numOfMessagesWithContentLength, "Responses");
    WRITE_STAT_LINE("Total body size (may be compressed)", collector.getResponseStats().totalContentLengthSize, "Bytes");
    WRITE_STAT_LINE("Average body size", collector.getResponseStats().averageContentLengthSize, "Bytes");
    file.close();
}

class packet_handler
{
private:
    int queue_len;                                          
    std::string table_prefix;                               
    Client client;                                          
    
    // ClickHouse数据列定义
    std::shared_ptr<ColumnUInt32> col_hash_val;             
    std::shared_ptr<ColumnInt32> col_packet_len;            
    std::shared_ptr<ColumnUInt32> col_tv_sec;               
    std::shared_ptr<ColumnUInt32> col_tv_nsec;              
    std::shared_ptr<ColumnString> col_src_ip;               
    std::shared_ptr<ColumnString> col_dst_ip;               
    std::shared_ptr<ColumnUInt8> col_ttl;                   
    std::shared_ptr<ColumnUInt8> col_tos;                   
    std::shared_ptr<ColumnUInt16> col_id;                   
    std::shared_ptr<ColumnUInt16> col_offset;               
    std::shared_ptr<ColumnUInt8> col_protocol_ip;           
    std::shared_ptr<ColumnUInt8> col_protocol;              
    std::shared_ptr<ColumnUInt16> col_len_load;             
    std::shared_ptr<ColumnUInt16> col_src_port;             
    std::shared_ptr<ColumnUInt16> col_dst_port;             
    std::shared_ptr<ColumnUInt32> col_ack_num;              
    std::shared_ptr<ColumnUInt32> col_seq_num;              
    std::shared_ptr<ColumnUInt16> col_flag;                 
    std::shared_ptr<ColumnUInt16> col_window;               
    std::shared_ptr<ColumnUInt16> col_len_udp;              
    std::shared_ptr<ColumnUInt8> col_icmp_type;             
    std::shared_ptr<ColumnUInt8> col_icmp_code;             
    HttpStatsCollector collector;                           
 
public:
    packet_handler(std::string tb_prefix = ""):table_prefix(tb_prefix),client((ClientOptions().SetHost("localhost")))
    {
        col_hash_val    = std::make_shared<ColumnUInt32>();
        col_packet_len  =std::make_shared<ColumnInt32>();
        col_tv_sec      = std::make_shared<ColumnUInt32>();
        col_tv_nsec     =std::make_shared<ColumnUInt32>();
        col_src_ip      =std::make_shared<ColumnString>();
        col_dst_ip      =std::make_shared<ColumnString>();
        col_ttl         =std::make_shared<ColumnUInt8>();
        col_tos         =std::make_shared<ColumnUInt8>();
        col_id          =std::make_shared<ColumnUInt16>();
        col_offset      =std::make_shared<ColumnUInt16>();
        col_protocol_ip =std::make_shared<ColumnUInt8>();
        col_protocol     =std::make_shared<ColumnUInt8>();
        col_len_load     =std::make_shared<ColumnUInt16>();
        col_src_port     =std::make_shared<ColumnUInt16>();
        col_dst_port     =std::make_shared<ColumnUInt16>();
        col_ack_num     =std::make_shared<ColumnUInt32>();
        col_seq_num     =std::make_shared<ColumnUInt32>();
        col_flag       =std::make_shared<ColumnUInt16>();
        col_window     =std::make_shared<ColumnUInt16>();
        col_len_udp    =std::make_shared<ColumnUInt16>();
        col_icmp_type  =std::make_shared<ColumnUInt8>();
        col_icmp_code  =std::make_shared<ColumnUInt8>();
        HttpStatsCollector collector();
    };

    ~packet_handler() {
        if(queue_len!=0) send2db();  // 线程池销毁时执行
    }

    void handle_packet(pcpp::Packet *packet)
    {
        queue_len++;
 
        unifiedPacketAttr packet_attr(packet);
        col_hash_val->Append(packet_attr.hash_val);
        col_packet_len->Append(packet_attr.packet_len);
        col_tv_sec->Append(packet_attr.tv_sec);
        col_tv_nsec->Append(packet_attr.tv_nsec);
        col_src_ip->Append(packet_attr.src_ip);
        col_dst_ip->Append(packet_attr.dst_ip);
        col_ttl->Append(packet_attr.ttl);
        col_tos->Append(packet_attr.tos);
        col_id->Append(packet_attr.id);
        col_offset->Append(packet_attr.offset);
        col_protocol_ip->Append(packet_attr.protocol_ip);
        col_protocol->Append(packet_attr.protocol);
        col_len_load->Append(packet_attr.len_load);
        col_src_port->Append(packet_attr.src_port);
        col_dst_port->Append(packet_attr.dst_port);
        col_ack_num->Append(packet_attr.ack_num);
        col_seq_num->Append(packet_attr.seq_num);
        col_flag->Append(packet_attr.flag);
        col_window->Append(packet_attr.window);
        col_len_udp->Append(packet_attr.len_udp);
        col_icmp_type->Append(packet_attr.icmp_type);
        col_icmp_code->Append(packet_attr.icmp_code);
        collector.collectStats(packet);
        
        if (queue_len >= MAX_QUEUE_LENGTH) // 如果达到一定数量，就批量存入db
        {
            send2db();
            flush();
            queue_len = 0;
            StatsSummary2File(collector,"/home/njust/program/example-app/stats.txt");
            collector.clear();
        }
    }

    void flush()
    {
        col_hash_val   ->Clear();
        col_packet_len ->Clear();
        col_tv_sec     ->Clear();
        col_tv_nsec    ->Clear();
        col_src_ip     ->Clear();
        col_dst_ip     ->Clear();
        col_ttl        ->Clear();
        col_tos        ->Clear();
        col_id         ->Clear();
        col_offset     ->Clear();
        col_protocol_ip->Clear();
        col_protocol   ->Clear();
        col_len_load   ->Clear();
        col_src_port   ->Clear();
        col_dst_port   ->Clear();
        col_ack_num    ->Clear();
        col_seq_num    ->Clear();
        col_flag       ->Clear();
        col_window     ->Clear();
        col_len_udp    ->Clear();
        col_icmp_type  ->Clear();
        col_icmp_code  ->Clear();
    };

    void send2db()
    { 
        // 构建主数据块
        Block block;
        block.AppendColumn("hashVal",       col_hash_val   );
        block.AppendColumn("packetLen",     col_packet_len );
        block.AppendColumn("tvSec",         col_tv_sec     );
        block.AppendColumn("tvNsec",        col_tv_nsec    );
        block.AppendColumn("srcIp",         col_src_ip     );
        block.AppendColumn("dstIp",         col_dst_ip     );
        block.AppendColumn("ttl",           col_ttl        );
        block.AppendColumn("tos",           col_tos        );
        block.AppendColumn("id",            col_id         );
        block.AppendColumn("offset",        col_offset     );
        block.AppendColumn("protocolIP",    col_protocol_ip);
        block.AppendColumn("protocol",      col_protocol   );
        block.AppendColumn("lenLoad",       col_len_load   );
        block.AppendColumn("srcPort",       col_src_port   );
        block.AppendColumn("dstPort",       col_dst_port   );
        block.AppendColumn("ackNum",        col_ack_num    );
        block.AppendColumn("seqNum",        col_seq_num    );
        block.AppendColumn("flag",          col_flag       );
        block.AppendColumn("window",        col_window     );
        block.AppendColumn("lenUdp",        col_len_udp    );
        block.AppendColumn("icmpType",      col_icmp_type  );
        block.AppendColumn("icmpCode",      col_icmp_code  );
 
        // 构建源IP表数据块
        Block block_src_ip;
        block_src_ip.AppendColumn("srcIp",col_src_ip);
        block_src_ip.AppendColumn("hashVal",col_hash_val);
 
        // 构建协议表数据块
        Block block_protocol;
        block_protocol.AppendColumn("protocol",col_protocol);
        block_protocol.AppendColumn("hashVal",col_hash_val);
 
        // 插入主包表
        try
        {
            client.Insert("default."+table_prefix+"_packets", block);
        }
        catch(const std::exception& e)
        {
            std::cerr <<"Insert failue.retrying... error msg:" <<e.what() << '\n';
            try
            {
                client.ResetConnection();
                client.Insert("default."+table_prefix+"_packets", block);
            }
            catch(const std::exception& e) {
                std::cerr<<"insert block failed!"<<std::endl;
            }
        }
 
        // 插入源IP哈希表
        try
        {
            client.Insert("_src_ip_hash_table", block_src_ip);
        }
        catch(const std::exception& e)
        {
            std::cerr <<"Insert failue.retrying... error msg:" <<e.what() << '\n';
            try
            {
                client.ResetConnection();
                client.Insert("_src_ip_hash_table", block_src_ip);
            }
            catch(const std::exception& e)
            {
                std::cerr <<"Insert failue and drop .error msg:" <<e.what() << '\n';
            }
        }
 
        // 插入协议哈希表
        try
        {
            client.Insert("_protocol_hash_table", block_protocol);
        }
        catch(const std::exception& e)
        {
            std::cerr <<"Insert failue.retrying... error msg:" <<e.what() << '\n';
            try
            {
                client.ResetConnection();
                client.Insert("_protocol_hash_table", block_protocol);
            }
            catch(const std::exception& e)
            {
                std::cerr <<"Insert failue and drop .error msg:" <<e.what() << '\n';
            }
        }
    };
};