support ipv6 for dns lookup

pythops · Oct 16, 2024 · 619736d · 619736d
1 parent de9d981
commit 619736d
Show file tree

Hide file tree

Showing 8 changed files with 149 additions and 47 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Release.md b/Release.md
@@ -1,3 +1,9 @@
+## v0.5 - TBA
+
+### Added
+
+- stats: top 10 websites support ipv6
+
 ## v0.4 - 2024-10-13
 
 ### Added

diff --git a/oryx-tui/Cargo.toml b/oryx-tui/Cargo.toml
@@ -19,7 +19,6 @@ oryx-common = { path = "../oryx-common" }
 mio = { version = "1", features = ["os-poll", "os-ext"] }
 itertools = "0.13"
 dirs = "5"
-dns-lookup = "2"
 kanal = "0.1.0-pre8"
 mimalloc = "0.1"
 clap = { version = "4", features = ["derive", "cargo"] }

diff --git a/oryx-tui/src/dns.rs b/oryx-tui/src/dns.rs
@@ -0,0 +1,83 @@
+use libc::{
+    c_char, getnameinfo, sockaddr_in, sockaddr_in6, socklen_t, NI_MAXHOST, NI_NAMEREQD,
+    NI_NUMERICSERV,
+};
+use std::ffi::CStr;
+use std::mem;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+
+use crate::app::AppResult;
+
+pub fn get_hostname(ip: &IpAddr) -> AppResult<String> {
+    match ip {
+        IpAddr::V4(v) => get_hostname_v4(v),
+        IpAddr::V6(v) => get_hostname_v6(v),
+    }
+}
+
+fn get_hostname_v4(ip: &Ipv4Addr) -> AppResult<String> {
+    let sockaddr = sockaddr_in {
+        sin_family: libc::AF_INET as u16,
+        sin_port: 0,
+        sin_addr: libc::in_addr {
+            s_addr: ip.to_bits(),
+        },
+        sin_zero: [0; 8],
+    };
+
+    let mut host: [c_char; NI_MAXHOST as usize] = [0; NI_MAXHOST as usize];
+
+    let result = unsafe {
+        getnameinfo(
+            &sockaddr as *const _ as *const _,
+            mem::size_of::<sockaddr_in>() as socklen_t,
+            host.as_mut_ptr(),
+            NI_MAXHOST,
+            std::ptr::null_mut(),
+            0,
+            NI_NUMERICSERV | NI_NAMEREQD,
+        )
+    };
+
+    if result != 0 {
+        return Err("Failed to get hostname".into());
+    }
+
+    let host_str = unsafe { CStr::from_ptr(host.as_ptr()).to_string_lossy().into_owned() };
+
+    Ok(host_str)
+}
+
+fn get_hostname_v6(ip: &Ipv6Addr) -> AppResult<String> {
+    let sockaddr = sockaddr_in6 {
+        sin6_family: libc::AF_INET6 as u16,
+        sin6_port: 0,
+        sin6_addr: libc::in6_addr {
+            s6_addr: ip.octets(),
+        },
+        sin6_flowinfo: 0,
+        sin6_scope_id: 0,
+    };
+
+    let mut host: [c_char; NI_MAXHOST as usize] = [0; NI_MAXHOST as usize];
+
+    let result = unsafe {
+        getnameinfo(
+            &sockaddr as *const _ as *const _,
+            mem::size_of::<sockaddr_in6>() as socklen_t,
+            host.as_mut_ptr(),
+            NI_MAXHOST,
+            std::ptr::null_mut(),
+            0,
+            NI_NUMERICSERV | NI_NAMEREQD,
+        )
+    };
+
+    if result != 0 {
+        return Err("Failed to get hostname".into());
+    }
+
+    let host_str = unsafe { CStr::from_ptr(host.as_ptr()).to_string_lossy().into_owned() };
+
+    Ok(host_str)
+}
diff --git a/oryx-tui/src/handler.rs b/oryx-tui/src/handler.rs
@@ -4,7 +4,7 @@ use crate::{
     app::{ActivePopup, App, AppResult},
     event::Event,
     filter::FocusedBlock,
-    section::FocusedSection,
+    section::{stats::Stats, FocusedSection},
 };
 use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers};
 
