feat: ✨ Remote Code Execution

.gitignore

@@ -12,3 +12,8 @@ Cargo.lock
 
 # MSVC Windows builds of rustc generate these, which store debugging information
 *.pdb
+
+
+# Added by cargo
+
+/target

Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "rust-remote"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+futures-util = "0.3.30"
+serde = { version = "1.0.204", features = ["derive"] }
+serde_json = "1.0.120"
+tokio = { version = "1.38.0", features = ["full"] }
+tokio-tungstenite = "0.23.1"

configs/config.txt

@@ -0,0 +1,2 @@
+server_address:127.0.0.1
+port:2444

src/client.rs

@@ -0,0 +1,106 @@
+use std::process::Output;
+
+use futures_util::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
+use tokio::process::Command;
+use tokio_tungstenite::{connect_async, tungstenite::Message, WebSocketStream};
+
+use crate::{Config, Payload, Report};
+
+type WebSocketSender =
+    SplitSink<WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>, Message>;
+type WebSocketReceiver =
+    SplitStream<WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>>;
+
+pub async fn start(config: Config) {
+    let (ws_sender, ws_receiver) = match connect(config).await {
+        Some((ws_sender, ws_receiver)) => (ws_sender, ws_receiver),
+        None => return,
+    };
+
+    serve((ws_sender, ws_receiver)).await;
+}
+
+pub async fn connect(config: Config) -> Option<(WebSocketSender, WebSocketReceiver)> {
+    let ws_connection =
+        match connect_async(format!("ws://{}:{}", config.server_address, config.port)).await {
+            Ok(ws_connection) => ws_connection,
+            Err(err_val) => {
+                eprintln!("Error: WebSocket Connection | {}", err_val);
+                return None;
+            }
+        };
+
+    let (ws_sender, ws_receiver) = ws_connection.0.split();
+    Some((ws_sender, ws_receiver))
+}
+
+pub async fn serve((mut ws_sender, mut ws_receiver): (WebSocketSender, WebSocketReceiver)) -> ! {
+    loop {
+        match receive(&mut ws_receiver).await {
+            Some(message) => {
+                match serde_json::from_str(&message[..]) {
+                    Ok(payload) => match execute(payload).await {
+                        Some(output) => send(output, &mut ws_sender).await,
+                        None => todo!(),
+                    },
+                    Err(err_val) => {
+                        eprintln!("Error: Message to Payload | {}", err_val);
+                        continue;
+                    }
+                };
+            }
+            None => continue,
+        }
+    }
+}
+
+pub async fn execute(payload: Payload) -> Option<Output> {
+    println!("{:#?}", payload);
+    match Command::new(payload.command)
+        .args(payload.args)
+        .output()
+        .await
+    {
+        Ok(output) => Some(output),
+        Err(err_val) => {
+            eprintln!("Error: Command Execution | {}", err_val);
+            return None;
+        }
+    }
+}
+
+pub async fn receive(ws_receiver: &mut WebSocketReceiver) -> Option<String> {
+    match ws_receiver.next().await {
+        Some(message) => match message {
+            Ok(message) => {
+                if let Message::Text(message) = message {
+                    Some(message)
+                } else {
+                    eprintln!("Error: Message Type | {:#?}", message);
+                    None
+                }
+            }
+            Err(err_val) => {
+                eprintln!("Error: Message | {}", err_val);
+                None
+            }
+        },
+        None => {
+            eprintln!("Error: WebSocket Receive");
+            None
+        }
+    }
+}
+
+pub async fn send(output: Output, ws_sender: &mut WebSocketSender) {
+    let report = Report {
+        status: output.status.to_string(),
+        stdout: String::from_utf8(output.stdout).unwrap(),
+        stderr: String::from_utf8(output.stderr).unwrap(),
+    };
+    let report = serde_json::json!(report);
+    let _ = ws_sender.send(report.to_string().into()).await;
+}

src/lib.rs

