Skip to main content

stackforge_core/pcap/
writer.rs

1//! PCAP and PcapNG file writer.
2
3use std::borrow::Cow;
4use std::fs::File;
5use std::io::{BufWriter, Write};
6use std::path::Path;
7use std::time::Duration;
8
9use pcap_file::pcap::{PcapHeader, PcapPacket, PcapWriter as PcapFileWriter};
10use pcap_file::pcapng::PcapNgWriter as PcapNgFileWriter;
11use pcap_file::pcapng::blocks::enhanced_packet::EnhancedPacketBlock;
12use pcap_file::pcapng::blocks::interface_description::InterfaceDescriptionBlock;
13
14use crate::error::{PacketError, Result};
15use crate::packet::Packet;
16
17use super::{CapturedPacket, PcapMetadata};
18
19// ---------------------------------------------------------------------------
20// Classic PCAP writer
21// ---------------------------------------------------------------------------
22
23/// Write captured packets to a PCAP file.
24///
25/// Preserves timestamps and original length from [`CapturedPacket`] metadata.
26pub fn wrpcap(path: impl AsRef<Path>, packets: &[CapturedPacket]) -> Result<()> {
27    let file = File::create(path.as_ref()).map_err(|e| {
28        PacketError::Io(format!(
29            "failed to create {}: {}",
30            path.as_ref().display(),
31            e
32        ))
33    })?;
34    let writer = BufWriter::new(file);
35
36    let header = PcapHeader::default();
37    let mut pcap_writer = PcapFileWriter::with_header(writer, header)
38        .map_err(|e| PacketError::Io(format!("PCAP write error: {e}")))?;
39
40    for cap in packets {
41        let pcap_pkt = PcapPacket::new(
42            cap.metadata.timestamp,
43            cap.metadata.orig_len,
44            cap.packet.as_bytes(),
45        );
46        pcap_writer
47            .write_packet(&pcap_pkt)
48            .map_err(|e| PacketError::Io(format!("PCAP write error: {e}")))?;
49    }
50
51    Ok(())
52}
53
54/// Write plain packets to a PCAP file (convenience function).
55///
56/// Timestamps are set to zero, `orig_len` matches each packet's data length.
57pub fn wrpcap_packets(path: impl AsRef<Path>, packets: &[Packet]) -> Result<()> {
58    let captured: Vec<CapturedPacket> = packets
59        .iter()
60        .map(|pkt| CapturedPacket {
61            packet: pkt.clone(),
62            metadata: PcapMetadata {
63                timestamp: Duration::ZERO,
64                orig_len: pkt.len() as u32,
65                ..Default::default()
66            },
67        })
68        .collect();
69    wrpcap(path, &captured)
70}
71
72/// PCAP writer for streaming writes.
73///
74/// Writes packets one at a time without buffering them all in memory.
75pub struct PcapStreamWriter<W: Write> {
76    inner: PcapFileWriter<W>,
77}
78
79impl PcapStreamWriter<BufWriter<File>> {
80    /// Create a new PCAP file for writing.
81    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
82        let file = File::create(path.as_ref()).map_err(|e| {
83            PacketError::Io(format!(
84                "failed to create {}: {}",
85                path.as_ref().display(),
86                e
87            ))
88        })?;
89        let writer = BufWriter::new(file);
90        Self::from_writer(writer)
91    }
92}
93
94impl<W: Write> PcapStreamWriter<W> {
95    /// Create a `PcapStreamWriter` from any writer.
96    pub fn from_writer(writer: W) -> Result<Self> {
97        let header = PcapHeader::default();
98        let pcap_writer = PcapFileWriter::with_header(writer, header)
99            .map_err(|e| PacketError::Io(format!("PCAP write error: {e}")))?;
100        Ok(Self { inner: pcap_writer })
101    }
102
103    /// Write a captured packet with metadata.
104    pub fn write(&mut self, cap: &CapturedPacket) -> Result<()> {
105        let pcap_pkt = PcapPacket::new(
106            cap.metadata.timestamp,
107            cap.metadata.orig_len,
108            cap.packet.as_bytes(),
109        );
110        self.inner
111            .write_packet(&pcap_pkt)
112            .map_err(|e| PacketError::Io(format!("PCAP write error: {e}")))?;
113        Ok(())
114    }
115
116    /// Write a plain packet (timestamp=0, `orig_len=data` length).
117    pub fn write_packet(&mut self, pkt: &Packet) -> Result<()> {
118        let pcap_pkt = PcapPacket::new(Duration::ZERO, pkt.len() as u32, pkt.as_bytes());
119        self.inner
120            .write_packet(&pcap_pkt)
121            .map_err(|e| PacketError::Io(format!("PCAP write error: {e}")))?;
122        Ok(())
123    }
124}
125
126// ---------------------------------------------------------------------------
127// PcapNG writer
128// ---------------------------------------------------------------------------
129
130/// Write captured packets to a PcapNG file.
131///
132/// Writes SHB + IDB (Ethernet) header, then each packet as an Enhanced Packet Block.
133pub fn wrpcapng(path: impl AsRef<Path>, packets: &[CapturedPacket]) -> Result<()> {
134    let file = File::create(path.as_ref()).map_err(|e| {
135        PacketError::Io(format!(
136            "failed to create {}: {}",
137            path.as_ref().display(),
138            e
139        ))
140    })?;
141    let writer = BufWriter::new(file);
142    let mut ng_writer = PcapNgStreamWriter::from_writer(writer)?;
143
144    for cap in packets {
145        ng_writer.write(cap)?;
146    }
147
148    Ok(())
149}
150
151/// Write plain packets to a PcapNG file (convenience function).
152///
153/// Timestamps are set to zero, `orig_len` matches each packet's data length.
154pub fn wrpcapng_packets(path: impl AsRef<Path>, packets: &[Packet]) -> Result<()> {
155    let captured: Vec<CapturedPacket> = packets
156        .iter()
157        .map(|pkt| CapturedPacket {
158            packet: pkt.clone(),
159            metadata: PcapMetadata {
160                timestamp: Duration::ZERO,
161                orig_len: pkt.len() as u32,
162                ..Default::default()
163            },
164        })
165        .collect();
166    wrpcapng(path, &captured)
167}
168
169/// PcapNG writer for streaming writes.
170///
171/// Writes packets one at a time. Auto-writes SHB + IDB on first packet.
172pub struct PcapNgStreamWriter<W: Write> {
173    inner: PcapNgFileWriter<W>,
174    interface_written: bool,
175}
176
177impl PcapNgStreamWriter<BufWriter<File>> {
178    /// Create a new PcapNG file for writing.
179    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
180        let file = File::create(path.as_ref()).map_err(|e| {
181            PacketError::Io(format!(
182                "failed to create {}: {}",
183                path.as_ref().display(),
184                e
185            ))
186        })?;
187        let writer = BufWriter::new(file);
188        Self::from_writer(writer)
189    }
190}
191
192impl<W: Write> PcapNgStreamWriter<W> {
193    /// Create a `PcapNgStreamWriter` from any writer.
194    ///
195    /// The SHB is written immediately. The IDB is written on the first packet.
196    pub fn from_writer(writer: W) -> Result<Self> {
197        let ng_writer = PcapNgFileWriter::new(writer)
198            .map_err(|e| PacketError::Io(format!("PcapNG write error: {e}")))?;
199        Ok(Self {
200            inner: ng_writer,
201            interface_written: false,
202        })
203    }
204
205    /// Ensure at least one Interface Description Block has been written.
206    fn ensure_interface(&mut self) -> Result<()> {
207        if !self.interface_written {
208            let idb = InterfaceDescriptionBlock {
209                linktype: pcap_file::DataLink::ETHERNET,
210                snaplen: 0xFFFF,
211                options: vec![],
212            };
213            self.inner
214                .write_pcapng_block(idb)
215                .map_err(|e| PacketError::Io(format!("PcapNG write error: {e}")))?;
216            self.interface_written = true;
217        }
218        Ok(())
219    }
220
221    /// Write a captured packet with metadata as an Enhanced Packet Block.
222    pub fn write(&mut self, cap: &CapturedPacket) -> Result<()> {
223        self.ensure_interface()?;
224        let epb = EnhancedPacketBlock {
225            interface_id: cap.metadata.interface_id.unwrap_or(0),
226            timestamp: cap.metadata.timestamp,
227            original_len: cap.metadata.orig_len,
228            data: Cow::Borrowed(cap.packet.as_bytes()),
229            options: vec![],
230        };
231        self.inner
232            .write_pcapng_block(epb)
233            .map_err(|e| PacketError::Io(format!("PcapNG write error: {e}")))?;
234        Ok(())
235    }
236
237    /// Write a plain packet (timestamp=0, `orig_len=data` length).
238    pub fn write_packet(&mut self, pkt: &Packet) -> Result<()> {
239        self.ensure_interface()?;
240        let epb = EnhancedPacketBlock {
241            interface_id: 0,
242            timestamp: Duration::ZERO,
243            original_len: pkt.len() as u32,
244            data: Cow::Borrowed(pkt.as_bytes()),
245            options: vec![],
246        };
247        self.inner
248            .write_pcapng_block(epb)
249            .map_err(|e| PacketError::Io(format!("PcapNG write error: {e}")))?;
250        Ok(())
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use crate::pcap::reader::{PcapIterator, PcapNgIterator};
258    use std::io::Cursor;
259    use std::time::Duration;
260
261    fn sample_ethernet_packet() -> Vec<u8> {
262        vec![
263            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // dst: broadcast
264            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, // src
265            0x08, 0x00, // EtherType: IPv4
266            0x00, 0x00, 0x00, 0x00, // dummy payload
267        ]
268    }
269
270    #[test]
271    fn test_wrpcap_roundtrip() {
272        let eth = sample_ethernet_packet();
273        let pkt = Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth));
274        let cap = CapturedPacket {
275            packet: pkt,
276            metadata: PcapMetadata {
277                timestamp: Duration::from_secs(42),
278                orig_len: eth.len() as u32,
279                ..Default::default()
280            },
281        };
282
283        let mut buf = Vec::new();
284        {
285            let mut writer = PcapStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
286            writer.write(&cap).unwrap();
287        }
288
289        let iter = PcapIterator::from_reader(Cursor::new(buf)).unwrap();
290        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
291        assert_eq!(packets.len(), 1);
292        assert_eq!(packets[0].metadata.timestamp, Duration::from_secs(42));
293        assert_eq!(packets[0].packet.as_bytes(), eth.as_slice());
294    }
295
296    #[test]
297    fn test_wrpcap_multiple_packets() {
298        let eth = sample_ethernet_packet();
299
300        let caps: Vec<CapturedPacket> = (0..5)
301            .map(|i| CapturedPacket {
302                packet: Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth)),
303                metadata: PcapMetadata {
304                    timestamp: Duration::from_secs(i),
305                    orig_len: eth.len() as u32,
306                    ..Default::default()
307                },
308            })
309            .collect();
310
311        let mut buf = Vec::new();
312        {
313            let mut writer = PcapStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
314            for cap in &caps {
315                writer.write(cap).unwrap();
316            }
317        }
318
319        let iter = PcapIterator::from_reader(Cursor::new(buf)).unwrap();
320        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
321        assert_eq!(packets.len(), 5);
322        for (i, pkt) in packets.iter().enumerate() {
323            assert_eq!(pkt.metadata.timestamp, Duration::from_secs(i as u64));
324        }
325    }
326
327    #[test]
328    fn test_write_packet_convenience() {
329        let eth = sample_ethernet_packet();
330        let pkt = Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth));
331
332        let mut buf = Vec::new();
333        {
334            let mut writer = PcapStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
335            writer.write_packet(&pkt).unwrap();
336        }
337
338        let iter = PcapIterator::from_reader(Cursor::new(buf)).unwrap();
339        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
340        assert_eq!(packets.len(), 1);
341        assert_eq!(packets[0].metadata.timestamp, Duration::ZERO);
342        assert_eq!(packets[0].metadata.orig_len, eth.len() as u32);
343    }
344
345    #[test]
346    fn test_pcapng_writer_roundtrip() {
347        let eth = sample_ethernet_packet();
348        let cap = CapturedPacket {
349            packet: Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth)),
350            metadata: PcapMetadata {
351                timestamp: Duration::from_secs(100),
352                orig_len: eth.len() as u32,
353                interface_id: Some(0),
354                comment: None,
355            },
356        };
357
358        let mut buf = Vec::new();
359        {
360            let mut writer = PcapNgStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
361            writer.write(&cap).unwrap();
362        }
363
364        let iter = PcapNgIterator::from_reader(Cursor::new(buf)).unwrap();
365        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
366        assert_eq!(packets.len(), 1);
367        assert_eq!(packets[0].metadata.timestamp, Duration::from_secs(100));
368        assert_eq!(packets[0].packet.as_bytes(), eth.as_slice());
369        assert_eq!(packets[0].metadata.interface_id, Some(0));
370    }
371
372    #[test]
373    fn test_pcapng_writer_multiple_packets() {
374        let eth = sample_ethernet_packet();
375
376        let caps: Vec<CapturedPacket> = (0..3)
377            .map(|i| CapturedPacket {
378                packet: Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth)),
379                metadata: PcapMetadata {
380                    timestamp: Duration::from_secs(i * 10),
381                    orig_len: eth.len() as u32,
382                    ..Default::default()
383                },
384            })
385            .collect();
386
387        let mut buf = Vec::new();
388        {
389            let mut writer = PcapNgStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
390            for cap in &caps {
391                writer.write(cap).unwrap();
392            }
393        }
394
395        let iter = PcapNgIterator::from_reader(Cursor::new(buf)).unwrap();
396        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
397        assert_eq!(packets.len(), 3);
398        assert_eq!(packets[0].metadata.timestamp, Duration::from_secs(0));
399        assert_eq!(packets[1].metadata.timestamp, Duration::from_secs(10));
400        assert_eq!(packets[2].metadata.timestamp, Duration::from_secs(20));
401    }
402
403    #[test]
404    fn test_pcapng_write_packet_convenience() {
405        let eth = sample_ethernet_packet();
406        let pkt = Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth));
407
408        let mut buf = Vec::new();
409        {
410            let mut writer = PcapNgStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
411            writer.write_packet(&pkt).unwrap();
412        }
413
414        let iter = PcapNgIterator::from_reader(Cursor::new(buf)).unwrap();
415        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
416        assert_eq!(packets.len(), 1);
417        assert_eq!(packets[0].metadata.timestamp, Duration::ZERO);
418    }
419}