wireparse 0.1.0

use crate::io::{
    ReadData, ReadDataUnchecked, Reader, UncheckedReader, UncheckedWriter, WriteDataUnchecked,
};
use crate::ArpPayload;
use crate::EthFrameError;
use crate::MacAddress;
use crate::NetpackError;
use crate::Result;

/// Parsed ethernet ethernet frame struct
///
/// Use with pcap
/// ```no_run
/// use pcap::{Capture, Device};
///
/// let device = match Device::lookup()?;
///
/// let mut cap = Capture::from_device(device)?
///    .timeout(1000)
///    .promisc(true)
///    .snaplen(5000)
///    .open()?;
///
/// while let Ok(raw_packet) = cap.next_packet() {
///    let eth_packet = EthernetFrame::from_slice(raw_packet.data);
///
///    // use the parsed ethernet data
/// }
/// ```
#[derive(Clone, Debug, Eq, PartialEq, Default)]
pub struct EthernetFrame<'a> {
    pub header: EthernetFrameHeader<'a>,
    pub payload: EthernetFramePayload<'a>,
}

impl<'a> EthernetFrame<'a> {
    pub fn size(&self) -> usize {
        self.header.size() + self.payload.size()
    }

    pub fn from_slice(data: &'a [u8]) -> Result<Self> {
        let reader = &mut data[..].as_ref();
        reader.read()
    }

    pub fn write_to_slice(&self, buffer: &mut [u8]) -> Result<()> {
        if buffer.len() < self.size() {
            return Err(NetpackError::BufferTooSmall {
                provided_size: buffer.len(),
                requried_size: self.size(),
            });
        }

        let writer = &mut buffer[..].as_mut();
        writer.write_unchecked(self);

        Ok(())
    }

    #[cfg(feature = "alloc")]
    pub fn encode(&self) -> Result<alloc::vec::Vec<u8>> {
        let mut buffer = alloc::vec![0;self.size()];
        self.write_to_slice(&mut buffer)?;
        Ok(buffer)
    }
}

impl<'a> ReadData<'a> for EthernetFrame<'a> {
    fn read(reader: &mut impl crate::io::Reader<'a>) -> Result<Self> {
        let header: EthernetFrameHeader = reader.read()?;
        let payload = EthernetFramePayload::from_reader(header.ether_type, reader)?;
        Ok(Self { header, payload })
    }
}

impl<'a> WriteDataUnchecked for &EthernetFrame<'a> {
    fn write_to_unchecked(self, writer: &mut impl UncheckedWriter) {
        writer.write_unchecked(&self.header);
        writer.write_unchecked(&self.payload);
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Default)]
pub struct EthernetFrameHeader<'a> {
    pub destination: MacAddress<'a>,
    pub source: MacAddress<'a>,
    // TODO: double tagging
    pub dot1q_header: Option<Dot1qHeader>,
    pub ether_type: EtherProtocal,
}

impl<'a> EthernetFrameHeader<'a> {
    pub const LEN: usize = 14;
    pub const LEN_WITH_DOT1Q: usize = 18;

    pub fn size(&self) -> usize {
        if self.dot1q_header.is_some() {
            Self::LEN_WITH_DOT1Q
        } else {
            Self::LEN
        }
    }

    pub fn from_slice(data: &'a [u8]) -> Result<Self> {
        if data.len() < Self::LEN {
            return Err(EthFrameError::NotEnoughBytes {
                provided_size: data.len(),
                expected_size: Self::LEN,
            }
            .into());
        }

        let reader = &mut data[..].as_ref();

        let destination = reader.read_unchecked();
        let source = reader.read_unchecked();
        let ether_type = reader.read_unchecked();

        let (dot1q_header, ether_type) = if matches!(ether_type, EtherProtocal::VLanTaggedFrame) {
            if data.len() < Self::LEN_WITH_DOT1Q {
                return Err(EthFrameError::NotEnoughBytes {
                    provided_size: data.len(),
                    expected_size: Self::LEN_WITH_DOT1Q,
                }
                .into());
            }

            let reader = &mut data[12..].as_ref();

            let dot1q_header = reader.read_unchecked();
            let ether_type = reader.read_unchecked();

            (Some(dot1q_header), ether_type)
        } else {
            (None, ether_type)
        };

        Ok(Self {
            destination,
            source,
            dot1q_header,
            ether_type,
        })
    }
}

