feat(localip): add module to print the current ipv4 address (#3289)

6 files changed, 159 insertions, 0 deletions

diff --git a/src/configs/localip.rs b/src/configs/localip.rs
new file mode 100644
index 000000000..98655afe7
--- /dev/null
+++ b/src/configs/localip.rs
@@ -0,0 +1,23 @@
+use crate::config::ModuleConfig;
+
+use serde::Serialize;
+use starship_module_config_derive::ModuleConfig;
+
+#[derive(Clone, ModuleConfig, Serialize)]
+pub struct LocalipConfig<'a> {
+    pub ssh_only: bool,
+    pub format: &'a str,
+    pub style: &'a str,
+    pub disabled: bool,
+}
+
+impl<'a> Default for LocalipConfig<'a> {
+    fn default() -> Self {
+        LocalipConfig {
+            ssh_only: true,
+            format: "[$localipv4]($style) ",
+            style: "yellow bold",
+            disabled: true,
+        }
+    }
+}
diff --git a/src/configs/mod.rs b/src/configs/mod.rs
index 252f9c7b2..02bdfb3dd 100644
--- a/src/configs/mod.rs
+++ b/src/configs/mod.rs
@@ -40,6 +40,7 @@ pub mod julia;
 pub mod kotlin;
 pub mod kubernetes;
 pub mod line_break;
+pub mod localip;
 pub mod lua;
 pub mod memory_usage;
 pub mod nim;
@@ -122,6 +123,7 @@ pub struct FullConfig<'a> {
     kotlin: kotlin::KotlinConfig<'a>,
     kubernetes: kubernetes::KubernetesConfig<'a>,
     line_break: line_break::LineBreakConfig,
+    localip: localip::LocalipConfig<'a>,
     lua: lua::LuaConfig<'a>,
     memory_usage: memory_usage::MemoryConfig<'a>,
     nim: nim::NimConfig<'a>,
