pcapforge-core 0.0.1

use anyhow::{Result, Context, bail};
use pcap::Capture;
use pcap_file::{pcap::PcapReader, pcapng::PcapNgReader};
use std::fs::File;
use std::path::Path;
use std::io::Read;

pub enum CaptureSource {
    Pcap(Capture<pcap::Offline>),
    PcapFile(Vec<u8>),
    PcapNg(Vec<u8>),
}

pub struct PacketCapture {
    source: CaptureSource,
    filter: Option<String>,
}

impl PacketCapture {
    /// Open a capture file (supports both pcap and pcapng)
    pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        let extension = path.extension()
            .and_then(|s| s.to_str())
            .unwrap_or("");

        match extension.to_lowercase().as_str() {
            "pcap" | "cap" => Self::from_pcap_file(path),
            "pcapng" | "ntar" => Self::from_pcapng_file(path),
            _ => {
                // Try to detect format by reading magic bytes
                Self::auto_detect(path)
            }
        }
    }

    fn from_pcap_file(path: &Path) -> Result<Self> {
        let capture = Capture::from_file(path)
            .context("Failed to open pcap file")?;

        Ok(Self {
            source: CaptureSource::Pcap(capture),
            filter: None,
        })
    }

    fn from_pcapng_file(path: &Path) -> Result<Self> {
        let mut file = File::open(path)?;
        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer)?;

        Ok(Self {
            source: CaptureSource::PcapNg(buffer),
            filter: None,
        })
    }

    fn auto_detect(path: &Path) -> Result<Self> {
        let mut file = File::open(path)?;
        let mut magic = [0u8; 4];
        file.read_exact(&mut magic)?;

        // Check magic numbers
        match &magic {
            // pcap magic numbers
            [0xa1, 0xb2, 0xc3, 0xd4] | [0xd4, 0xc3, 0xb2, 0xa1] |
            [0xa1, 0xb2, 0x3c, 0x4d] | [0x4d, 0x3c, 0xb2, 0xa1] => {
                drop(file);
                Self::from_pcap_file(path)
            }
            // pcapng magic number
            [0x0a, 0x0d, 0x0d, 0x0a] => {
                drop(file);
                Self::from_pcapng_file(path)
            }
            _ => bail!("Unknown file format. Expected pcap or pcapng")
        }
    }

    /// Apply a BPF filter (only works with pcap backend currently)
    pub fn set_filter(&mut self, filter: &str) -> Result<()> {
        match &mut self.source {
            CaptureSource::Pcap(capture) => {
                capture.filter(filter, true)
                    .context("Failed to apply BPF filter")?;
                self.filter = Some(filter.to_string());
                Ok(())
            }
            _ => {
                // For pcapng, we'll need to filter manually during iteration
                self.filter = Some(filter.to_string());
                Ok(())
            }
        }
    }

    /// Get statistics about the capture file
    pub fn stats(&mut self) -> Result<CaptureStats> {
        let mut stats = CaptureStats::default();

        match &mut self.source {
            CaptureSource::Pcap(capture) => {
                while let Ok(packet) = capture.next_packet() {
                    stats.update(&packet.data, packet.header.len);
                }
            }
            CaptureSource::PcapFile(data) => {
                let mut reader = PcapReader::new(&data[..])?;
                while let Some(pkt) = reader.next_packet() {
                    let pkt = pkt?;
                    stats.update(&pkt.data, pkt.data.len() as u32);
                }
            }
            CaptureSource::PcapNg(data) => {
                let mut reader = PcapNgReader::new(&data[..])?;
                while let Some(block) = reader.next_block() {
                    match block {
                        Ok(pcap_file::pcapng::Block::EnhancedPacket(pkt)) => {
                            stats.update(&pkt.data, pkt.data.len() as u32);
                        }
                        Ok(pcap_file::pcapng::Block::SimplePacket(pkt)) => {
                            stats.update(&pkt.data, pkt.data.len() as u32);
                        }
                        _ => {}
                    }
                }
            }
        }

        Ok(stats)
    }

    /// Process packets with a callback function
    pub fn process_packets<F>(&mut self, mut callback: F) -> Result<()>
    where
        F: FnMut(ProcessedPacket) -> Result<()>,
    {
        match &mut self.source {
            CaptureSource::Pcap(capture) => {
                while let Ok(packet) = capture.next_packet() {
                    let processed = ProcessedPacket {
                        timestamp: packet.header.ts.tv_sec as u64 * 1_000_000
                            + packet.header.ts.tv_usec as u64,
                        data: packet.data.to_vec(),
                        len: packet.header.len,
                        caplen: packet.header.caplen,
                    };
                    callback(processed)?;
                }
            }
            CaptureSource::PcapFile(data) => {
                let mut reader = PcapReader::new(&data[..])?;
                while let Some(pkt) = reader.next_packet() {
                    let pkt = pkt?;
                    let processed = ProcessedPacket {
                        timestamp: pkt.timestamp.as_micros() as u64,
                        data: pkt.data.to_vec(),
                        len: pkt.orig_len,
                        caplen: pkt.data.len() as u32,
                    };
                    callback(processed)?;
                }
            }
            CaptureSource::PcapNg(data) => {
                let mut reader = PcapNgReader::new(&data[..])?;
                while let Some(block) = reader.next_block() {
                    match block {
                        Ok(pcap_file::pcapng::Block::EnhancedPacket(pkt)) => {
                            let processed = ProcessedPacket {
                                timestamp: pkt.timestamp.as_micros() as u64,
                                data: pkt.data.to_vec(),
                                len: pkt.original_len,
                                caplen: pkt.data.len() as u32,
                            };
                            callback(processed)?;
                        }
                        Ok(pcap_file::pcapng::Block::SimplePacket(pkt)) => {
                            let processed = ProcessedPacket {
                                timestamp: 0, // Simple packets don't have timestamps
                                data: pkt.data.to_vec(),
                                len: pkt.original_len,
                                caplen: pkt.data.len() as u32,
                            };
                            callback(processed)?;
                        }
                        _ => {}
                    }
                }
            }
        }
        Ok(())
    }
}

pub struct ProcessedPacket {
    pub timestamp: u64, // microseconds since epoch
    pub data: Vec<u8>,
    pub len: u32,
    pub caplen: u32,
}

#[derive(Default, Debug)]
pub struct CaptureStats {
    pub total_packets: u64,
    pub total_bytes: u64,
    pub avg_packet_size: u64,
    pub max_packet_size: u32,
    pub min_packet_size: u32,
}

impl CaptureStats {
    fn update(&mut self, _data: &[u8], len: u32) {
        self.total_packets += 1;
        self.total_bytes += len as u64;

        if len > self.max_packet_size {
            self.max_packet_size = len;
        }
        if len < self.min_packet_size || self.min_packet_size == 0 {
            self.min_packet_size = len;
        }

        if self.total_packets > 0 {
            self.avg_packet_size = self.total_bytes / self.total_packets;
        }
    }
}