main.cpp

//
// Created by lyw on 1/13/25.
// 数据库流量解析系统 - 主程序入口
//

#include "constants.h"
#include "packet_parser.h"
#include <arpa/inet.h>
#include <iostream>
#include <stdexcept>
#include <cstdlib>

/**
 * @brief 从环境变量获取配置信息
 * @return 解析器配置
 */
ParserConfig GetConfigurationFromEnvironment() {
    ParserConfig config;
    
    // 读取数据库类型
    if (const char* db_type_env = std::getenv("DB_TYPE"); db_type_env) {
        std::string db_type_str = db_type_env;
        if (db_type_str == "hive") {
            config.db_type = DB_HIVE;
        } else if (db_type_str == "mysql") {
            config.db_type = DB_MYSQL;
        } else if (db_type_str == "pg") {
            config.db_type = DB_POSTGRESQL;
        } else {
            std::cerr << "Unknown database type: " << db_type_str << std::endl;
            std::cerr << "Supported types: mysql, pg, hive" << std::endl;
            exit(1);
        }
    } else {
        std::cerr << "DB_TYPE environment variable is required" << std::endl;
        exit(1);
    }
    
    // 读取PCAP文件路径
    if (const char* pcap_file_env = std::getenv("PCAP_FILE"); pcap_file_env) {
        config.pcap_file = pcap_file_env;
        std::cout << "PCAP file path: " << config.pcap_file << std::endl;
    } else {
        std::cerr << "PCAP_FILE environment variable is required" << std::endl;
        exit(1);
    }
    
    // 读取BPF过滤器
    if (const char* bpf_filter_env = std::getenv("BPF_FILTER"); bpf_filter_env) {
        config.bpf_filter = bpf_filter_env;
        std::cout << "BPF filter: " << config.bpf_filter << std::endl;
    }
    
    // 读取解析器地址
    if (const char* parser_addr_env = std::getenv("PARSER_ADDR"); parser_addr_env) {
        config.parser_addr = parser_addr_env;
    } else {
        std::cerr << "PARSER_ADDR environment variable is required" << std::endl;
        exit(1);
    }
    
    // 读取解析器端口
    if (const char* parser_port_env = std::getenv("PARSER_PORT"); parser_port_env) {
        config.parser_port = std::stoi(parser_port_env);
    } else {
        std::cerr << "PARSER_PORT environment variable is required" << std::endl;
        exit(1);
    }
    
    // 读取服务器地址
    if (const char* server_addr_env = std::getenv("SERVER_ADDR"); server_addr_env) {
        uint32_t server_ip = 0;
        if (inet_pton(AF_INET, server_addr_env, &server_ip) == 1) {
            config.server_addr = server_ip;
        } else {
            std::cerr << "Invalid SERVER_ADDR format: " << server_addr_env << std::endl;
            exit(1);
        }
    }
    
    // 读取服务器端口
    if (const char* server_port_env = std::getenv("SERVER_PORT"); server_port_env) {
        config.server_port = std::stoi(server_port_env);
    }
    
    return config;
}


void PrintUsage() {
    std::cout << "Database Traffic Parser" << std::endl;
    std::cout << "Required environment variables:" << std::endl;
    std::cout << "  DB_TYPE      - Database type (mysql, pg, hive)" << std::endl;
    std::cout << "  PCAP_FILE    - Path to PCAP file" << std::endl;
    std::cout << "  PARSER_ADDR  - Parser server IP address" << std::endl;
    std::cout << "  PARSER_PORT  - Parser server port" << std::endl;
    std::cout << "Optional environment variables:" << std::endl;
    std::cout << "  BPF_FILTER   - BPF filter expression" << std::endl;
    std::cout << "  SERVER_ADDR  - Database server IP address" << std::endl;
    std::cout << "  SERVER_PORT  - Database server port" << std::endl;
}

int main(int argc, char* argv[]) {
    if (argc > 1 && std::string(argv[1]) == "--help") {
        PrintUsage();
        return 0;
    }
    
    try {
        // 获取配置
        const auto config = GetConfigurationFromEnvironment();
        
        // 创建并运行解析器
        PacketParser parser(config);
        parser.ReadAndSendPackets();
        
        std::cout << "Packet parsing completed successfully" << std::endl;
        return 0;
        
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }
}