diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs index e0a5a3962..ae737145b 100644 --- a/src/network/sniffer.rs +++ b/src/network/sniffer.rs @@ -1,8 +1,8 @@ use ::std::boxed::Box; use ::pnet_bandwhich_fork::datalink::{DataLinkReceiver, NetworkInterface}; -use ::pnet_bandwhich_fork::packet::ethernet::{EtherType, EthernetPacket}; -use ::pnet_bandwhich_fork::packet::ip::IpNextHeaderProtocol; +use ::pnet_bandwhich_fork::packet::ethernet::{EtherTypes, EthernetPacket}; +use ::pnet_bandwhich_fork::packet::ip::IpNextHeaderProtocols; use ::pnet_bandwhich_fork::packet::ipv4::Ipv4Packet; use ::pnet_bandwhich_fork::packet::tcp::TcpPacket; use ::pnet_bandwhich_fork::packet::udp::UdpPacket; @@ -56,51 +56,61 @@ impl Sniffer { } pub fn next(&mut self) -> Option { let bytes = self.network_frames.next().ok()?; - let packet = EthernetPacket::new(bytes)?; - match packet.get_ethertype() { - EtherType(2048) => { - let ip_packet = Ipv4Packet::new(packet.payload())?; - let (protocol, source_port, destination_port, data_length) = - match ip_packet.get_next_level_protocol() { - IpNextHeaderProtocol(6) => { - let message = TcpPacket::new(ip_packet.payload())?; - ( - Protocol::Tcp, - message.get_source(), - message.get_destination(), - ip_packet.payload().len() as u128, - ) - } - IpNextHeaderProtocol(17) => { - let datagram = UdpPacket::new(ip_packet.payload())?; - ( - Protocol::Udp, - datagram.get_source(), - datagram.get_destination(), - ip_packet.payload().len() as u128, - ) - } - _ => return None, - }; - let interface_name = self.network_interface.name.clone(); - let direction = Direction::new(&self.network_interface.ips, &ip_packet); - let from = SocketAddr::new(IpAddr::V4(ip_packet.get_source()), source_port); - let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port); + let ip_packet = Ipv4Packet::new(&bytes)?; + let version = ip_packet.get_version(); - let connection = match direction { - Direction::Download => { - Connection::new(from, to.ip(), destination_port, protocol)? + match version { + 4 => Self::handle_v4(ip_packet, &self.network_interface), + 6 => None, // FIXME v6 support! + _ => { + let pkg = EthernetPacket::new(bytes)?; + match pkg.get_ethertype() { + EtherTypes::Ipv4 => { + Self::handle_v4(Ipv4Packet::new(pkg.payload())?, &self.network_interface) } - Direction::Upload => Connection::new(to, from.ip(), source_port, protocol)?, - }; - Some(Segment { - interface_name, - connection, - data_length, - direction, - }) + _ => None, + } } - _ => None, } } + fn handle_v4(ip_packet: Ipv4Packet, network_interface: &NetworkInterface) -> Option { + let (protocol, source_port, destination_port, data_length) = + match ip_packet.get_next_level_protocol() { + IpNextHeaderProtocols::Tcp => { + let message = TcpPacket::new(ip_packet.payload())?; + ( + Protocol::Tcp, + message.get_source(), + message.get_destination(), + ip_packet.payload().len() as u128, + ) + } + IpNextHeaderProtocols::Udp => { + let datagram = UdpPacket::new(ip_packet.payload())?; + ( + Protocol::Udp, + datagram.get_source(), + datagram.get_destination(), + ip_packet.payload().len() as u128, + ) + } + _ => return None, + }; + + let interface_name = network_interface.name.clone(); + let direction = Direction::new(&network_interface.ips, &ip_packet); + let from = SocketAddr::new(IpAddr::V4(ip_packet.get_source()), source_port); + let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port); + + let connection = match direction { + Direction::Download => Connection::new(from, to.ip(), destination_port, protocol)?, + Direction::Upload => Connection::new(to, from.ip(), source_port, protocol)?, + }; + Some(Segment { + interface_name, + connection, + data_length, + direction, + }) + } } diff --git a/src/tests/cases/raw_mode.rs b/src/tests/cases/raw_mode.rs index 5353fd033..03c3684ed 100644 --- a/src/tests/cases/raw_mode.rs +++ b/src/tests/cases/raw_mode.rs @@ -36,6 +36,23 @@ fn build_tcp_packet( pkt.packet().to_vec() } +fn build_ip_tcp_packet( + source_ip: &str, + destination_ip: &str, + source_port: u16, + destination_port: u16, + payload: &'static [u8], +) -> Vec { + let mut pkt_buf = [0u8; 1500]; + let pkt = packet_builder!( + pkt_buf, + ipv4({set_source => ipv4addr!(source_ip), set_destination => ipv4addr!(destination_ip) }) / + tcp({set_source => source_port, set_destination => destination_port }) / + payload(payload) + ); + pkt.packet().to_vec() +} + fn format_raw_output(output: Vec) -> String { let stdout_utf8 = String::from_utf8(output).unwrap(); use regex::Regex; @@ -44,6 +61,25 @@ fn format_raw_output(output: Vec) -> String { format!("{}", replaced) } +#[test] +fn one_ip_packet_of_traffic() { + let network_frames = vec![NetworkFrames::new(vec![Some(build_ip_tcp_packet( + "10.0.0.2", + "1.1.1.1", + 443, + 12345, + b"I am a fake tcp packet", + ))]) as Box]; + let (_, _, backend) = test_backend_factory(190, 50); + let stdout = Arc::new(Mutex::new(Vec::new())); + let os_input = os_input_output_stdout(network_frames, 2, Some(stdout.clone())); + let opts = opts_raw(); + start(backend, os_input, opts); + let stdout = Arc::try_unwrap(stdout).unwrap().into_inner().unwrap(); + let formatted = format_raw_output(stdout); + assert_snapshot!(formatted); +} + #[test] fn one_packet_of_traffic() { let network_frames = vec![NetworkFrames::new(vec![Some(build_tcp_packet( diff --git a/src/tests/cases/snapshots/raw_mode__one_ip_packet_of_traffic.snap b/src/tests/cases/snapshots/raw_mode__one_ip_packet_of_traffic.snap new file mode 100644 index 000000000..f5b78d93f --- /dev/null +++ b/src/tests/cases/snapshots/raw_mode__one_ip_packet_of_traffic.snap @@ -0,0 +1,8 @@ +--- +source: src/tests/cases/raw_mode.rs +expression: formatted +--- +process: "1" up/down Bps: 42/0 connections: 1 +connection: :443 => 1.1.1.1:12345 (tcp) up/down Bps: 42/0 process: "1" +remote_address: 1.1.1.1 up/down Bps: 42/0 connections: 1 +