wiretap/
tcp_packet.rs

1use pnet::packet::tcp::TcpPacket as pnet_TcpPacket;
2use pnet::packet::Packet;
3use std::ops::{Deref, DerefMut};
4use std::sync::Arc;
5
6/// Wrapper around pnet's TcpPacket for adding additional funcitonality
7#[derive(Debug)]
8pub struct TcpSegment<'a>(pnet_TcpPacket<'a>);
9
10impl<'a> From<pnet_TcpPacket<'a>> for TcpSegment<'a> {
11    fn from(ipv4_packet: pnet_TcpPacket<'a>) -> Self {
12        TcpSegment(ipv4_packet)
13    }
14}
15
16impl<'a> Deref for TcpSegment<'a> {
17    type Target = pnet_TcpPacket<'a>;
18
19    fn deref(&self) -> &Self::Target {
20        &self.0
21    }
22}
23
24impl TcpSegment<'_> {
25    pub fn new<'a>(packet: &'a [u8]) -> Option<TcpSegment<'a>>{
26        pnet_TcpPacket::new(packet).map(TcpSegment::from)
27    }
28
29    /// Return true if the TCP segment has a payload
30    pub fn has_payload(&self) -> bool {
31        !&self.payload().is_empty()
32    }
33
34    pub fn create_clone<'a>(&self) -> TcpSegment<'a> {
35        TcpSegment::from(pnet_TcpPacket::owned(self.packet().to_vec()).unwrap())
36    }
37
38    fn is_answered_by(&self, other: &TcpSegment<'_>) -> bool {
39        self.get_source() == other.get_destination()
40            && self.get_destination() == other.get_source()
41            && self.get_sequence() as usize + self.payload().len()
42                == other.get_acknowledgement() as usize
43    }
44}
45
46/// Wrapper around an Arc<[TcpSegment]> for additional functionality
47#[derive(Debug)]
48pub struct TcpSegmentCollection<'a>(Arc<[TcpSegment<'a>]>);
49
50impl<'a> FromIterator<TcpSegment<'a>> for TcpSegmentCollection<'a> {
51    fn from_iter<I: IntoIterator<Item = TcpSegment<'a>>>(iter: I) -> Self {
52        TcpSegmentCollection(iter.into_iter().collect())
53    }
54}
55
56impl<'a> Deref for TcpSegmentCollection<'a> {
57    type Target = Arc<[TcpSegment<'a>]>;
58
59    fn deref(&self) -> &Self::Target {
60        &self.0
61    }
62}
63
64impl<'a> From<crate::Ipv4PacketCollection<'a>> for TcpSegmentCollection<'a> {
65    fn from(ipv4_packet_collection: crate::Ipv4PacketCollection) -> Self {
66        ipv4_packet_collection
67            .iter()
68            .filter(|ipv4_packet| pnet_TcpPacket::new(ipv4_packet.payload()).is_some())
69            .map(|ipv4_packet| {
70                TcpSegment::from(pnet_TcpPacket::owned(ipv4_packet.payload().to_vec()).unwrap())
71            })
72            .collect::<TcpSegmentCollection>()
73    }
74}
75
76impl<'a> TcpSegmentCollection<'a> {
77    /// Get a collection of TcpSegment with TCP payloads
78    ///
79    /// Returns a new TcpSegmentCollection containing only the segments that have a TCP payload
80    pub fn filter_no_payload(&'a self) -> TcpSegmentCollection<'a> {
81        TcpSegmentCollection(
82            self.iter()
83                .filter(|s| s.has_payload())
84                .map(|s| s.create_clone())
85                .collect::<Arc<[TcpSegment]>>(),
86        )
87    }
88
89    /// Couple the challenge / response pairs in a collection of TCP segments
90    ///
91    /// Returns a new TcpSegmentCollection containing only the segments that have a TCP payload
92    pub fn find_challenge_response_pairs(
93        &'a mut self,
94    ) -> (TcpChallengeResponseCollection<'a>, TcpSegmentCollection<'a>) {
95        let mut matched = Vec::new();
96        let mut unmatched = self
97            .iter()
98            .map(|s| s.create_clone())
99            .collect::<Vec<TcpSegment<'a>>>();
100        let mut i = 0;
101        while i < unmatched.len() {
102            let challenge = unmatched[i].create_clone();
103            let mut j = 0;
104            let mut found_match = false;
105            while j < unmatched.len() - 1 {
106                j += 1;
107                let candidate = unmatched[j].create_clone();
108                if challenge.is_answered_by(&candidate) {
109                    matched.push(TcpChallengeResponse::new(
110                        challenge.create_clone(),
111                        candidate.create_clone(),
112                    ));
113                    if j > i {
114                        unmatched.remove(j);
115                        unmatched.remove(i);
116                    } else {
117                        unmatched.remove(i);
118                        unmatched.remove(j);
119                    }
120                    found_match = true;
121                    break;
122                }
123            }
124            if !found_match {
125                i += 1;
126            }
127        }
128        (
129            TcpChallengeResponseCollection(matched.into()),
130            TcpSegmentCollection(unmatched.into()),
131        )
132    }
133}
134
135/// Container for TCP segments where the "challenge" was answered by the "response"
136#[derive(Debug)]
137pub struct TcpChallengeResponse<'a> {
138    pub challenge: TcpSegment<'a>,
139    pub response: TcpSegment<'a>,
140}
141
142impl<'a> TcpChallengeResponse<'a> {
143    fn new(challenge: TcpSegment<'a>, response: TcpSegment<'a>) -> TcpChallengeResponse<'a> {
144        TcpChallengeResponse {
145            challenge,
146            response,
147        }
148    }
149}
150
151/// Wrapper around an Arc<[TcpChallengeResponse]> for additional functionality
152#[derive(Debug)]
153pub struct TcpChallengeResponseCollection<'a>(Arc<[TcpChallengeResponse<'a>]>);
154
155impl<'a> FromIterator<TcpChallengeResponse<'a>> for TcpChallengeResponseCollection<'a> {
156    fn from_iter<I: IntoIterator<Item = TcpChallengeResponse<'a>>>(iter: I) -> Self {
157        TcpChallengeResponseCollection(iter.into_iter().collect())
158    }
159}
160
161impl<'a> Deref for TcpChallengeResponseCollection<'a> {
162    type Target = Arc<[TcpChallengeResponse<'a>]>;
163
164    fn deref(&self) -> &Self::Target {
165        &self.0
166    }
167}
168
169impl DerefMut for TcpChallengeResponseCollection<'_> {
170    fn deref_mut(&mut self) -> &mut Self::Target {
171        &mut self.0
172    }
173}
wiretap/tcp_packet.rs

wiretap/
tcp_packet.rs