@@ -202,6 +204,7 @@ impl<'a> Default for FullConfig<'a> {
             kotlin: Default::default(),
             kubernetes: Default::default(),
             line_break: Default::default(),
+            localip: Default::default(),
             lua: Default::default(),
             memory_usage: Default::default(),
             nim: Default::default(),
diff --git a/src/configs/starship_root.rs b/src/configs/starship_root.rs
index b5ea06af4..86f12798d 100644
--- a/src/configs/starship_root.rs
+++ b/src/configs/starship_root.rs
@@ -20,6 +20,7 @@ pub struct StarshipRootConfig {
 pub const PROMPT_ORDER: &[&str] = &[
     "username",
     "hostname",
+    "localip",
     "shlvl",
     "singularity",
     "kubernetes",
diff --git a/src/module.rs b/src/module.rs
index 925e569c0..65a1d67f7 100644
--- a/src/module.rs
+++ b/src/module.rs
@@ -45,6 +45,7 @@ pub const ALL_MODULES: &[&str] = &[
     "kotlin",
     "kubernetes",
     "line_break",
+    "localip",
     "lua",
     "memory_usage",
     "nim",
diff --git a/src/modules/localip.rs b/src/modules/localip.rs
new file mode 100644
index 000000000..8d02f0914
--- /dev/null
+++ b/src/modules/localip.rs
@@ -0,0 +1,128 @@
+use super::{Context, Module};
+
+use crate::config::RootModuleConfig;
+use crate::configs::localip::LocalipConfig;
+use crate::formatter::StringFormatter;
+
+/// Creates a module with the ipv4 address of the local machine.
+///
+/// The `local_ipaddress` crate is used to determine the local IP address of your machine.
+/// An accurate and fast way, especially if there are multiple IP addresses available,
+/// is to connect a UDP socket and then reading its local endpoint.
+///
+/// Will display the ip if all of the following criteria are met:
+///     - localip.disabled is absent or false
+///     - localip.ssh_only is false OR the user is currently connected as an SSH session (`$SSH_CONNECTION`)
+pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
+    let mut module = context.new_module("localip");
+    let config: LocalipConfig = LocalipConfig::try_load(module.config);
+
+    let ssh_connection = context.get_env("SSH_CONNECTION");
+    if config.ssh_only && ssh_connection.is_none() {
+        return None;
+    }
+
+    let localip = local_ipaddress::get().unwrap_or_default();
+    if localip.is_empty() {
+        log::warn!("unable to determine local ipv4 address");
+        return None;
+    }
+
+    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
+        formatter
+            .map_style(|variable| match variable {
+                "style" => Some(Ok(config.style)),
+                _ => None,
+            })
+            .map(|variable| match variable {
+                "localipv4" => Some(Ok(&localip)),
+                _ => None,
+            })
+            .parse(None, Some(context))
+    });
+
+    module.set_segments(match parsed {
+        Ok(segments) => segments,
+        Err(error) => {
+            log::warn!("Error in module `localip`:\n{}", error);
+            return None;
+        }
+    });
+
+    Some(module)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::test::ModuleRenderer;
+    use ansi_term::{Color, Style};
+
+    macro_rules! get_localip {
+        () => {
+            if let Some(localip) = local_ipaddress::get() {
+                localip
+            } else {
+                println!(
+                    "localip was not tested because socket connection failed! \
+                     This could be caused by an unconventional network setup."
+                );
+                return;
+            }
+        };
+    }
+
+    #[test]
+    fn is_ipv4_format() {
+        let localip = get_localip!();
+        assert!(regex::Regex::new(r"^(?:[0-9]{1,3}\.){3}[0-9]{1,3}$")
+            .unwrap()
+            .is_match(&localip));
+    }
+
+    #[test]
+    fn ssh_only_false() {
+        let localip = get_localip!();
+        let actual = ModuleRenderer::new("localip")
+            .config(toml::toml! {
+                [localip]
+                ssh_only = false
+            })
+            .collect();
+        let expected = Some(format!("{} ", style().paint(localip)));
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn no_ssh() {
+        let actual = ModuleRenderer::new("localip")
+            .config(toml::toml! {
+                [localip]
+                ssh_only = true
+            })
+            .collect();
+        let expected = None;
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn ssh() {
+        let localip = get_localip!();
+        let actual = ModuleRenderer::new("localip")
+            .config(toml::toml! {
+                [localip]
+                ssh_only = true
+                trim_at = ""
+            })
+            .env("SSH_CONNECTION", "something")
+            .collect();
+        let expected = Some(format!("{} ", style().paint(localip)));
+
+        assert_eq!(expected, actual);
+    }
+
+    fn style() -> Style {
+        Color::Yellow.bold()
+    }
+}
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 91333ccb4..da724a059 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -35,6 +35,7 @@ mod julia;
 mod kotlin;
 mod kubernetes;
 mod line_break;
+mod localip;
 mod lua;
 mod memory_usage;
 mod nim;
@@ -121,6 +122,7 @@ pub fn handle<'a>(module: &str, context: &'a Context) -> Option<Module<'a>> {
             "kotlin" => kotlin::module(context),
             "kubernetes" => kubernetes::module(context),
             "line_break" => line_break::module(context),
+            "localip" => localip::module(context),
             "lua" => lua::module(context),
             "memory_usage" => memory_usage::module(context),
             "nim" => nim::module(context),
@@ -212,6 +214,7 @@ pub fn description(module: &str) -> &'static str {
         "kotlin" => "The currently installed version of Kotlin",
         "kubernetes" => "The current Kubernetes context name and, if set, the namespace",
         "line_break" => "Separates the prompt into two lines",
+        "localip" => "The currently assigned ipv4 address",
         "lua" => "The currently installed version of Lua",
         "memory_usage" => "Current system memory and swap usage",
         "nim" => "The currently installed version of Nim",