impl<'a> ReadData<'a> for EthernetFrameHeader<'a> {
    fn read(reader: &mut impl Reader<'a>) -> crate::Result<Self> {
        if reader.length() < Self::LEN {
            return Err(EthFrameError::NotEnoughBytes {
                provided_size: reader.length(),
                expected_size: Self::LEN,
            }
            .into());
        }

        let destination = reader.read_unchecked();
        let source = reader.read_unchecked();
        let ether_type = reader.read_unchecked();

        let (dot1q_header, ether_type) = if matches!(ether_type, EtherProtocal::VLanTaggedFrame) {
            if reader.length() < Self::LEN_WITH_DOT1Q {
                return Err(EthFrameError::NotEnoughBytes {
                    provided_size: reader.length(),
                    expected_size: Self::LEN_WITH_DOT1Q,
                }
                .into());
            }

            let dot1q_header = Dot1qHeader {
                tpid: ether_type.to_u16(),
                tci: reader.read_unchecked(),
            };
            let ether_type = reader.read_unchecked();

            (Some(dot1q_header), ether_type)
        } else {
            (None, ether_type)
        };

        Ok(Self {
            destination,
            source,
            dot1q_header,
            ether_type,
        })
    }
}

impl<'a> WriteDataUnchecked for &EthernetFrameHeader<'a> {
    fn write_to_unchecked(self, writer: &mut impl crate::io::UncheckedWriter) {
        writer.write_unchecked(&self.destination);
        writer.write_unchecked(&self.source);
        if let Some(ref dotq) = self.dot1q_header {
            writer.write_unchecked(dotq);
        }
        writer.write_unchecked(self.ether_type);
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Default)]
pub struct Dot1qHeader {
    pub tpid: u16,
    pub tci: Dot1qTCI,
}

impl<'a> ReadDataUnchecked<'a> for Dot1qHeader {
    fn read_unchecked(reader: &mut impl crate::io::UncheckedReader<'a>) -> Self {
        let tpid: u16 = reader.read_unchecked();
        let tci: u16 = reader.read_unchecked();
        Self {
            tpid,
            tci: Dot1qTCI::from_u16(tci),
        }
    }
}

impl WriteDataUnchecked for &Dot1qHeader {
    fn write_to_unchecked(self, writer: &mut impl crate::io::UncheckedWriter) {
        writer.write_unchecked(self.tpid);
        writer.write_unchecked(self.tci.to_u16());
    }
}

impl Dot1qHeader {
    pub const LEN: usize = 2 + Dot1qTCI::LEN;
}

#[derive(Clone, Debug, Eq, PartialEq, Default)]
pub struct Dot1qTCI {
    pub pcp: u8,
    pub dei: u8,
    pub vid: u16,
}
impl Dot1qTCI {
    pub const LEN: usize = 2;

    pub fn from_u16(data: u16) -> Self {
        let pcp = ((data >> 13) & 0x_111) as u8;
        let dei = ((data >> 10) & 0b_111) as u8;
        let vid = (data) & 0b_111_111_111_111;
        Self { pcp, dei, vid }
    }

    pub fn to_u16(&self) -> u16 {
        let mut tci = 0;
        tci |= (self.pcp as u16) << 13;
        tci |= (self.dei as u16) << 10;
        tci |= self.vid;
        tci
    }
}

impl<'a> ReadDataUnchecked<'a> for Dot1qTCI {
    fn read_unchecked(reader: &mut impl UncheckedReader<'a>) -> Self {
        Self::from_u16(reader.read_unchecked())
    }
}

