1212 feature automate pubkey based on callback uri

docs/_docs/user-guide/imix.md

@@ -16,7 +16,7 @@ Imix has compile-time configuration, that may be specified using environment var
 | Env Var | Description | Default | Required |
 | ------- | ----------- | ------- | -------- |
 | IMIX_CALLBACK_URI | URI for initial callbacks (must specify a scheme, e.g. `http://`) | `http://127.0.0.1:8000` | No |
-| IMIX_SERVER_PUBKEY | The public key for the tavern server (obtain from server using `curl $IMIX_CALLBACK_URI/status`). | - | Yes |
+| IMIX_SERVER_PUBKEY | The public key for the tavern server (obtain from server using `curl $IMIX_CALLBACK_URI/status`). | automatic | Yes |
 | IMIX_CALLBACK_INTERVAL | Duration between callbacks, in seconds. | `5` | No |
 | IMIX_RETRY_INTERVAL | Duration to wait before restarting the agent loop if an error occurs, in seconds. | `5` | No |
 | IMIX_PROXY_URI | Overide system settings for proxy URI over HTTP(S) (must specify a scheme, e.g. `https://`) | No proxy | No |
@@ -98,6 +98,15 @@ These flags are passed to cargo build Eg.:
 - `--features grpc-doh` - Enable DNS over HTTP using cloudflare DNS for the grpc transport
 - `--features http1 --no-default-features` - Changes the default grpc transport to use HTTP/1.1. Requires running the http redirector.
 
+## Setting encryption key
+
+By default imix will automatically collect the IMIX_CALLBACK_URI server's public key during the build process. This can be overridden by manually setinng the `IMIX_SERVER_PUBKEY` environment variable but should only be necesarry when using redirectors. Redirectors have no visibliity into the realm encryption by design, this means that agents must be compiled with the upstream tavern instance's public key.
+
+A server's public key can be found using:
+```bash
+export IMIX_SERVER_PUBKEY="$(curl $IMIX_CALLBACK_URI/status | jq -r '.Pubkey')"
+```
+
 ### Linux
 
 ```bash
@@ -107,9 +116,6 @@ sudo apt update
 sudo apt install musl-tools
 cd realm/implants/imix/
 export IMIX_CALLBACK_URI="http://localhost"
-# To get a servers pubkey:
-# curl $IMIX_CALLBACK_URI/status | jq -r '.Pubkey'
-export IMIX_SERVER_PUBKEY="<SERVER_PUBKEY>"
 
 cargo build --release --bin imix --target=x86_64-unknown-linux-musl
 ```
@@ -144,9 +150,7 @@ Modify .devcontainer/devcontainer.json by uncommenting the MacOSX.sdk mount. Thi
 cd realm/implants/imix/
 # Tell the linker to use the MacOSX.sdk
 export RUSTFLAGS="-Clink-arg=-isysroot -Clink-arg=/MacOSX.sdk -Clink-arg=-F/MacOSX.sdk/System/Library/Frameworks -Clink-arg=-L/MacOSX.sdk/usr/lib -Clink-arg=-lresolv"
-export IMIX_CALLBACK_URI="http://localhost"
-# To get a servers pubkey:
-# curl $IMIX_CALLBACK_URI/status | jq -r '.Pubkey'
+
 export IMIX_SERVER_PUBKEY="<SERVER_PUBKEY>"
 
 cargo zigbuild  --release --target aarch64-apple-darwin
@@ -160,9 +164,6 @@ cargo zigbuild  --release --target aarch64-apple-darwin
 cd realm/implants/imix/
 
 export IMIX_CALLBACK_URI="http://localhost"
-# To get a servers pubkey:
-# curl $IMIX_CALLBACK_URI/status | jq -r '.Pubkey'
-export IMIX_SERVER_PUBKEY="<SERVER_PUBKEY>"
 
 # Build imix.exe
  cargo build --release --target=x86_64-pc-windows-gnu

implants/lib/pb/Cargo.toml

@@ -32,3 +32,5 @@ whoami = { workspace = true }
 tonic-build = { workspace = true, features = ["prost"] }
 which = { workspace = true }
 home = "=0.5.11"
+reqwest = { workspace = true, features = ["blocking", "json", "rustls-tls"] }
+serde_json = { workspace = true }

implants/lib/pb/build.rs

@@ -2,7 +2,71 @@ use std::env;
 use std::path::PathBuf;
 use which::which;
 
+fn get_pub_key() {
+    // Check if IMIX_SERVER_PUBKEY is already set
+    if std::env::var("IMIX_SERVER_PUBKEY").is_ok() {
+        println!("cargo:warning=IMIX_SERVER_PUBKEY already set, skipping fetch");
+        return;
+    }
+
+    // Get the callback URI from environment variable, default to http://127.0.0.1:8000
+    let callback_uri =
+        std::env::var("IMIX_CALLBACK_URI").unwrap_or_else(|_| "http://127.0.0.1:8000".to_string());
+
+    // Construct the status endpoint URL
+    let status_url = format!("{}/status", callback_uri);
+
+    // Make a GET request to /status
+    let response = match reqwest::blocking::get(&status_url) {
+        Ok(resp) => resp,
+        Err(e) => {
+            println!("cargo:warning=Failed to connect to {}: {}", status_url, e);
+            return;
+        }
+    };
+
+    if !response.status().is_success() {
+        println!(
+            "cargo:warning=Failed to fetch status from {}: HTTP {}",
+            status_url,
+            response.status()
+        );
+        return;
+    }
+
+    let json = match response.json::<serde_json::Value>() {
+        Ok(json) => json,
+        Err(e) => {
+            println!(
+                "cargo:warning=Failed to parse JSON response from {}: {}",
+                status_url, e
+            );
+            return;
+        }
+    };
+
+    let pubkey = match json.get("Pubkey").and_then(|v| v.as_str()) {
+        Some(key) => key,
+        None => {
+            println!(
+                "cargo:warning=Pubkey field not found in response from {}",
+                status_url
+            );
+            return;
+        }
+    };
+
+    // Set the IMIX_SERVER_PUBKEY environment variable for the build
+    println!("cargo:rustc-env=IMIX_SERVER_PUBKEY={}", pubkey);
+    println!(
+        "cargo:warning=Successfully fetched server public key from {}",
+        status_url
+    );
+}
+
 fn main() -> Result<(), Box<dyn std::error::Error>> {
+    get_pub_key();
+
     // Skip if no `protoc` can be found
     match env::var_os("PROTOC")
         .map(PathBuf::from)