From 91052d125a6727e9f46c199d0b2c497b89065a1d Mon Sep 17 00:00:00 2001
From: Nicholas O'Brien <ndo9903@rit.edu>
Date: Thu, 14 Mar 2024 18:34:14 -0400
Subject: [PATCH] Get Agent IP every callback (#711)

* added ip to each callback

* move ip refresh to own function

* cleaning

* comment

* oops
---
 .vscode/settings.json       |  1 +
 implants/imix/src/agent.rs  |  5 +++++
 implants/imix/src/config.rs | 22 ++++++++++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 75e520d59..4fad2d189 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -4,4 +4,5 @@
         "rust-analyzer"
     ],
     "rust-analyzer.check.command": "clippy",
+    "rust-analyzer.showUnlinkedFileNotification": false,
 }
diff --git a/implants/imix/src/agent.rs b/implants/imix/src/agent.rs
index de64ba1cc..6f779903e 100644
--- a/implants/imix/src/agent.rs
+++ b/implants/imix/src/agent.rs
@@ -91,6 +91,11 @@ impl Agent {
         loop {
             let start = Instant::now();
 
+            // Sometimes Imix starts too quickly in a boot sequence, a NIC is down during the initial callback,
+            // or the box Imix is on changes its IP. In any case, for each callback we should refresh our claimed
+            // IP.
+            self.cfg.refresh_primary_ip();
+
             match self.callback().await {
                 Ok(_) => {}
                 Err(_err) => {
diff --git a/implants/imix/src/config.rs b/implants/imix/src/config.rs
index 4f972d3c9..ae9f3c299 100644
--- a/implants/imix/src/config.rs
+++ b/implants/imix/src/config.rs
@@ -107,6 +107,28 @@ impl Default for Config {
     }
 }
 
+impl Config {
+    pub fn refresh_primary_ip(&mut self) {
+        let fresh_ip = get_primary_ip();
+        if self
+            .info
+            .host
+            .as_ref()
+            .is_some_and(|h| h.primary_ip != fresh_ip)
+        {
+            match self.info.host.as_mut() {
+                Some(h) => {
+                    h.primary_ip = fresh_ip;
+                }
+                None => {
+                    #[cfg(debug_assertions)]
+                    log::error!("host struct was never initialized, failed to set primary ip");
+                }
+            }
+        }
+    }
+}
+
 /*
  * Returns which Platform imix has been compiled for.
  */