src/ip_packet.rs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use std::net::IpAddr;

pub enum IpPacket<'a> {
    V4(Ipv4Packet<'a>),
    V6(Ipv6Packet<'a>),
}

impl<'a> IpPacket<'a> {
    pub fn new(packet: &'a [u8]) -> Option<Self> {
        if packet.len() == 0 {
            return None;
        }

        match packet[0] >> 4 {
            4 => Ipv4Packet::new(&packet).map(|packet| IpPacket::V4(packet)),
            6 => Ipv6Packet::new(&packet).map(|packet| IpPacket::V6(packet)),
            _ => None
        }
    }

    pub fn get_source(&self) -> IpAddr {
        match *self {
            IpPacket::V4(ref packet) => packet.get_source().into(),
            IpPacket::V6(ref packet) => packet.get_source().into(),
        }
    }

    pub fn get_destination(&self) -> IpAddr {
        match *self {
            IpPacket::V4(ref packet) => packet.get_destination().into(),
            IpPacket::V6(ref packet) => packet.get_destination().into(),
        }
    }
}