Add ipv4

src/packet/protocols/icmp.rs

@@ -23,37 +23,22 @@ impl<T> IcmpPacket<T> {
     }
 }
 
-impl<T> IcmpPacket<T>
+impl<T> TryFrom<Vec<u8>> for IcmpPacket<T>
 where
-    T: From<Vec<u8>>,
+    T: TryFrom<Vec<u8>, Error = PacketError>,
 {
-    /// Construct a new ICMPv6 packet from raw bytes
-    pub fn new_from_bytes(bytes: &[u8]) -> Result<Self, PacketError> {
+    type Error = PacketError;
+
+    fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
         // Parse the packet
         let packet =
-            pnet_packet::icmp::IcmpPacket::new(bytes).ok_or(PacketError::TooShort(bytes.len()))?;
+            pnet_packet::icmp::IcmpPacket::new(&bytes).ok_or(PacketError::TooShort(bytes.len()))?;
 
         // Return the packet
         Ok(Self {
             icmp_type: packet.get_icmp_type(),
             icmp_code: packet.get_icmp_code(),
-            payload: packet.payload().to_vec().into(),
-        })
-    }
-}
-
-impl IcmpPacket<Vec<u8>> {
-    /// Construct a new ICMPv6 packet with a raw payload from raw bytes
-    pub fn new_from_bytes_raw_payload(bytes: &[u8]) -> Result<Self, PacketError> {
-        // Parse the packet
-        let packet =
-            pnet_packet::icmp::IcmpPacket::new(bytes).ok_or(PacketError::TooShort(bytes.len()))?;
-
-        // Return the packet
-        Ok(Self {
-            icmp_type: packet.get_icmp_type(),
-            icmp_code: packet.get_icmp_code(),
-            payload: packet.payload().to_vec(),
+            payload: packet.payload().to_vec().try_into()?,
         })
     }
 }

src/packet/protocols/ipv4.rs

@@ -0,0 +1,131 @@
+use std::net::Ipv4Addr;
+
+use pnet_packet::{
+    ip::IpNextHeaderProtocol,
+    ipv4::{Ipv4Option, Ipv4OptionPacket},
+    Packet,
+};
+
+use crate::packet::error::PacketError;
+
+#[derive(Debug, Clone)]
+pub struct Ipv4Packet<T> {
+    pub dscp: u8,
+    pub ecn: u8,
+    pub identification: u16,
+    pub flags: u8,
+    pub fragment_offset: u16,
+    pub ttl: u8,
+    pub protocol: IpNextHeaderProtocol,
+    pub source_address: Ipv4Addr,
+    pub destination_address: Ipv4Addr,
+    pub options: Vec<Ipv4Option>,
+    pub payload: T,
+}
+
+impl<T> Ipv4Packet<T> {
+    /// Construct a new IPv4 packet
+    pub fn new(
+        dscp: u8,
+        ecn: u8,
+        identification: u16,
+        flags: u8,
+        fragment_offset: u16,
+        ttl: u8,
+        protocol: IpNextHeaderProtocol,
+        source_address: Ipv4Addr,
+        destination_address: Ipv4Addr,
+        options: Vec<Ipv4Option>,
+        payload: T,
+    ) -> Self {
+        Self {
+            dscp,
+            ecn,
+            identification,
+            flags,
+            fragment_offset,
+            ttl,
+            protocol,
+            source_address,
+            destination_address,
+            options,
+            payload,
+        }
+    }
+
+    fn options_length_words(&self) -> u8 {
+        self.options
+            .iter()
+            .map(|option| Ipv4OptionPacket::packet_size(option) as u8)
+            .sum::<u8>()
+            / 4
+    }
+}
+
+impl<T> TryFrom<Vec<u8>> for Ipv4Packet<T>
+where
+    T: From<Vec<u8>>,
+{
+    type Error = PacketError;
+
+    fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
+        // Parse the packet
+        let packet =
+            pnet_packet::ipv4::Ipv4Packet::new(&bytes).ok_or(PacketError::TooShort(bytes.len()))?;
+
+        // Return the packet
+        Ok(Self {
+            dscp: packet.get_dscp(),
+            ecn: packet.get_ecn(),
+            identification: packet.get_identification(),
+            flags: packet.get_flags(),
+            fragment_offset: packet.get_fragment_offset(),
+            ttl: packet.get_ttl(),
+            protocol: packet.get_next_level_protocol(),
+            source_address: packet.get_source(),
+            destination_address: packet.get_destination(),
+            options: packet.get_options(),
+            payload: packet.payload().to_vec().into(),
+        })
+    }
+}
+
+impl<T> Into<Vec<u8>> for Ipv4Packet<T>
+where
+    T: Into<Vec<u8>> + Clone,
+{
+    fn into(self) -> Vec<u8> {
+        // Convert the payload into raw bytes
+        let payload: Vec<u8> = self.payload.clone().into();
+
+        // Build the packet
+        let total_length = 20 + (self.options_length_words() as usize * 4) + payload.len();
+        let mut packet =
+            pnet_packet::ipv4::MutableIpv4Packet::owned(vec![0u8; total_length]).unwrap();
+
+        // Set the fields
+        packet.set_version(4);
+        packet.set_header_length(5 + self.options_length_words());
+        packet.set_dscp(self.dscp);
+        packet.set_ecn(self.ecn);
+        packet.set_total_length(total_length.try_into().unwrap());
+        packet.set_identification(self.identification);
+        packet.set_flags(self.flags);
+        packet.set_fragment_offset(self.fragment_offset);
+        packet.set_ttl(self.ttl);
+        packet.set_next_level_protocol(self.protocol);
+        packet.set_source(self.source_address);
+        packet.set_destination(self.destination_address);
+        packet.set_options(&self.options);
+
+        // Set the payload
+        packet.set_payload(&payload);
+
+        // Calculate the checksum
+        packet.set_checksum(0);
+        packet.set_checksum(pnet_packet::ipv4::checksum(&packet.to_immutable()));
+
+        // Return the packet
+        packet.to_immutable().packet().to_vec()
+    }
+}

