diff options
-rw-r--r-- | Cargo.lock | 24 | ||||
-rw-r--r-- | Cargo.toml | 1 | ||||
-rw-r--r-- | src/display/ui_state.rs | 7 | ||||
-rw-r--r-- | src/lib.rs | 67 | ||||
-rw-r--r-- | src/main.rs | 4 | ||||
-rw-r--r-- | src/network/connection.rs | 92 | ||||
-rw-r--r-- | src/network/dns_queue.rs | 36 | ||||
-rw-r--r-- | src/network/mod.rs | 4 | ||||
-rw-r--r-- | src/network/resolve_connections.rs | 26 | ||||
-rw-r--r-- | src/os/linux.rs | 8 | ||||
-rw-r--r-- | tests/cli.rs | 128 | ||||
-rw-r--r-- | tests/fakes/fake_input.rs | 33 | ||||
-rw-r--r-- | tests/snapshots/cli__traffic_with_host_names-2.snap | 55 | ||||
-rw-r--r-- | tests/snapshots/cli__traffic_with_host_names.snap | 55 |
14 files changed, 479 insertions, 61 deletions
@@ -221,6 +221,16 @@ dependencies = [ ] [[package]] +name = "dns-lookup" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "libc 0.2.60 (registry+https://github.com/rust-lang/crates.io-index)", + "socket2 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)", + "winapi 0.3.7 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] name = "dtoa" version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -897,6 +907,17 @@ version = "0.6.10" source = "registry+https://github.com/rust-lang/crates.io-index" [[package]] +name = "socket2" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "cfg-if 0.1.9 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.60 (registry+https://github.com/rust-lang/crates.io-index)", + "redox_syscall 0.1.56 (registry+https://github.com/rust-lang/crates.io-index)", + "winapi 0.3.7 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] name = "strsim" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1161,6 +1182,7 @@ version = "0.1.0" dependencies = [ "assert_cmd 0.10.2 (registry+https://github.com/rust-lang/crates.io-index)", "cargo-insta 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)", + "dns-lookup 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)", "insta 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)", "ipnetwork 0.14.0 (registry+https://github.com/rust-lang/crates.io-index)", "packet-builder 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)", @@ -1252,6 +1274,7 @@ dependencies = [ "checksum derive-new 0.5.8 (registry+https://github.com/rust-lang/crates.io-index)" = "71f31892cd5c62e414316f2963c5689242c43d8e7bbcaaeca97e5e28c95d91d9" "checksum difference 2.0.0 (registry+https://github.com/rust-lang/crates.io-index)" = "524cbf6897b527295dff137cec09ecf3a05f4fddffd7dfcd1585403449e74198" "checksum digest 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "f3d0c8c8752312f9713efd397ff63acb9f85585afbf179282e720e7704954dd5" +"checksum dns-lookup 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)" = "13988670860b076248c74e1b54444efc4f1dec70c8bb25da4b7c0024396b72bf" "checksum dtoa 0.4.4 (registry+https://github.com/rust-lang/crates.io-index)" = "ea57b42383d091c85abcc2706240b94ab2a8fa1fc81c10ff23c4de06e2a90b5e" "checksum either 1.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "5527cfe0d098f36e3f8839852688e63c8fff1c90b2b405aef730615f9a7bcf7b" "checksum encode_unicode 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)" = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" @@ -1334,6 +1357,7 @@ dependencies = [ "checksum serde_yaml 0.8.9 (registry+https://github.com/rust-lang/crates.io-index)" = "38b08a9a90e5260fe01c6480ec7c811606df6d3a660415808c3c3fa8ed95b582" "checksum sha-1 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "23962131a91661d643c98940b20fcaffe62d776a823247be80a48fcb8b6fce68" "checksum smallvec 0.6.10 (registry+https://github.com/rust-lang/crates.io-index)" = "ab606a9c5e214920bb66c458cd7be8ef094f813f20fe77a54cc7dbfff220d4b7" +"checksum socket2 0.3.11 (registry+https://github.com/rust-lang/crates.io-index)" = "e8b74de517221a2cb01a53349cf54182acdc31a074727d3079068448c0676d85" "checksum strsim 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)" = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" "checksum structopt 0.2.18 (registry+https://github.com/rust-lang/crates.io-index)" = "16c2cdbf9cc375f15d1b4141bc48aeef444806655cd0e904207edc8d68d86ed7" "checksum structopt 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)" = "2ac9d6e93dd792b217bf89cda5c14566e3043960c6f9da890c2ba5d09d07804c" @@ -11,6 +11,7 @@ ipnetwork = "0.14.0" tui = "0.5" termion = "1.5" structopt = "0.3" +dns-lookup = "*" [dev-dependencies] assert_cmd = "0.10" diff --git a/src/display/ui_state.rs b/src/display/ui_state.rs index b2be517..1d3233b 100644 --- a/src/display/ui_state.rs +++ b/src/display/ui_state.rs @@ -1,5 +1,4 @@ use ::std::collections::{BTreeMap, HashMap}; -use ::std::net::Ipv4Addr; use crate::network::{Connection, Utilization}; @@ -42,7 +41,7 @@ impl Bandwidth for NetworkData { pub struct UIState { pub processes: BTreeMap<String, NetworkData>, - pub remote_ips: BTreeMap<Ipv4Addr, NetworkData>, + pub remote_ips: BTreeMap<String, NetworkData>, pub connections: BTreeMap<Connection, ConnectionData>, } @@ -52,14 +51,14 @@ impl UIState { network_utilization: &Utilization, ) -> Self { let mut processes: BTreeMap<String, NetworkData> = BTreeMap::new(); - let mut remote_ips: BTreeMap<Ipv4Addr, NetworkData> = BTreeMap::new(); + let mut remote_ips: BTreeMap<String, NetworkData> = BTreeMap::new(); let mut connections: BTreeMap<Connection, ConnectionData> = BTreeMap::new(); for (connection, process_name) in connections_to_procs { if let Some(connection_bandwidth_utilization) = network_utilization.connections.get(&connection) { let data_for_remote_ip = remote_ips - .entry(connection.remote_socket.ip.clone()) + .entry(connection.remote_socket.clone_host_or_ip()) .or_default(); let connection_data = connections.entry(connection).or_default(); let data_for_process = processes.entry(process_name.clone()).or_default(); @@ -2,7 +2,9 @@ mod display; pub mod network; use display::display_loop; -use network::{Connection, Sniffer, Utilization}; +use network::{Connection, Sniffer, Utilization, DnsQueue, resolve_connections}; + +use ::std::net::IpAddr; use ::pnet::datalink::{DataLinkReceiver, NetworkInterface}; use ::std::collections::HashMap; @@ -18,43 +20,74 @@ pub struct OsInput { pub network_frames: Box<DataLinkReceiver>, pub get_open_sockets: fn() -> HashMap<Connection, String>, pub keyboard_events: Box<Iterator<Item = Event> + Send + Sync + 'static>, + pub lookup_addr: Box<Fn(&IpAddr) -> Option<String> + Send + Sync + 'static> } pub fn start<B>(terminal_backend: B, os_input: OsInput) where B: Backend + Send + 'static, { - let r = Arc::new(AtomicBool::new(true)); - let displaying = r.clone(); - let running = r.clone(); + let running = Arc::new(AtomicBool::new(true)); let keyboard_events = os_input.keyboard_events; // TODO: as methods in os_interface let get_open_sockets = os_input.get_open_sockets; + let lookup_addr = os_input.lookup_addr; - let stdin_handler = thread::spawn(move || { - for evt in keyboard_events { - match evt { - Event::Key(Key::Ctrl('c')) | Event::Key(Key::Char('q')) => { - // TODO: exit faster - r.store(false, Ordering::Relaxed); - break; - } - _ => (), - }; + let stdin_handler = thread::spawn({ + let running = running.clone(); + move || { + for evt in keyboard_events { + match evt { + Event::Key(Key::Ctrl('c')) | Event::Key(Key::Char('q')) => { + // TODO: exit faster + running.store(false, Ordering::Relaxed); + break; + } + _ => (), + }; + } } }); let mut sniffer = Sniffer::new(os_input.network_interface, os_input.network_frames); let network_utilization = Arc::new(Mutex::new(Utilization::new())); + let dns_queue = Arc::new(DnsQueue::new()); + let ip_to_host = Arc::new(Mutex::new(HashMap::new())); + + let dns_handler = thread::spawn({ + let running = running.clone(); + let dns_queue = dns_queue.clone(); + let ip_to_host = ip_to_host.clone(); + move || { + while running.load(Ordering::Relaxed) { + let jobs = dns_queue.wait_for_jobs(); + for ip in jobs { + if let Some(addr) = lookup_addr(&IpAddr::V4(ip.clone())) { + ip_to_host.lock().unwrap().insert(ip, addr); + } + } + } + } + }); + let display_handler = thread::spawn({ + let running = running.clone(); let network_utilization = network_utilization.clone(); + let ip_to_host = ip_to_host.clone(); + let dns_queue = dns_queue.clone(); move || { let mut terminal = Terminal::new(terminal_backend).unwrap(); terminal.clear().unwrap(); terminal.hide_cursor().unwrap(); - while displaying.load(Ordering::Relaxed) { - let connections_to_procs = get_open_sockets(); + while running.load(Ordering::Relaxed) { + let connections_to_procs = { + let open_sockets = get_open_sockets(); + let ip_to_host = ip_to_host.lock().unwrap(); + let (unresolved_ips, connections_to_procs) = resolve_connections(open_sockets, &ip_to_host); + dns_queue.add_ips_to_resolve(unresolved_ips); + connections_to_procs + }; { let mut network_utilization = network_utilization.lock().unwrap(); let utilization = network_utilization.clone_and_reset(); @@ -64,6 +97,7 @@ where } terminal.clear().unwrap(); terminal.show_cursor().unwrap(); + dns_queue.end(); } }); @@ -77,4 +111,5 @@ where display_handler.join().unwrap(); sniffing_handler.join().unwrap(); stdin_handler.join().unwrap(); + dns_handler.join().unwrap(); } diff --git a/src/main.rs b/src/main.rs index faab23c..c85b555 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,7 +19,7 @@ fn main() { "Sorry, no implementations for platforms other than linux yet :( - PRs welcome!" ); - use os::{get_datalink_channel, get_interface, get_open_sockets, KeyboardEvents}; + use os::{get_datalink_channel, get_interface, get_open_sockets, KeyboardEvents, lookup_addr}; let opt = Opt::from_args(); let stdout = io::stdout().into_raw_mode().unwrap(); @@ -28,12 +28,14 @@ fn main() { let keyboard_events = Box::new(KeyboardEvents); let network_interface = get_interface(&opt.interface).unwrap(); let network_frames = get_datalink_channel(&network_interface); + let lookup_addr = Box::new(lookup_addr); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(terminal_backend, os_input) diff --git a/src/network/connection.rs b/src/network/connection.rs index 8676ec1..34194c1 100644 --- a/src/network/connection.rs +++ b/src/network/connection.rs @@ -4,7 +4,10 @@ use ::std::net::Ipv4Addr; use ::std::mem::swap; use ::std::net::SocketAddr; -#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)] +use std::cmp::Ordering; +use std::hash::{Hash, Hasher}; + +#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)] pub enum Protocol { Tcp, Udp, @@ -19,19 +22,67 @@ impl fmt::Display for Protocol { } } -#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)] +#[derive(Clone)] pub struct Socket { pub ip: Ipv4Addr, pub port: u16, + host_addr: Option<String> +} + +impl Socket { + pub fn clone_host_or_ip(&self) -> String { + match &self.host_addr { + Some(host_addr) => host_addr.clone(), + None => self.ip.to_string() + } + } +} + +impl Ord for Socket { + fn cmp(&self, other: &Self) -> Ordering { + let ip_eq = self.ip.cmp(&other.ip); // TODO: also port + match ip_eq { + Ordering::Equal => self.port.cmp(&other.port), + _ => ip_eq + } + } +} + +impl PartialOrd for Socket { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + Some(self.cmp(other)) + } +} + +impl PartialEq for Socket { + fn eq(&self, other: &Self) -> bool { + self.ip == other.ip && self.port == other.port + } +} + +impl Eq for Socket {} + +impl Hash for Socket { + fn hash<H: Hasher>(&self, state: &mut H) { + self.ip.hash(state); + self.port.hash(state); + } } impl fmt::Display for Socket { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}:{}", self.ip, self.port) + match &self.host_addr { + Some(host_addr) => { + write!(f, "{}:{}", host_addr, self.port) + }, + None => { + write!(f, "{}:{}", self.ip, self.port) + } + } } } -#[derive(PartialEq, Hash, Eq, Debug, Clone, PartialOrd, Ord)] +#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)] pub struct Connection { pub local_socket: Socket, pub remote_socket: Socket, @@ -55,22 +106,31 @@ impl Connection { protocol: Protocol, ) -> Option<Self> { match (local_socket, remote_socket) { - (SocketAddr::V4(local_socket), SocketAddr::V4(remote_socket)) => Some(Connection { - // we use our own Socket here because SocketAddr is not sorable - local_socket: Socket { - ip: *local_socket.ip(), - port: local_socket.port(), - }, - remote_socket: Socket { - ip: *remote_socket.ip(), - port: remote_socket.port(), - }, - protocol, - }), + (SocketAddr::V4(local_socket), SocketAddr::V4(remote_socket)) => { + Some(Connection { + local_socket: Socket { + ip: *local_socket.ip(), + port: local_socket.port(), + host_addr: None, + }, + remote_socket: Socket { + ip: *remote_socket.ip(), + port: remote_socket.port(), + host_addr: None, + }, + protocol, + }) + }, (_, _) => None, } } pub fn swap_direction(&mut self) { swap(&mut self.local_socket, &mut self.remote_socket); } + pub fn set_local_host_addr(&mut self, addr: &str) { + self.local_socket.host_addr = Some(String::from(addr)); + } + pub fn set_remote_host_addr(&mut self, addr: &str) { + self.remote_socket.host_addr = Some(String::from(addr)); + } } diff --git a/src/network/dns_queue.rs b/src/network/dns_queue.rs new file mode 100644 index 0000000..93728e8 --- /dev/null +++ b/src/network/dns_queue.rs @@ -0,0 +1,36 @@ + +use ::std::net::Ipv4Addr; + +use ::std::mem::swap; +use ::std::sync::{Mutex, Condvar}; + +pub struct DnsQueue { + jobs: Mutex<Vec<Ipv4Addr>>, + cvar: Condvar +} + +impl DnsQueue { + pub fn new() -> Self { + DnsQueue { + jobs: Mutex::new(Vec::new()), + cvar: Condvar::new() + } + } +} + +impl DnsQueue { + pub fn add_ips_to_resolve(&self, unresolved_ips: Vec<Ipv4Addr>) { + let mut queue = self.jobs.lock().unwrap(); + queue.append(&mut unresolved_ips.into_iter().collect()); + self.cvar.notify_all(); + } + pub fn wait_for_jobs(&self) -> Vec<Ipv4Addr> { + let mut jobs = self.cvar.wait(self.jobs.lock().unwrap()).unwrap(); + let mut new_jobs = Vec::new(); + swap(&mut new_jobs, &mut jobs); + new_jobs + } + pub fn end(&self) { + self.cvar.notify_all(); + } +} diff --git a/src/network/mod.rs b/src/network/mod.rs index e8e3207..3bb54d7 100644 --- a/src/network/mod.rs +++ b/src/network/mod.rs @@ -1,7 +1,11 @@ mod connection; mod sniffer; mod utilization; +mod dns_queue; +mod resolve_connections; pub use connection::*; pub use sniffer::*; pub use utilization::*; +pub use dns_queue::*; +pub use resolve_connections::*; diff --git a/src/network/resolve_connections.rs b/src/network/resolve_connections.rs new file mode 100644 index 0000000..ee70ade --- /dev/null +++ b/src/network/resolve_connections.rs @@ -0,0 +1,26 @@ +use crate::Connection; + +use ::std::net::Ipv4Addr; +use ::std::collections::HashMap; + +pub fn resolve_connections ( + open_sockets: HashMap<Connection, String>, + ip_to_host: &HashMap<Ipv4Addr, String> +) -> (Vec<Ipv4Addr>, HashMap<Connection, String>) { + let mut unresolved_ips = vec![]; + let mut resolved_connections_to_procs: HashMap<Connection, String> = HashMap::new(); + for connection in open_sockets.keys() { + let mut connection = connection.clone(); + match ip_to_host.get(&connection.local_socket.ip) { + Some(local_host_addr) => connection.set_local_host_addr(local_host_addr), + None => unresolved_ips.push(connection.local_socket.ip.clone()), + } + match ip_to_host.get(&connection.remote_socket.ip) { + Some(remote_host_addr) => connection.set_remote_host_addr(remote_host_addr), + None => unresolved_ips.push(connection.remote_socket.ip.clone()), + } + let connection_value = open_sockets.get(&connection).unwrap().clone(); + &resolved_connections_to_procs.insert(connection, connection_value); + } + (unresolved_ips, resolved_connections_to_procs) +} diff --git a/src/os/linux.rs b/src/os/linux.rs index aa31452..e14a351 100644 --- a/src/os/linux.rs +++ b/src/os/linux.rs @@ -6,6 +6,7 @@ use ::termion::event::Event; use ::termion::input::TermRead; use ::std::collections::HashMap; +use ::std::net::IpAddr; use ::procfs::FDTarget; @@ -77,3 +78,10 @@ pub fn get_open_sockets() -> HashMap<Connection, String> { } open_sockets } + +pub fn lookup_addr(ip: &IpAddr) -> Option<String> { + if let Ok(addr) = ::dns_lookup::lookup_addr(ip) { + return Some(addr) + } + None +} diff --git a/tests/cli.rs b/tests/cli.rs index 8d2384c..83f298b 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -1,12 +1,15 @@ mod fakes; use fakes::TerminalEvent::*; -use fakes::{get_interface, get_open_sockets, KeyboardEvents, NetworkFrames, TestBackend}; +use fakes::{get_interface, get_open_sockets, KeyboardEvents, NetworkFrames, TestBackend, create_fake_lookup_addr}; use ::insta::assert_snapshot; use ::std::sync::{Arc, Mutex}; use ::termion::event::{Event, Key}; +use ::std::collections::HashMap; +use ::std::net::IpAddr; + use packet_builder::payload::PayloadData; use packet_builder::*; use pnet::packet::Packet; @@ -58,12 +61,14 @@ fn basic_startup() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -97,12 +102,14 @@ fn one_packet_of_traffic() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -148,12 +155,14 @@ fn bi_directional_traffic() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -199,12 +208,14 @@ fn multiple_packets_of_traffic_from_different_connections() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -250,12 +261,14 @@ fn multiple_packets_of_traffic_from_single_connection() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -301,12 +314,14 @@ fn one_process_with_multiple_connections() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -366,12 +381,14 @@ fn multiple_processes_with_multiple_connections() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -417,12 +434,14 @@ fn multiple_connections_from_remote_ip() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -470,12 +489,14 @@ fn sustained_traffic_from_one_process() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -537,12 +558,14 @@ fn sustained_traffic_from_multiple_processes() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); @@ -632,12 +655,115 @@ fn sustained_traffic_from_multiple_processes_bi_directional() { let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); let network_interface = get_interface(); + let lookup_addr = create_fake_lookup_addr(HashMap::new()); + + let os_input = what::OsInput { + network_interface, + network_frames, + get_open_sockets, + keyboard_events, + lookup_addr + }; + what::start(backend, os_input); + + let terminal_events_mirror = terminal_events.mirror.lock().unwrap(); + let terminal_draw_events_mirror = terminal_draw_events.mirror.lock().unwrap(); + + let expected_terminal_events = vec![ + Clear, HideCursor, Draw, Flush, Draw, Flush, Draw, Flush, Clear, ShowCursor, + ]; + assert_eq!(&terminal_events_mirror[..], &expected_terminal_events[..]); + + assert_eq!(terminal_draw_events_mirror.len(), 3); + assert_snapshot!(&terminal_draw_events_mirror[1]); + assert_snapshot!(&terminal_draw_events_mirror[2]); +} + +#[test] +fn traffic_with_host_names() { + let keyboard_events = Box::new(KeyboardEvents::new(vec![ + None, // sleep + None, // sleep + None, // sleep + Some(Event::Key(Key::Ctrl('c'))), + ])); + let network_frames = NetworkFrames::new(vec![ + Some(build_tcp_packet( + "10.0.0.2", + "3.3.3.3", + 443, + 1337, + b"omw to 3.3.3.3", + )), + Some(build_tcp_packet( + "3.3.3.3", + "10.0.0.2", + 1337, + 443, + b"I was just there!", + )), + Some(build_tcp_packet( + "1.1.1.1", + "10.0.0.2", + 12345, + 443, + b"Is it nice there? I think 1.1.1.1 is dull", + )), + Some(build_tcp_packet( + "10.0.0.2", + "1.1.1.1", + 443, + 12345, + b"Well, I heard 1.1.1.1 is all the rage", + )), + None, // sleep + Some(build_tcp_packet( + "10.0.0.2", + "3.3.3.3", + 443, + 1337, + b"Wait for me!", + )), + Some(build_tcp_packet( + "3.3.3.3", + "10.0.0.2", + 1337, + 443, + b"They're waiting for you...", + )), + Some(build_tcp_packet( + "1.1.1.1", + "10.0.0.2", + 12345, + 443, + b"1.1.1.1 forever!", + )), + Some(build_tcp_packet( + "10.0.0.2", + "1.1.1.1", + 443, + 12345, + b"10.0.0.2 forever!", + )), + ]); + + let terminal_events = LogWithMirror::new(Vec::new()); + let terminal_draw_events = LogWithMirror::new(Vec::new()); + + let backend = TestBackend::new(terminal_events.write, terminal_draw_events.write); + let network_interface = get_interface(); + let mut ips_to_hostnames = HashMap::new(); + ips_to_hostnames.insert(IpAddr::V4("1.1.1.1".parse().unwrap()), String::from("one.one.one.one")); + ips_to_hostnames.insert(IpAddr::V4("3.3.3.3".parse().unwrap()), String::from("three.three.three.three")); + ips_to_hostnames.insert(IpAddr::V4("10.0.0.2".parse().unwrap()), String::from("i-like-cheese.com")); + let lookup_addr = create_fake_lookup_addr(ips_to_hostnames); let os_input = what::OsInput { network_interface, network_frames, get_open_sockets, keyboard_events, + lookup_addr }; what::start(backend, os_input); diff --git a/tests/fakes/fake_input.rs b/tests/fakes/fake_input.rs index 02df931..9ab1624 100644 --- a/tests/fakes/fake_input.rs +++ b/tests/fakes/fake_input.rs @@ -2,8 +2,7 @@ use ::ipnetwork::IpNetwork; use ::pnet::datalink::DataLinkReceiver; use ::pnet::datalink::NetworkInterface; use ::std::collections::HashMap; -use ::std::net::IpAddr; -use ::std::net::{Ipv4Addr, SocketAddr}; +use ::std::net::{IpAddr, Ipv4Addr, SocketAddr}; use ::std::{thread, time}; use ::termion::event::Event; @@ -83,27 +82,6 @@ impl DataLinkReceiver for NetworkFrames { } } -// fn create_fake_socket( -// associated_pids: Vec<u32>, -// local_ip: IpAddr, -// remote_ip: IpAddr, -// local_port: u16, -// remote_port: u16, -// ) -> SocketInfo { -// let protocol_socket_info = TcpSocketInfo { -// local_addr: local_ip, -// remote_addr: remote_ip, -// local_port: local_port, -// remote_port: remote_port, -// state: TcpState::Listen, -// }; -// SocketInfo { -// protocol_socket_info: ProtocolSocketInfo::Tcp(protocol_socket_info), -// associated_pids: associated_pids, -// inode: 2, -// } -// } - pub fn get_open_sockets() -> HashMap<Connection, String> { let mut open_sockets = HashMap::new(); open_sockets.insert( @@ -164,3 +142,12 @@ pub fn get_interface() -> NetworkInterface { }; interface } + +pub fn create_fake_lookup_addr(ips_to_hosts: HashMap<IpAddr, String>) -> Box<Fn(&IpAddr) -> Option<String> + Send + Sync + 'static> { + Box::new(move |ip| { + match ips_to_hosts.get(ip) { + Some(host) => Some(host.clone()), + None => None + } + }) +} diff --git a/tests/snapshots/cli__traffic_with_host_names-2.snap b/tests/snapshots/cli__traffic_with_host_names-2.snap new file mode 100644 index 0000000..9c963d1 --- /dev/null +++ b/tests/snapshots/cli__traffic_with_host_names-2.snap @@ -0,0 +1,55 @@ +--- +source: tests/cli.rs +expression: "&terminal_draw_events_mirror[2]" +--- + + + + 1 16 1 16 + 2 26 2 26 + + + + + + + + + |