Add local_ip field to Connection/RawConnection struct

5 files changed, 46 insertions, 21 deletions

diff --git a/src/network/connection.rs b/src/network/connection.rs
index b4e0de2..a8a17cb 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::Ipv4Addr;
+use ::std::net::{Ipv4Addr, IpAddr};
 
 use ::std::net::SocketAddr;
 
@@ -29,17 +29,23 @@ impl fmt::Display for Protocol {
     }
 }
 
-#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash)]
+#[derive(Clone, Ord, PartialOrd, PartialEq, Eq, Hash, Debug)]
 pub struct Socket {
     pub ip: Ipv4Addr,
     pub port: u16,
 }
 
-#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord)]
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
+pub struct LocalSocket {
+    pub ip: IpAddr,
+    pub port: u16,
+    pub protocol: Protocol,
+}
+
+#[derive(PartialEq, Hash, Eq, Clone, PartialOrd, Ord, Debug)]
 pub struct Connection {
     pub remote_socket: Socket,
-    pub protocol: Protocol,
-    pub local_port: u16,
+    pub local_socket: LocalSocket,
 }
 
 pub fn display_ip_or_host(ip: Ipv4Addr, ip_to_host: &HashMap<Ipv4Addr, String>) -> String {
@@ -57,23 +63,26 @@ pub fn display_connection_string(
     format!(
         "<{}>:{} => {}:{} ({})",
         interface_name,
-        connection.local_port,
+        connection.local_socket.port,
         display_ip_or_host(connection.remote_socket.ip, ip_to_host),
         connection.remote_socket.port,
-        connection.protocol,
+        connection.local_socket.protocol,
     )
 }
 
 impl Connection {
-    pub fn new(remote_socket: SocketAddr, local_port: u16, protocol: Protocol) -> Option<Self> {
+    pub fn new(remote_socket: SocketAddr, 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(),
                 },
-                protocol,
-                local_port,
+                local_socket: LocalSocket {
+                    ip: local_ip,
+                    port: local_port,
+                    protocol,
+                },
             }),
             _ => None,
         }