@@ -18,6 +18,10 @@ pub fn handle_key_events(
         match key_event.code {
             KeyCode::Enter => {
                 if app.filter.focused_block == FocusedBlock::Apply {
+                    app.section.stats = Some(Stats::new(
+                        app.packets.clone(),
+                        app.filter.interface.selected_interface.clone(),
+                    ));
                     app.filter
                         .start(event_sender.clone(), app.data_channel_sender.clone())?;
 

diff --git a/oryx-tui/src/lib.rs b/oryx-tui/src/lib.rs
@@ -25,3 +25,5 @@ pub mod bandwidth;
 pub mod packet;
 
 pub mod section;
+
+pub mod dns;
diff --git a/oryx-tui/src/section.rs b/oryx-tui/src/section.rs
@@ -38,7 +38,7 @@ pub enum FocusedSection {
 pub struct Section {
     pub focused_section: FocusedSection,
     pub inspection: Inspection,
-    pub stats: Stats,
+    pub stats: Option<Stats>,
     pub alert: Alert,
     pub firewall: Firewall,
 }
@@ -51,7 +51,7 @@ impl Section {
         Self {
             focused_section: FocusedSection::Inspection,
             inspection: Inspection::new(packets.clone()),
-            stats: Stats::new(packets.clone()),
+            stats: None,
             alert: Alert::new(packets.clone()),
             firewall: Firewall::new(firewall_chans.ingress.sender, firewall_chans.egress.sender),
         }
@@ -254,7 +254,11 @@ impl Section {
 
         match self.focused_section {
             FocusedSection::Inspection => self.inspection.render(frame, section_block),
-            FocusedSection::Stats => self.stats.render(frame, section_block, network_interace),
+            FocusedSection::Stats => {
+                if let Some(stats) = &self.stats {
+                    stats.render(frame, section_block, network_interace)
+                }
+            }
             FocusedSection::Alerts => self.alert.render(frame, section_block),
             FocusedSection::Firewall => self.firewall.render(frame, section_block),
         }

diff --git a/oryx-tui/src/section/stats.rs b/oryx-tui/src/section/stats.rs
@@ -1,7 +1,6 @@
-use dns_lookup::lookup_addr;
 use std::{
     collections::HashMap,
-    net::{IpAddr, Ipv4Addr},
+    net::IpAddr,
     sync::{Arc, Mutex},
     thread,
     time::Duration,
@@ -17,6 +16,8 @@ use ratatui::{
 
 use crate::{
     bandwidth::Bandwidth,
+    dns::get_hostname,
+    interface::NetworkInterface,
     packet::{
         network::{IpPacket, IpProto},
         AppPacket,
@@ -30,7 +31,7 @@ pub struct PacketStats {
     pub network: NetworkStats,
     pub transport: TransportStats,
     pub link: LinkStats,
-    pub addresses: HashMap<Ipv4Addr, (Option<String>, usize)>,
+    pub addresses: HashMap<IpAddr, (Option<String>, usize)>,
 }
 
 #[derive(Debug)]
@@ -40,7 +41,7 @@ pub struct Stats {
 }
 
 impl Stats {
-    pub fn new(packets: Arc<Mutex<Vec<AppPacket>>>) -> Self {
+    pub fn new(packets: Arc<Mutex<Vec<AppPacket>>>, selected_interface: NetworkInterface) -> Self {
         let packet_stats: Arc<Mutex<PacketStats>> = Arc::new(Mutex::new(PacketStats::default()));
 
         thread::spawn({
@@ -67,20 +68,22 @@ impl Stats {
                                     if !ipv4_packet.dst_ip.is_private()
                                         && !ipv4_packet.dst_ip.is_loopback()
                                     {
-                                        if let Some((_, counts)) =
-                                            packet_stats.addresses.get_mut(&ipv4_packet.dst_ip)
+                                        if let Some((_, counts)) = packet_stats
+                                            .addresses
+                                            .get_mut(&IpAddr::V4(ipv4_packet.dst_ip))
                                         {
                                             *counts += 1;
                                         } else if let Ok(host) =
-                                            lookup_addr(&IpAddr::V4(ipv4_packet.dst_ip))
+                                            get_hostname(&IpAddr::V4(ipv4_packet.dst_ip))
                                         {
-                                            packet_stats
-                                                .addresses
-                                                .insert(ipv4_packet.dst_ip, (Some(host), 1));
+                                            packet_stats.addresses.insert(
+                                                IpAddr::V4(ipv4_packet.dst_ip),
+                                                (Some(host), 1),
+                                            );
                                         } else {
                                             packet_stats
                                                 .addresses
-                                                .insert(ipv4_packet.dst_ip, (None, 1));
+                                                .insert(IpAddr::V4(ipv4_packet.dst_ip), (None, 1));
                                         }
                                     }
 
@@ -98,6 +101,30 @@ impl Stats {
                                 }
                                 IpPacket::V6(ipv6_packet) => {
                                     packet_stats.network.ipv6 += 1;
+
+                                    if !selected_interface
+                                        .addresses
+                                        .contains(&IpAddr::V6(ipv6_packet.dst_ip))
+                                    {
+                                        if let Some((_, counts)) = packet_stats
+                                            .addresses
+                                            .get_mut(&IpAddr::V6(ipv6_packet.dst_ip))
+                                        {
+                                            *counts += 1;
+                                        } else if let Ok(host) =
+                                            get_hostname(&IpAddr::V6(ipv6_packet.dst_ip))
+                                        {
+                                            packet_stats.addresses.insert(
+                                                IpAddr::V6(ipv6_packet.dst_ip),
+                                                (Some(host), 1),
+                                            );
+                                        } else {
+                                            packet_stats
+                                                .addresses
+                                                .insert(IpAddr::V6(ipv6_packet.dst_ip), (None, 1));
+                                        }
+                                    }
+
                                     match ipv6_packet.proto {
                                         IpProto::Tcp(_) => {
                                             packet_stats.transport.tcp += 1;
@@ -128,9 +155,9 @@ impl Stats {
     }
     pub fn get_top_10(
         &self,
-        addresses: HashMap<Ipv4Addr, (Option<String>, usize)>,
-    ) -> Vec<(Ipv4Addr, (Option<String>, usize))> {
-        let mut items: Vec<(Ipv4Addr, (Option<String>, usize))> = addresses.into_iter().collect();
+        addresses: HashMap<IpAddr, (Option<String>, usize)>,
+    ) -> Vec<(IpAddr, (Option<String>, usize))> {
+        let mut items: Vec<(IpAddr, (Option<String>, usize))> = addresses.into_iter().collect();
         items.sort_by(|a, b| b.1 .1.cmp(&a.1 .1));
         items.into_iter().take(10).collect()
     }
@@ -302,7 +329,7 @@ impl Stats {
                     .title_alignment(Alignment::Center)
                     .padding(Padding::horizontal(1))
                     .padding(Padding::right(3))
-                    .title_bottom("Top visited websites (Ipv4 only)"),
+                    .title_bottom("Top visited websites"),
             );
 
         frame.render_widget(addresses_chart, address_block);