net_trace/
sessions.rs

1use std::collections::HashMap;
2use serde::Serialize;
3use chrono::DateTime;
4use chrono::Utc;
5use crate::NetworkPacket;
6
7/// Represents a TCP session
8#[derive(Debug, Serialize)]
9pub struct TCPSession {
10    source_port: u16,
11    destination_port: u16,
12    source_ip: [u8; 4],
13    destination_ip: [u8; 4],
14    start_timestamp: DateTime<Utc>,
15    end_timestamp: Option<DateTime<Utc>>,
16    packets: Vec<NetworkPacket>,
17}
18
19pub fn find_tcp_sessions(packets: &[NetworkPacket]) -> Vec<TCPSession> {
20    let mut sessions: HashMap<(u16, u16, [u8; 4], [u8; 4]), (DateTime<Utc>, Option<DateTime<Utc>>, Vec<NetworkPacket>)> = HashMap::new();
21
22    for packet in packets {
23        if packet.tcp_layer.flags.syn {
24            // Start of a new session
25            let key = (
26                packet.tcp_layer.source_port,
27                packet.tcp_layer.destination_port,
28                packet.ip_layer.source_ip,
29                packet.ip_layer.destination_ip,
30            );
31            sessions
32                .entry(key)
33                .or_insert((packet.timestamp, None, vec![]))
34                .2
35                .push(packet.clone());
36        }
37
38        if packet.tcp_layer.flags.fin {
39            // End of an existing session
40            let key = (
41                packet.tcp_layer.destination_port,
42                packet.tcp_layer.source_port,
43                packet.ip_layer.destination_ip,
44                packet.ip_layer.source_ip,
45            );
46
47            if let Some(session) = sessions.get_mut(&key) {
48                session.1 = Some(packet.timestamp);
49                session.2.push(packet.clone());
50            }
51        }
52
53        // Track packets for ongoing sessions
54        for session in sessions.values_mut() {
55            session.2.push(packet.clone());
56        }
57    }
58
59    // Convert to vector format
60    sessions
61        .into_iter()
62        .map(|(key, (start, end, packets))| TCPSession {
63            source_port: key.0,
64            destination_port: key.1,
65            source_ip: key.2,
66            destination_ip: key.3,
67            start_timestamp: start,
68            end_timestamp: end,
69            packets,
70        })
71        .collect()
72}