pcapforge_core/
feature_store.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc};
3use polars::prelude::*;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::path::Path;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct PacketFeatures {
10    pub timestamp: DateTime<Utc>,
11    pub src_ip: String,
12    pub dst_ip: String,
13    pub src_port: Option<u16>,
14    pub dst_port: Option<u16>,
15    pub protocol: String,
16    pub packet_size: u32,
17    pub flags: HashMap<String, bool>,
18    pub payload_size: u32,
19    pub ttl: Option<u8>,
20    pub window_size: Option<u16>,
21    pub custom_features: HashMap<String, serde_json::Value>,
22}
23
24pub struct FeatureStore {
25    features: Vec<PacketFeatures>,
26    dataframe: Option<DataFrame>,
27}
28
29impl FeatureStore {
30    pub fn new() -> Self {
31        Self {
32            features: Vec::new(),
33            dataframe: None,
34        }
35    }
36
37    pub fn add_features(&mut self, features: PacketFeatures) {
38        self.features.push(features);
39    }
40
41    pub fn build_dataframe(&mut self) -> Result<()> {
42        if self.features.is_empty() {
43            return Ok(());
44        }
45
46        let mut timestamps = Vec::new();
47        let mut src_ips = Vec::new();
48        let mut dst_ips = Vec::new();
49        let mut src_ports = Vec::new();
50        let mut dst_ports = Vec::new();
51        let mut protocols = Vec::new();
52        let mut packet_sizes = Vec::new();
53        let mut payload_sizes = Vec::new();
54
55        for feature in &self.features {
56            timestamps.push(feature.timestamp.timestamp_millis());
57            src_ips.push(feature.src_ip.clone());
58            dst_ips.push(feature.dst_ip.clone());
59            src_ports.push(feature.src_port);
60            dst_ports.push(feature.dst_port);
61            protocols.push(feature.protocol.clone());
62            packet_sizes.push(feature.packet_size);
63            payload_sizes.push(feature.payload_size);
64        }
65
66        let df = DataFrame::new(vec![
67            Column::new("timestamp".into(), timestamps),
68            Column::new("src_ip".into(), src_ips),
69            Column::new("dst_ip".into(), dst_ips),
70            Column::new("src_port".into(), src_ports.into_iter().map(|p| p.unwrap_or(0) as u32).collect::<Vec<_>>()),
71            Column::new("dst_port".into(), dst_ports.into_iter().map(|p| p.unwrap_or(0) as u32).collect::<Vec<_>>()),
72            Column::new("protocol".into(), protocols),
73            Column::new("packet_size".into(), packet_sizes),
74            Column::new("payload_size".into(), payload_sizes),
75        ])?;
76
77        self.dataframe = Some(df);
78        Ok(())
79    }
80
81    pub fn save_parquet<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
82        if self.dataframe.is_none() {
83            self.build_dataframe()?;
84        }
85
86        if let Some(df) = &mut self.dataframe {
87            let file = std::fs::File::create(path)?;
88            ParquetWriter::new(file).finish(df)?;
89        }
90
91        Ok(())
92    }
93
94    pub fn save_json<P: AsRef<Path>>(&self, path: P) -> Result<()> {
95        let file = std::fs::File::create(path)?;
96        serde_json::to_writer_pretty(file, &self.features)?;
97        Ok(())
98    }
99
100    pub fn save_csv<P: AsRef<Path>>(&mut self, path: P) -> Result<()> {
101        if self.dataframe.is_none() {
102            self.build_dataframe()?;
103        }
104
105        if let Some(df) = &mut self.dataframe {
106            let mut file = std::fs::File::create(path)?;
107            CsvWriter::new(&mut file).finish(df)?;
108        }
109
110        Ok(())
111    }
112
113    pub fn get_statistics(&self) -> FeatureStatistics {
114        let total_packets = self.features.len();
115        let mut protocol_counts: HashMap<String, usize> = HashMap::new();
116        let mut total_bytes = 0u64;
117
118        for feature in &self.features {
119            *protocol_counts.entry(feature.protocol.clone()).or_insert(0) += 1;
120            total_bytes += feature.packet_size as u64;
121        }
122
123        FeatureStatistics {
124            total_packets,
125            total_bytes,
126            protocol_distribution: protocol_counts,
127            unique_src_ips: self.features.iter()
128                .map(|f| f.src_ip.clone())
129                .collect::<std::collections::HashSet<_>>()
130                .len(),
131            unique_dst_ips: self.features.iter()
132                .map(|f| f.dst_ip.clone())
133                .collect::<std::collections::HashSet<_>>()
134                .len(),
135        }
136    }
137}
138
139#[derive(Debug, Serialize)]
140pub struct FeatureStatistics {
141    pub total_packets: usize,
142    pub total_bytes: u64,
143    pub protocol_distribution: HashMap<String, usize>,
144    pub unique_src_ips: usize,
145    pub unique_dst_ips: usize,
146}
pcapforge_core/feature_store.rs

pcapforge_core/
feature_store.rs
