Sniff Ipv6 packets

9 files changed, 105 insertions, 71 deletions

diff --git a/src/display/components/table.rs b/src/display/components/table.rs
index 2b823d7..ff5a148 100644
--- a/src/display/components/table.rs
+++ b/src/display/components/table.rs
@@ -9,7 +9,7 @@ use ::tui::widgets::{Block, Borders, Row, Widget};
 use crate::display::{Bandwidth, DisplayBandwidth, UIState};
 use crate::network::{display_connection_string, display_ip_or_host};
 
-use ::std::net::Ipv4Addr;
+use ::std::net::IpAddr;
 use std::iter::FromIterator;
 
 const FIRST_WIDTH_BREAKPOINT: u16 = 50;
@@ -54,10 +54,7 @@ pub struct Table<'a> {
 }
 
 impl<'a> Table<'a> {
-    pub fn create_connections_table(
-        state: &UIState,
-        ip_to_host: &HashMap<Ipv4Addr, String>,
-    ) -> Self {
+    pub fn create_connections_table(state: &UIState, ip_to_host: &HashMap<IpAddr, String>) -> Self {
         let mut connections_list = Vec::from_iter(&state.connections);
         sort_by_bandwidth(&mut connections_list);
         let connections_rows = connections_list
@@ -105,7 +102,7 @@ impl<'a> Table<'a> {
     }
     pub fn create_remote_addresses_table(
         state: &UIState,
-        ip_to_host: &HashMap<Ipv4Addr, String>,
+        ip_to_host: &HashMap<IpAddr, String>,
     ) -> Self {
         let mut remote_addresses_list = Vec::from_iter(&state.remote_addresses);
         sort_by_bandwidth(&mut remote_addresses_list);
diff --git a/src/display/ui.rs b/src/display/ui.rs
index 0c4d36a..2c44d23 100644
--- a/src/display/ui.rs
+++ b/src/display/ui.rs
@@ -7,7 +7,7 @@ use crate::display::components::{HelpText, Layout, Table, TotalBandwidth};
 use crate::display::UIState;
 use crate::network::{display_connection_string, display_ip_or_host, LocalSocket, Utilization};
 
-use ::std::net::Ipv4Addr;
+use ::std::net::IpAddr;
 
 use chrono::prelude::*;
 
@@ -17,7 +17,7 @@ where
 {
     terminal: Terminal<B>,
     state: UIState,
-    ip_to_host: HashMap<Ipv4Addr, String>,
+    ip_to_host: HashMap<IpAddr, String>,
 }
 
 impl<B> Ui<B>
@@ -101,7 +101,7 @@ where
         &mut self,
         connections_to_procs: HashMap<LocalSocket, String>,
         utilization: Utilization,
-        ip_to_host: HashMap<Ipv4Addr, String>,
+        ip_to_host: HashMap<IpAddr, String>,
     ) {
         self.state.update(connections_to_procs, utilization);
         self.ip_to_host.extend(ip_to_host);
diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs
index 7a22c6d..7e8bb78 100644
--- a/src/display/ui_state.rs
+++ b/src/display/ui_state.rs
@@ -65,7 +65,7 @@ pub struct UtilizationData {
 #[derive(Default)]
 pub struct UIState {
     pub processes: BTreeMap<String, NetworkData>,
-    pub remote_addresses: BTreeMap<Ipv4Addr, NetworkData>,
+    pub remote_addresses: BTreeMap<IpAddr, NetworkData>,
     pub connections: BTreeMap<Connection, ConnectionData>,
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
@@ -106,7 +106,7 @@ impl UIState {
             self.utilization_data.pop_front();
         }
         let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new();
-        let mut remote_addresses: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new();
+        let mut remote_addresses: BTreeMap<IpAddr, NetworkData> = BTreeMap::new();
         let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new();
         let mut total_bytes_downloaded: u128 = 0;
         let mut total_bytes_uploaded: u128 = 0;
diff --git a/src/network/connection.rs b/src/network/connection.rs
index bdbf448..c0c35ea 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -1,6 +1,6 @@
 use ::std::collections::HashMap;
 use ::std::fmt;
-use ::std::net::{IpAddr, Ipv4Addr};
+use ::std::net::IpAddr;
 
 use ::std::net::SocketAddr;
 
@@ -35,7 +35,7 @@ impl fmt::Display for Protocol {
 
 #[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash, Debug, Copy)]
 pub struct Socket {
-    pub ip: Ipv4Addr,
+    pub ip: IpAddr,
     pub port: u16,
 }
 
@@ -52,7 +52,7 @@ pub struct Connection {
     pub local_socket: LocalSocket,
 }
 
-pub fn display_ip_or_host(ip: Ipv4Addr, ip_to_host: &HashMap<Ipv4Addr, String>) -> String {
+pub fn display_ip_or_host(ip: IpAddr, ip_to_host: &HashMap<IpAddr, String>) -> String {
     match ip_to_host.get(&ip) {
         Some(host) => host.clone(),
         None => ip.to_string(),
@@ -61,7 +61,7 @@ pub fn display_ip_or_host(ip: Ipv4Addr, ip_to_host: &HashMap<Ipv4Addr, String>)
 
 pub fn display_connection_string(
     connection: &Connection,
-    ip_to_host: &HashMap<Ipv4Addr, String>,
+    ip_to_host: &HashMap<IpAddr, String>,
     interface_name: &str,
 ) -> String {
     format!(
@@ -80,20 +80,17 @@ impl Connection {
         local_ip: IpAddr,
         local_port: u16,
         protocol: Protocol,
-    ) -> Option<Self> {
-        match remote_socket {
-            SocketAddr::V4(remote_socket) => Some(Connection {
-                remote_socket: Socket {
-                    ip: *remote_socket.ip(),
-                    port: remote_socket.port(),
-                },
-                local_socket: LocalSocket {
-                    ip: local_ip,
-                    port: local_port,
-                    protocol,
-                },
-            }),
-            _ => None,
+    ) -> Self {
+        Connection {
+            remote_socket: Socket {
+                ip: remote_socket.ip(),
+                port: remote_socket.port(),
+            },
+            local_socket: LocalSocket {
+                ip: local_ip,
+                port: local_port,
+                protocol,
+            },
         }
     }
 }
diff --git a/src/network/dns/client.rs b/src/network/dns/client.rs
index f9d12ed..9e13604 100644
--- a/src/network/dns/client.rs
+++ b/src/network/dns/client.rs
@@ -1,7 +1,7 @@
 use crate::network::dns::{resolver::Lookup, IpTable};
 use std::{
     collections::HashSet,
-    net::Ipv4Addr,
+    net::IpAddr,
     sync::{Arc, Mutex},
     thread::{Builder, JoinHandle},
 };
@@ -10,14 +10,14 @@ use tokio::{
     sync::mpsc::{self, Sender},
 };
 
-type PendingAddrs = HashSet<Ipv4Addr>;
+type PendingAddrs = HashSet<IpAddr>;
 
 const CHANNEL_SIZE: usize = 1_000;
 
 pub struct Client {
     cache: Arc<Mutex<IpTable>>,
     pending: Arc<Mutex<PendingAddrs>>,
-    tx: Option<Sender<Vec<Ipv4Addr>>>,
+    tx: Option<Sender<Vec<IpAddr>>>,
     handle: Option<JoinHandle<()>>,
 }
 
@@ -28,7 +28,7 @@ impl Client {
     {
         let cache = Arc::new(Mutex::new(IpTable::new()));
         let pending = Arc::new(Mutex::new(PendingAddrs::new()));
-        let (tx, mut rx) = mpsc::channel::<Vec<Ipv4Addr>>(CHANNEL_SIZE);
+        let (tx, mut rx) = mpsc::channel::<Vec<IpAddr>>(CHANNEL_SIZE);
 
         let handle = Builder::new().name("resolver".into()).spawn({
             let cache = cache.clone();
@@ -65,7 +65,7 @@ impl Client {
         })
     }
 
-    pub fn resolve(&mut self, ips: Vec<Ipv4Addr>) {
+    pub fn resolve(&mut self, ips: Vec<IpAddr>) {
         // Remove ips that are already being resolved
         let ips = ips
             .into_iter()
diff --git a/src/network/dns/mod.rs b/src/network/dns/mod.rs
index be1464e..2c911b0 100644
--- a/src/network/dns/mod.rs
+++ b/src/network/dns/mod.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, net::Ipv4Addr};
+use std::{collections::HashMap, net::IpAddr};
 
 mod client;
 mod resolver;
@@ -6,4 +6,4 @@ mod resolver;
 pub use client::*;
 pub use resolver::*;
 
-pub type IpTable = HashMap<Ipv4Addr, String>;
+pub type IpTable = HashMap<IpAddr, String>;
diff --git a/src/network/dns/resolver.rs b/src/network/dns/resolver.rs
index 007c485..074f46e 100644
--- a/src/network/dns/resolver.rs
+++ b/src/network/dns/resolver.rs
@@ -1,11 +1,11 @@
 use async_trait::async_trait;
-use std::net::Ipv4Addr;
+use std::net::IpAddr;
 use tokio::runtime::Handle;
 use trust_dns_resolver::{error::ResolveErrorKind, TokioAsyncResolver};
 
 #[async_trait]
 pub trait Lookup {
-    async fn lookup(&self, ip: Ipv4Addr) -> Option<String>;
+    async fn lookup(&self, ip: IpAddr) -> Option<String>;
 }
 
 pub struct Resolver(TokioAsyncResolver);
@@ -19,8 +19,8 @@ impl Resolver {
 
 #[async_trait]
 impl Lookup for Resolver {
-    async fn lookup(&self, ip: Ipv4Addr) -> Option<String> {
-        let lookup_future = self.0.reverse_lookup(ip.into());
+    async fn lookup(&self, ip: IpAddr) -> Option<String> {
+        let lookup_future = self.0.reverse_lookup(ip);
         match lookup_future.await {
             Ok(names) => {
                 // Take the first result and convert it to a string
diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs
index ae73714..ab23906 100644
--- a/src/network/sniffer.rs
+++ b/src/network/sniffer.rs
@@ -2,8 +2,9 @@ use ::std::boxed::Box;
 
 use ::pnet_bandwhich_fork::datalink::{DataLinkReceiver, NetworkInterface};
 use ::pnet_bandwhich_fork::packet::ethernet::{EtherTypes, EthernetPacket};
-use ::pnet_bandwhich_fork::packet::ip::IpNextHeaderProtocols;
+use ::pnet_bandwhich_fork::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
 use ::pnet_bandwhich_fork::packet::ipv4::Ipv4Packet;
+use ::pnet_bandwhich_fork::packet::ipv6::Ipv6Packet;
 use ::pnet_bandwhich_fork::packet::tcp::TcpPacket;
 use ::pnet_bandwhich_fork::packet::udp::UdpPacket;
 use ::pnet_bandwhich_fork::packet::Packet;
@@ -27,10 +28,10 @@ pub enum Direction {
 }
 
 impl Direction {
-    pub fn new(network_interface_ips: &[IpNetwork], ip_packet: &Ipv4Packet) -> Self {
+    pub fn new(network_interface_ips: &[IpNetwork], source: IpAddr) -> Self {
         if network_interface_ips
             .iter()
-            .any(|ip_network| ip_network.ip() == ip_packet.get_source())
+            .any(|ip_network| ip_network.ip() == source)
         {
             Direction::Upload
         } else {
@@ -39,6 +40,42 @@ impl Direction {
     }
 }
 
+trait NextLevelProtocol {
+    fn get_next_level_protocol(&self) -> IpNextHeaderProtocol;
+}
+
+impl NextLevelProtocol for Ipv6Packet<'_> {
+    fn get_next_level_protocol(&self) -> IpNextHeaderProtocol {
+        self.get_next_header()
+    }
+}
+
+macro_rules! extract_transport_protocol {
+    (  $ip_packet: ident ) => {{
+        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,
+        }
+    }};
+}
+
 pub struct Sniffer {
     network_interface: NetworkInterface,
     network_frames: Box<dyn DataLinkReceiver>,
@@ -61,50 +98,53 @@ impl Sniffer {
 
         match version {
             4 => Self::handle_v4(ip_packet, &self.network_interface),
-            6 => None, // FIXME v6 support!
+            6 => Self::handle_v6(Ipv6Packet::new(&bytes)?, &self.network_interface),
             _ => {
                 let pkg = EthernetPacket::new(bytes)?;
                 match pkg.get_ethertype() {
                     EtherTypes::Ipv4 => {
                         Self::handle_v4(Ipv4Packet::new(pkg.payload())?, &self.network_interface)
                     }
+                    EtherTypes::Ipv6 => {
+                        Self::handle_v6(Ipv6Packet::new(pkg.payload())?, &self.network_interface)
+                    }
                     _ => None,
                 }
             }
         }
     }
+    fn handle_v6(ip_packet: Ipv6Packet, network_interface: &NetworkInterface) -> Option<Segment> {
+        let (protocol, source_port, destination_port, data_length) =
+            extract_transport_protocol!(ip_packet);
+
+        let interface_name = network_interface.name.clone();
+        let direction = Direction::new(&network_interface.ips, ip_packet.get_source().into());
+        let from = SocketAddr::new(ip_packet.get_source().into(), source_port);
+        let to = SocketAddr::new(ip_packet.get_destination().into(), 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,
+        })
+    }
     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,
-            };
+            extract_transport_protocol!(ip_packet);
 
         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 direction = Direction::new(&network_interface.ips, ip_packet.get_source().into());
+        let from = SocketAddr::new(ip_packet.get_source().into(), source_port);
+        let to = SocketAddr::new(ip_packet.get_destination().into(), 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)?,
+            Direction::Download => Connection::new(from, to.ip(), destination_port, protocol),
+            Direction::Upload => Connection::new(to, from.ip(), source_port, protocol),
         };
         Some(Segment {
             interface_name,
diff --git a/src/os/lsof.rs b/src/os/lsof.rs
index 68b4bad..31e3266 100644
--- a/src/os/lsof.rs
+++ b/src/os/lsof.rs
@@ -29,7 +29,7 @@ pub(crate) fn get_open_sockets() -> OpenSockets {
         let local_port = raw_connection.get_local_port();
 
         let socket_addr = SocketAddr::new(remote_ip, remote_port);
-        let connection = Connection::new(socket_addr, local_ip, local_port, protocol).unwrap();
+        let connection = Connection::new(socket_addr, local_ip, local_port, protocol);
 
         open_sockets.insert(connection.local_socket, raw_connection.process_name.clone());
         connections_vec.push(connection);