@@ -0,0 +1,32 @@
+use std::net::IpAddr;
+
+use serde::{Deserialize, Serialize};
+
+pub mod client;
+pub mod server;
+pub mod utils;
+
+pub enum Runner {
+    Server,
+    Client,
+}
+
+#[derive(Debug)]
+pub struct Config {
+    pub server_address: IpAddr,
+    pub port: u16,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Payload {
+    pub sudo: bool,
+    pub command: String,
+    pub args: Vec<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Report {
+    pub status: String,
+    pub stdout: String,
+    pub stderr: String,
+}

src/main.rs

@@ -0,0 +1,28 @@
+use rust_remote::{
+    client, server,
+    utils::{read_config, take_args},
+};
+
+#[tokio::main]
+async fn main() {
+    println!("Hello, world!");
+
+    let config = match read_config() {
+        Some(config) => config,
+        None => {
+            eprintln!("Error: Read Config");
+            return;
+        }
+    };
+
+    match take_args() {
+        Some(runner) => match runner {
+            rust_remote::Runner::Server => server::start(config).await,
+            rust_remote::Runner::Client => client::start(config).await,
+        },
+        None => {
+            eprintln!("Error: Take Args");
+            return;
+        }
+    }
+}

src/server.rs

@@ -0,0 +1,113 @@
+use std::io::stdin;
+
+use futures_util::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
+use tokio::net::{TcpListener, TcpStream};
+use tokio_tungstenite::{accept_async, tungstenite::Message, WebSocketStream};
+
+use crate::{Config, Payload};
+type WebSocketSender = SplitSink<WebSocketStream<TcpStream>, Message>;
+type WebSocketReceiver = SplitStream<WebSocketStream<TcpStream>>;
+
+pub async fn start(config: Config) {
+    let listener =
+        match TcpListener::bind(format!("{}:{}", config.server_address, config.port)).await {
+            Ok(listener) => listener,
+            Err(err_val) => {
+                eprintln!("Error: Listener | {}", err_val);
+                return;
+            }
+        };
+
+    loop {
+        if let Ok(connection) = listener.accept().await {
+            let ws_connection = match accept_async(connection.0).await {
+                Ok(ws_connection) => ws_connection,
+                Err(err_val) => {
+                    eprintln!("Error: WebSocket Upgrade | {}", err_val);
+                    continue;
+                }
+            };
+
+            let (mut ws_sender, mut ws_receiver) = ws_connection.split();
+            loop {
+                match payload_from_input().await {
+                    Some(payload) => {
+                        send(payload, &mut ws_sender).await;
+                        let report = receive(&mut ws_receiver).await;
+                        println!("{:#?}", report);
+                    }
+                    None => continue,
+                }
+            }
+        }
+    }
+}
+
+pub async fn payload_from_input() -> Option<Payload> {
+    match get_input() {
+        Some(input) => {
+            let mut args: Vec<String> = input.split_ascii_whitespace().map(String::from).collect();
+            if args.is_empty() {
+                None
+            } else {
+                let mut sudo = false;
+                let mut command = args.remove(0);
+                if command == "sudo" {
+                    if args.is_empty() {
+                        return None;
+                    }
+                    sudo = true;
+                    command = args.remove(0);
+                }
+                Some(Payload {
+                    sudo,
+                    command,
+                    args,
+                })
+            }
+        }
+        None => None,
+    }
+}
+
+pub fn get_input() -> Option<String> {
+    let mut payload_input: String = String::new();
+    match stdin().read_line(&mut payload_input) {
+        Ok(_) => Some(payload_input.trim().to_string()),
+        Err(err_val) => {
+            eprintln!("Error: Read Input | {}", err_val);
+            None
+        }
+    }
+}
+
+pub async fn receive(ws_receiver: &mut WebSocketReceiver) -> Option<String> {
+    match ws_receiver.next().await {
+        Some(message) => match message {
+            Ok(message) => {
+                if let Message::Text(message) = message {
+                    Some(message)
+                } else {
+                    eprintln!("Error: Message Type | {:#?}", message);
+                    None
+                }
+            }
+            Err(err_val) => {
+                eprintln!("Error: Message | {}", err_val);
+                None
+            }
+        },
+        None => {
+            eprintln!("Error: WebSocket Receive");
+            None
+        }
+    }
+}
+
+pub async fn send(payload: Payload, ws_sender: &mut WebSocketSender) {
+    let payload = serde_json::json!(payload);
+    let _ = ws_sender.send(payload.to_string().into()).await;
+}

src/utils.rs

@@ -0,0 +1,48 @@
+use std::{env, fs::File, io::Read};
+
+use crate::{Config, Runner};
+
+pub fn take_args() -> Option<Runner> {
+    let args: Vec<String> = env::args().collect();
+    if args.len() > 1 {
+        match &args[1][..] {
+            "--server" => Some(Runner::Server),
+            "--client" => Some(Runner::Client),
+            _ => None,
+        }
+    } else {
+        None
+    }
+}
+
+pub fn read_config() -> Option<Config> {
+    let mut config_file = match File::open("configs/config.txt") {
+        Ok(config_file) => config_file,
+        Err(_) => return None,
+    };
+    let mut configs = String::new();
+    match config_file.read_to_string(&mut configs) {
+        Ok(_) => {
+            let configs: Vec<String> = configs.split('\n').map(|x| x.to_string()).collect();
+            let server_address = match configs[0].split(':').last() {
+                Some(server_address_unchecked) => match server_address_unchecked.parse() {
+                    Ok(server_address) => server_address,
+                    Err(_) => return None,
+                },
+                None => return None,
+            };
+            let port = match configs[1].split(':').last() {
+                Some(port_unchecked) => match port_unchecked.parse() {
+                    Ok(port) => port,
+                    Err(_) => return None,
+                },
+                None => return None,
+            };
+            Some(Config {
+                server_address,
+                port,
+            })
+        }
+        Err(_) => None,
+    }
+}