stackforge_core/pcap/
writer.rs1use 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
15pub 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
46pub 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
63pub struct PcapStreamWriter<W: Write> {
67 inner: PcapFileWriter<W>,
68}
69
70impl PcapStreamWriter<BufWriter<File>> {
71 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 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 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 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, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, ]
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(ð));
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 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 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(ð)),
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(ð));
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}