Skip to main content

stackforge_core/pcap/
writer.rs

1//! PCAP file writer.
2
3use std::fs::File;
4use std::io::{BufWriter, Write};
5use std::path::Path;
6use std::time::Duration;
7
8use pcap_file::pcap::{PcapHeader, PcapPacket, PcapWriter as PcapFileWriter};
9
10use crate::error::{PacketError, Result};
11use crate::packet::Packet;
12
13use super::{CapturedPacket, PcapMetadata};
14
15/// Write captured packets to a PCAP file.
16///
17/// Preserves timestamps and original length from [`CapturedPacket`] metadata.
18pub fn wrpcap(path: impl AsRef<Path>, packets: &[CapturedPacket]) -> Result<()> {
19    let file = File::create(path.as_ref()).map_err(|e| {
20        PacketError::Io(format!(
21            "failed to create {}: {}",
22            path.as_ref().display(),
23            e
24        ))
25    })?;
26    let writer = BufWriter::new(file);
27
28    let header = PcapHeader::default();
29    let mut pcap_writer = PcapFileWriter::with_header(writer, header)
30        .map_err(|e| PacketError::Io(format!("PCAP write error: {}", e)))?;
31
32    for cap in packets {
33        let pcap_pkt = PcapPacket::new(
34            cap.metadata.timestamp,
35            cap.metadata.orig_len,
36            cap.packet.as_bytes(),
37        );
38        pcap_writer
39            .write_packet(&pcap_pkt)
40            .map_err(|e| PacketError::Io(format!("PCAP write error: {}", e)))?;
41    }
42
43    Ok(())
44}
45
46/// Write plain packets to a PCAP file (convenience function).
47///
48/// Timestamps are set to zero, `orig_len` matches each packet's data length.
49pub fn wrpcap_packets(path: impl AsRef<Path>, packets: &[Packet]) -> Result<()> {
50    let captured: Vec<CapturedPacket> = packets
51        .iter()
52        .map(|pkt| CapturedPacket {
53            packet: pkt.clone(),
54            metadata: PcapMetadata {
55                timestamp: Duration::ZERO,
56                orig_len: pkt.len() as u32,
57            },
58        })
59        .collect();
60    wrpcap(path, &captured)
61}
62
63/// PCAP writer for streaming writes.
64///
65/// Writes packets one at a time without buffering them all in memory.
66pub struct PcapStreamWriter<W: Write> {
67    inner: PcapFileWriter<W>,
68}
69
70impl PcapStreamWriter<BufWriter<File>> {
71    /// Create a new PCAP file for writing.
72    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
73        let file = File::create(path.as_ref()).map_err(|e| {
74            PacketError::Io(format!(
75                "failed to create {}: {}",
76                path.as_ref().display(),
77                e
78            ))
79        })?;
80        let writer = BufWriter::new(file);
81        Self::from_writer(writer)
82    }
83}
84
85impl<W: Write> PcapStreamWriter<W> {
86    /// Create a `PcapStreamWriter` from any writer.
87    pub fn from_writer(writer: W) -> Result<Self> {
88        let header = PcapHeader::default();
89        let pcap_writer = PcapFileWriter::with_header(writer, header)
90            .map_err(|e| PacketError::Io(format!("PCAP write error: {}", e)))?;
91        Ok(Self { inner: pcap_writer })
92    }
93
94    /// Write a captured packet with metadata.
95    pub fn write(&mut self, cap: &CapturedPacket) -> Result<()> {
96        let pcap_pkt = PcapPacket::new(
97            cap.metadata.timestamp,
98            cap.metadata.orig_len,
99            cap.packet.as_bytes(),
100        );
101        self.inner
102            .write_packet(&pcap_pkt)
103            .map_err(|e| PacketError::Io(format!("PCAP write error: {}", e)))?;
104        Ok(())
105    }
106
107    /// Write a plain packet (timestamp=0, orig_len=data length).
108    pub fn write_packet(&mut self, pkt: &Packet) -> Result<()> {
109        let pcap_pkt = PcapPacket::new(Duration::ZERO, pkt.len() as u32, pkt.as_bytes());
110        self.inner
111            .write_packet(&pcap_pkt)
112            .map_err(|e| PacketError::Io(format!("PCAP write error: {}", e)))?;
113        Ok(())
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use crate::pcap::reader::PcapIterator;
121    use std::io::Cursor;
122    use std::time::Duration;
123
124    fn sample_ethernet_packet() -> Vec<u8> {
125        vec![
126            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // dst: broadcast
127            0x00, 0x11, 0x22, 0x33, 0x44, 0x55, // src
128            0x08, 0x00, // EtherType: IPv4
129            0x00, 0x00, 0x00, 0x00, // dummy payload
130        ]
131    }
132
133    #[test]
134    fn test_wrpcap_roundtrip() {
135        let eth = sample_ethernet_packet();
136        let pkt = Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth));
137        let cap = CapturedPacket {
138            packet: pkt,
139            metadata: PcapMetadata {
140                timestamp: Duration::from_secs(42),
141                orig_len: eth.len() as u32,
142            },
143        };
144
145        // Write to in-memory buffer via PcapStreamWriter
146        let mut buf = Vec::new();
147        {
148            let mut writer = PcapStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
149            writer.write(&cap).unwrap();
150        }
151
152        // Read back
153        let iter = PcapIterator::from_reader(Cursor::new(buf)).unwrap();
154        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
155        assert_eq!(packets.len(), 1);
156        assert_eq!(packets[0].metadata.timestamp, Duration::from_secs(42));
157        assert_eq!(packets[0].packet.as_bytes(), eth.as_slice());
158    }
159
160    #[test]
161    fn test_wrpcap_multiple_packets() {
162        let eth = sample_ethernet_packet();
163
164        let caps: Vec<CapturedPacket> = (0..5)
165            .map(|i| CapturedPacket {
166                packet: Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth)),
167                metadata: PcapMetadata {
168                    timestamp: Duration::from_secs(i),
169                    orig_len: eth.len() as u32,
170                },
171            })
172            .collect();
173
174        let mut buf = Vec::new();
175        {
176            let mut writer = PcapStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
177            for cap in &caps {
178                writer.write(cap).unwrap();
179            }
180        }
181
182        let iter = PcapIterator::from_reader(Cursor::new(buf)).unwrap();
183        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
184        assert_eq!(packets.len(), 5);
185        for (i, pkt) in packets.iter().enumerate() {
186            assert_eq!(pkt.metadata.timestamp, Duration::from_secs(i as u64));
187        }
188    }
189
190    #[test]
191    fn test_write_packet_convenience() {
192        let eth = sample_ethernet_packet();
193        let pkt = Packet::from_bytes(bytes::Bytes::copy_from_slice(&eth));
194
195        let mut buf = Vec::new();
196        {
197            let mut writer = PcapStreamWriter::from_writer(Cursor::new(&mut buf)).unwrap();
198            writer.write_packet(&pkt).unwrap();
199        }
200
201        let iter = PcapIterator::from_reader(Cursor::new(buf)).unwrap();
202        let packets: Vec<_> = iter.collect::<std::result::Result<Vec<_>, _>>().unwrap();
203        assert_eq!(packets.len(), 1);
204        assert_eq!(packets[0].metadata.timestamp, Duration::ZERO);
205        assert_eq!(packets[0].metadata.orig_len, eth.len() as u32);
206    }
207}