diff --git a/src/network/sniffer.rs b/src/network/sniffer.rs
index 4c19731..b2f947f 100644
--- a/src/network/sniffer.rs
+++ b/src/network/sniffer.rs
@@ -88,8 +88,8 @@ impl Sniffer {
                 let to = SocketAddr::new(IpAddr::V4(ip_packet.get_destination()), destination_port);
 
                 let connection = match direction {
-                    Direction::Download => Connection::new(from, destination_port, protocol)?,
-                    Direction::Upload => Connection::new(to, 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_utils.rs b/src/os/lsof_utils.rs
index d1ca51b..62f4cce 100644
--- a/src/os/lsof_utils.rs
+++ b/src/os/lsof_utils.rs
@@ -7,7 +7,8 @@ use std::process::Command;
 
 #[derive(Debug, Clone)]
 pub struct RawConnection {
-    ip: String,
+    remote_ip: String,
+    local_ip: String,
     local_port: String,
     remote_port: String,
     protocol: String,
@@ -19,17 +20,21 @@ lazy_static! {
         Regex::new(r"([^\s]+).*(TCP|UDP).*:(.*)->(.*):(\d*)(\s|$)").unwrap();
 }
 
+#[allow(clippy::needless_return)]
 impl RawConnection {
     pub fn new(raw_line: &str) -> Option<RawConnection> {
         let raw_connection_iter = CONNECTION_REGEX.captures_iter(raw_line).filter_map(|cap| {
             let process_name = String::from(cap.get(1).unwrap().as_str()).replace("\\x20", " ");
             let protocol = String::from(cap.get(2).unwrap().as_str());
             let local_port = String::from(cap.get(3).unwrap().as_str());
-            let ip = String::from(cap.get(4).unwrap().as_str());
+            let remote_ip = String::from(cap.get(4).unwrap().as_str());
+            // TODO correctly parse local IP from lsof output
+            let local_ip = String::from("0.0.0.0");
             let remote_port = String::from(cap.get(5).unwrap().as_str());
             let connection = RawConnection {
                 process_name,
-                ip,
+                remote_ip,
+                local_ip,
                 local_port,
                 remote_port,
                 protocol,
@@ -48,14 +53,18 @@ impl RawConnection {
         return Protocol::from_str(&self.protocol).unwrap();
     }
 
-    pub fn get_ip_address(&self) -> IpAddr {
-        return IpAddr::V4(self.ip.parse().unwrap());
+    pub fn get_remote_ip(&self) -> IpAddr {
+        return IpAddr::V4(self.remote_ip.parse().unwrap());
     }
 
     pub fn get_remote_port(&self) -> u16 {
         return self.remote_port.parse::<u16>().unwrap();
     }
 
+    pub fn get_local_ip(&self) -> IpAddr {
+        return IpAddr::V4(self.local_ip.parse().unwrap());
+    }
+
     pub fn get_local_port(&self) -> u16 {
         return self.local_port.parse::<u16>().unwrap();
     }
@@ -160,7 +169,7 @@ com.apple   590 etoledom  204u  IPv4 0x28ffb9c04111253f      0t0  TCP 192.168.1.
     fn test_raw_connection_parse_ip_address() {
         let connection = RawConnection::new(LINE_RAW_OUTPUT).unwrap();
         assert_eq!(
-            connection.get_ip_address().to_string(),
+            connection.get_remote_ip().to_string(),
             String::from("198.252.206.25")
         );
     }
diff --git a/src/os/macos.rs b/src/os/macos.rs
index cf68d04..1ed6ecf 100644
--- a/src/os/macos.rs
+++ b/src/os/macos.rs
@@ -21,12 +21,13 @@ pub(crate) fn get_open_sockets() -> HashMap<Connection, String> {
 
     for raw_connection in connections {
         let protocol = raw_connection.get_protocol();
-        let ip_address = raw_connection.get_ip_address();
+        let remote_ip = raw_connection.get_remote_ip();
+        let local_ip = raw_connection.get_local_ip();
         let remote_port = raw_connection.get_remote_port();
         let local_port = raw_connection.get_local_port();
 
-        let socket_addr = SocketAddr::new(ip_address, remote_port);
-        let connection = Connection::new(socket_addr, local_port, protocol).unwrap();
+        let socket_addr = SocketAddr::new(remote_ip, remote_port);
+        let connection = Connection::new(socket_addr, local_ip, local_port, protocol).unwrap();
 
         open_sockets.insert(connection, raw_connection.process_name.clone());
     }
diff --git a/src/tests/fakes/fake_input.rs b/src/tests/fakes/fake_input.rs
index a3b84b3..7f384f6 100644
--- a/src/tests/fakes/fake_input.rs
+++ b/src/tests/fakes/fake_input.rs
@@ -87,9 +87,11 @@ impl DataLinkReceiver for NetworkFrames {
 
 pub fn get_open_sockets() -> HashMap<Connection, String> {
     let mut open_sockets = HashMap::new();
+    let local_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
     open_sockets.insert(
         Connection::new(
             SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 12345),
+            local_ip,
             443,
             Protocol::Tcp,
         )
@@ -99,6 +101,7 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
     open_sockets.insert(
         Connection::new(
             SocketAddr::new(IpAddr::V4(Ipv4Addr::new(2, 2, 2, 2)), 54321),
+            local_ip,
             443,
             Protocol::Tcp,
         )
@@ -108,6 +111,7 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
     open_sockets.insert(
         Connection::new(
             SocketAddr::new(IpAddr::V4(Ipv4Addr::new(3, 3, 3, 3)), 1337),
+            local_ip,
             443,
             Protocol::Tcp,
         )
@@ -117,6 +121,7 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
     open_sockets.insert(
         Connection::new(
             SocketAddr::new(IpAddr::V4(Ipv4Addr::new(4, 4, 4, 4)), 1337),
+            local_ip,
             443,
             Protocol::Tcp,
         )
@@ -126,6 +131,7 @@ pub fn get_open_sockets() -> HashMap<Connection, String> {
     open_sockets.insert(
         Connection::new(
             SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 12346),
+            local_ip,
             443,
             Protocol::Tcp,
         )