sniffnet-packet-parser 0.2.1

use std::net::IpAddr;

use etherparse::{ArpHardwareId, EtherType, LinkHeader, NetHeaders, TransportHeader, VlanHeader};

use crate::arp_type::ArpType;
use crate::headers::{LinkInfo, NetInfo, TransportInfo, get_sniffable_headers};
use crate::icmp_type::{IcmpTypeV4, IcmpTypeV6};
use crate::igmp_type::IgmpType;
use crate::link_type::LinkType;
use crate::protocol::Protocol;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
/// Parsed network packet with info extracted from its headers.
pub struct ParsedPacket {
    /// Info extracted from the data link layer header.
    pub link_info: LinkInfo,
    /// Info extracted from the network layer header.
    pub net_info: NetInfo,
    /// Info extracted from the transport layer header.
    pub transport_info: TransportInfo,
}

impl ParsedPacket {
    #[must_use]
    /// Parse one raw network packet from its link type and wire bytes.
    pub fn from_bytes(bytes: &[u8], link_type: LinkType) -> Option<ParsedPacket> {
        let headers = get_sniffable_headers(bytes, link_type)?;

        let vlan_header = headers.vlan();
        let link_info = analyze_link_header(headers.link, vlan_header);

        let is_arp = matches!(&headers.net, Some(NetHeaders::Arp(_)));

        let net_info = analyze_net_header(headers.net)?;

        let transport_info = if is_arp {
            Some(TransportInfo {
                src_port: None,
                dst_port: None,
                protocol: Protocol::Arp,
                icmp_type: None,
                igmp_type: None,
            })
        } else {
            analyze_transport_header(headers.transport)
        }?;

        Some(ParsedPacket {
            link_info,
            net_info,
            transport_info,
        })
    }

    #[must_use]
    /// Returns the total number of bytes in the packet.
    pub fn bytes_count(&self) -> usize {
        self.link_info.bytes.saturating_add(self.net_info.bytes)
    }
}

/// This function extracts info from the data link layer header passed as parameter.
fn analyze_link_header(
    link_header: Option<LinkHeader>,
    vlan_header: Option<VlanHeader>,
) -> LinkInfo {
    let (vlan_id, vlan_bytes) = match vlan_header {
        Some(VlanHeader::Single(single)) => (Some(single.vlan_id.value()), 4),
        Some(VlanHeader::Double(double)) => (Some(double.outer.vlan_id.value()), 8),
        None => (None, 0),
    };

    match link_header {
        Some(LinkHeader::Ethernet2(header)) => {
            let src_mac = Some(header.source);
            let dst_mac = Some(header.destination);
            let bytes = vlan_bytes + 14;
            LinkInfo {
                src_mac,
                dst_mac,
                vlan_id,
                bytes,
            }
        }
        Some(LinkHeader::LinuxSll(header)) => {
            let src_mac: Option<[u8; 6]> = if header.sender_address_valid_length == 6
                && header.arp_hrd_type == ArpHardwareId::ETHERNET
                && let Ok(sender) = header.sender_address[0..6].try_into()
            {
                Some(sender)
            } else {
                None
            };
            let bytes = vlan_bytes + 16;
            LinkInfo {
                src_mac,
                dst_mac: None,
                vlan_id,
                bytes,
            }
        }
        None => LinkInfo {
            src_mac: None,
            dst_mac: None,
            vlan_id,
            bytes: vlan_bytes,
        },
    }
}

