Merge pull request #98 from Ma27/fix-network-gathering

Fix parsing of IP-Packets from Layer3

3 files changed, 98 insertions, 44 deletions

diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs
index e0a5a39..ae73714 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<Segment> {
         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<Segment> {
+        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 5353fd0..03c3684 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<u8> {
+    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<u8>) -> String {
     let stdout_utf8 = String::from_utf8(output).unwrap();
     use regex::Regex;
@@ -45,6 +62,25 @@ fn format_raw_output(output: Vec<u8>) -> String {
 }
 
 #[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<dyn DataLinkReceiver>];
+    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(
         "10.0.0.2",
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 0000000..f5b78d9
--- /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: <TIMESTAMP_REMOVED> "1" up/down Bps: 42/0 connections: 1
+connection: <TIMESTAMP_REMOVED> <interface_name>:443 => 1.1.1.1:12345 (tcp) up/down Bps: 42/0 process: "1"
+remote_address: <TIMESTAMP_REMOVED> 1.1.1.1 up/down Bps: 42/0 connections: 1
+