src/packet/protocols/ipv6.rs

@@ -0,0 +1,95 @@
+use std::net::Ipv6Addr;
+
+use pnet_packet::{ip::IpNextHeaderProtocol, Packet};
+
+use crate::packet::error::PacketError;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct Ipv6Packet<T> {
+    pub traffic_class: u8,
+    pub flow_label: u32,
+    pub next_header: IpNextHeaderProtocol,
+    pub hop_limit: u8,
+    pub source_address: Ipv6Addr,
+    pub destination_address: Ipv6Addr,
+    pub payload: T,
+}
+
+impl<T> Ipv6Packet<T> {
+    /// Construct a new IPv6 packet
+    pub fn new(
+        traffic_class: u8,
+        flow_label: u32,
+        next_header: IpNextHeaderProtocol,
+        hop_limit: u8,
+        source_address: Ipv6Addr,
+        destination_address: Ipv6Addr,
+        payload: T,
+    ) -> Self {
+        Self {
+            traffic_class,
+            flow_label,
+            next_header,
+            hop_limit,
+            source_address,
+            destination_address,
+            payload,
+        }
+    }
+}
+
+impl<T> TryFrom<Vec<u8>> for Ipv6Packet<T>
+where
+    T: From<Vec<u8>>,
+{
+    type Error = PacketError;
+
+    fn try_from(bytes: Vec<u8>) -> Result<Self, Self::Error> {
+        // Parse the packet
+        let packet =
+            pnet_packet::ipv6::Ipv6Packet::new(&bytes).ok_or(PacketError::TooShort(bytes.len()))?;
+
+        // Return the packet
+        Ok(Self {
+            traffic_class: packet.get_traffic_class(),
+            flow_label: packet.get_flow_label(),
+            next_header: packet.get_next_header(),
+            hop_limit: packet.get_hop_limit(),
+            source_address: packet.get_source(),
+            destination_address: packet.get_destination(),
+            payload: packet.payload().to_vec().into(),
+        })
+    }
+}
+
+impl<T> Into<Vec<u8>> for Ipv6Packet<T>
+where
+    T: Into<Vec<u8>>,
+{
+    fn into(self) -> Vec<u8> {
+        // Convert the payload into raw bytes
+        let payload: Vec<u8> = self.payload.into();
+
+        // Allocate a mutable packet to write into
+        let total_length =
+            pnet_packet::ipv6::MutableIpv6Packet::minimum_packet_size() + payload.len();
+        let mut output =
+            pnet_packet::ipv6::MutableIpv6Packet::owned(vec![0u8; total_length]).unwrap();
+
+        // Write the header
+        output.set_version(6);
+        output.set_traffic_class(self.traffic_class);
+        output.set_flow_label(self.flow_label);
+        output.set_payload_length(payload.len() as u16);
+        output.set_next_header(self.next_header);
+        output.set_hop_limit(self.hop_limit);
+        output.set_source(self.source_address);
+        output.set_destination(self.destination_address);
+
+        // Write the payload
+        output.set_payload(&payload);
+
+        // Return the packet
+        output.to_immutable().packet().to_vec()
+    }
+}

src/packet/protocols/mod.rs

@@ -1,4 +1,6 @@
 pub mod icmp;
 pub mod icmpv6;
+pub mod ipv4;
+pub mod ipv6;
 pub mod tcp;
 pub mod udp;

src/packet/protocols/tcp.rs

@@ -188,8 +188,9 @@ where
         let payload: Vec<u8> = self.payload.into();
 
         // Allocate a mutable packet to write into
-        let total_length =
-            pnet_packet::tcp::MutableTcpPacket::minimum_packet_size() + payload.len();
+        let total_length = pnet_packet::tcp::MutableTcpPacket::minimum_packet_size()
+            + (self.options_length_words() as usize * 4)
+            + payload.len();
         let mut output =
             pnet_packet::tcp::MutableTcpPacket::owned(vec![0u8; total_length]).unwrap();