/// This function extracts info from the network layer header passed as parameter.
/// Returns `None` if packet has to be skipped.
fn analyze_net_header(network_header: Option<NetHeaders>) -> Option<NetInfo> {
    match network_header {
        Some(NetHeaders::Ipv4(ipv4header, _)) => {
            let src_ip = IpAddr::from(ipv4header.source);
            let dst_ip = IpAddr::from(ipv4header.destination);
            let bytes = usize::from(ipv4header.total_len);
            let ether_type = EtherType::IPV4.0;
            Some(NetInfo {
                src_ip,
                dst_ip,
                arp_type: None,
                ether_type,
                bytes,
            })
        }
        Some(NetHeaders::Ipv6(ipv6header, _)) => {
            let src_ip = IpAddr::from(ipv6header.source);
            let dst_ip = IpAddr::from(ipv6header.destination);
            let bytes = usize::from(ipv6header.payload_length) + 40;
            let ether_type = EtherType::IPV6.0;
            Some(NetInfo {
                src_ip,
                dst_ip,
                arp_type: None,
                ether_type,
                bytes,
            })
        }
        Some(NetHeaders::Arp(arp_packet)) => {
            let (src_ip, dst_ip) = match arp_packet.proto_addr_type {
                EtherType::IPV4 => {
                    let src_ip =
                        match TryInto::<[u8; 4]>::try_into(arp_packet.sender_protocol_addr()) {
                            Ok(source) => IpAddr::from(source),
                            Err(_) => return None,
                        };
                    let dst_ip =
                        match TryInto::<[u8; 4]>::try_into(arp_packet.target_protocol_addr()) {
                            Ok(destination) => IpAddr::from(destination),
                            Err(_) => return None,
                        };
                    (src_ip, dst_ip)
                }
                EtherType::IPV6 => {
                    let src_ip =
                        match TryInto::<[u8; 16]>::try_into(arp_packet.sender_protocol_addr()) {
                            Ok(source) => IpAddr::from(source),
                            Err(_) => return None,
                        };
                    let dst_ip =
                        match TryInto::<[u8; 16]>::try_into(arp_packet.target_protocol_addr()) {
                            Ok(destination) => IpAddr::from(destination),
                            Err(_) => return None,
                        };
                    (src_ip, dst_ip)
                }
                _ => return None,
            };
            let bytes = arp_packet.packet_len();
            let arp_type = Some(ArpType::from_etherparse(arp_packet.operation));
            let ether_type = EtherType::ARP.0;
            Some(NetInfo {
                src_ip,
                dst_ip,
                arp_type,
                ether_type,
                bytes,
            })
        }
        None => None,
    }
}

/// This function extracts info from the transport layer header passed as parameter.
/// Returns `None` if packet has to be skipped.
fn analyze_transport_header(transport_header: Option<TransportHeader>) -> Option<TransportInfo> {
    match transport_header {
        Some(TransportHeader::Udp(udp_header)) => {
            let src_port = Some(udp_header.source_port);
            let dst_port = Some(udp_header.destination_port);
            let protocol = Protocol::Udp;
            Some(TransportInfo {
                src_port,
                dst_port,
                protocol,
                icmp_type: None,
                igmp_type: None,
            })
        }
        Some(TransportHeader::Tcp(tcp_header)) => {
            let src_port = Some(tcp_header.source_port);
            let dst_port = Some(tcp_header.destination_port);
            let protocol = Protocol::Tcp;
            Some(TransportInfo {
                src_port,
                dst_port,
                protocol,
                icmp_type: None,
                igmp_type: None,
            })
        }
        Some(TransportHeader::Icmpv4(icmpv4_header)) => {
            let protocol = Protocol::Icmpv4;
            let icmp_type = Some(IcmpTypeV4::from_etherparse(&icmpv4_header.icmp_type));
            Some(TransportInfo {
                src_port: None,
                dst_port: None,
                protocol,
                icmp_type,
                igmp_type: None,
            })
        }
        Some(TransportHeader::Icmpv6(icmpv6_header)) => {
            let protocol = Protocol::Icmpv6;
            let icmp_type = Some(IcmpTypeV6::from_etherparse(&icmpv6_header.icmp_type));
            Some(TransportInfo {
                src_port: None,
                dst_port: None,
                protocol,
                icmp_type,
                igmp_type: None,
            })
        }
        Some(TransportHeader::Igmp(igmp_header)) => {
            let protocol = Protocol::Igmp;
            let igmp_type = Some(IgmpType::from_etherparse(&igmp_header.igmp_type));
            Some(TransportInfo {
                src_port: None,
                dst_port: None,
                protocol,
                icmp_type: None,
                igmp_type,
            })
        }
        None => None,
    }
}