impl WriteDataUnchecked for Dot1qTCI {
    fn write_to_unchecked(self, writer: &mut impl UncheckedWriter) {
        writer.write_unchecked(self.to_u16())
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum EtherProtocal {
    Unknown(u16),
    IPv4,
    Arp,
    VLanTaggedFrame,
}

impl Default for EtherProtocal {
    fn default() -> Self {
        Self::Unknown(0)
    }
}

impl EtherProtocal {
    #[inline]
    pub fn from_u16(val: u16) -> Self {
        match val {
            0x0800 => Self::IPv4,
            0x0806 => Self::Arp,
            0x8100 => Self::VLanTaggedFrame,
            _ => Self::Unknown(val),
        }
    }

    #[inline]
    pub fn to_u16(&self) -> u16 {
        match *self {
            Self::IPv4 => 0x0800,
            Self::Arp => 0x0806,
            Self::VLanTaggedFrame => 0x8100,
            Self::Unknown(v) => v,
        }
    }
}

impl<'a> ReadDataUnchecked<'a> for EtherProtocal {
    fn read_unchecked(reader: &mut impl crate::io::UncheckedReader<'a>) -> Self {
        Self::from_u16(reader.read_unchecked())
    }
}

impl WriteDataUnchecked for EtherProtocal {
    fn write_to_unchecked(self, writer: &mut impl crate::io::UncheckedWriter) {
        writer.write_unchecked(self.to_u16())
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum EthernetFramePayload<'a> {
    Raw(&'a [u8]),
    Arp(ArpPayload<'a>),
}

impl<'a> Default for EthernetFramePayload<'a> {
    fn default() -> Self {
        Self::Raw(&[])
    }
}

impl<'a> EthernetFramePayload<'a> {
    pub fn size(&self) -> usize {
        match self {
            Self::Raw(b) => b.len(),
            Self::Arp(arp) => arp.size(),
        }
    }
    pub fn from_reader(ether_type: EtherProtocal, reader: &mut impl Reader<'a>) -> Result<Self> {
        let payload = match ether_type {
            EtherProtocal::IPv4 => {
                // TODO: impliment IPv4
                Self::Raw(reader.get_slice())
            }
            EtherProtocal::Arp => Self::Arp(reader.read()?),
            EtherProtocal::Unknown(_) | EtherProtocal::VLanTaggedFrame => {
                Self::Raw(reader.get_slice())
            }
        };
        Ok(payload)
    }
}

impl<'a> WriteDataUnchecked for &EthernetFramePayload<'a> {
    fn write_to_unchecked(self, writer: &mut impl UncheckedWriter) {
        match self {
            EthernetFramePayload::Raw(b) => writer.write_slice_unchecked(b),
            EthernetFramePayload::Arp(arp) => writer.write_unchecked(arp),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_arp_parse() {
        let arp_packet = [
            255, 255, 255, 255, 255, 255, 200, 200, 200, 200, 200, 200, 8, 6, 0, 1, 8, 0, 6, 4, 0,
            1, 234, 0, 80, 227, 124, 225, 192, 168, 1, 230, 0, 0, 0, 0, 0, 0, 192, 168, 1, 1,
        ];

        let eth = EthernetFrame::from_slice(&arp_packet).unwrap();

        let mut out_buffer = [0u8; 42];
        eth.write_to_slice(&mut out_buffer).unwrap();

        assert_eq!(arp_packet, out_buffer);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn test_arp_encode() {
        let arp_packet = [
            255, 255, 255, 255, 255, 255, 200, 200, 200, 200, 200, 200, 8, 6, 0, 1, 8, 0, 6, 4, 0,
            1, 234, 0, 80, 227, 124, 225, 192, 168, 1, 230, 0, 0, 0, 0, 0, 0, 192, 168, 1, 1,
        ];

        let eth = EthernetFrame::from_slice(&arp_packet).unwrap();

        let out_buffer = eth.encode().unwrap();

        assert_eq!(arp_packet, out_buffer.as_slice());
    }
}