RUST Rcon manager

Cargo.toml

@@ -6,3 +6,23 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+bytes = "1.4.0"
+tokio = { version = "1", features = ["full"] }
+tokio-util = { version="0.6", features = ["codec"] }
+tokio-stream = "0.1"
+encoding_rs = "0.8.32"
+log = "0.4"
+env_logger = "0.10.0"
+clap = { version = "4.3.11", features = ["derive"] }
+async-trait = "0.1.71"
+async-channel = "1.9.0"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+futures-util = "0.3.28"
+url = "2.3.1"
+tokio-tungstenite = "0.19.0"
+chrono = "0.4.26"
+colored = "2.0.4"
+job_scheduler = "*"
+tokio-cron-scheduler = "*"
+serde_yaml = "0.9"

sample.yaml

@@ -0,0 +1,16 @@
+ip: SERVER_IP
+port: RCON_PORT
+password: RCON_PASSWORD
+server_log: true
+
+jobs:
+  - name: "Auto Message (every hour)"
+    cron: "0 0 * * * *"
+    commands:
+      - say This is an automatic message.
+  - name: "restart (UST 4:00)"
+    cron: "0 0 4 * * *"
+    commands:
+      - save
+      - restart
+

src/args.rs

@@ -0,0 +1,7 @@
+use clap::Parser;
+
+#[derive(Parser, Debug)]
+pub struct Args {
+    #[arg(short, long)]
+    pub config: String,
+}

src/config.rs

@@ -0,0 +1,28 @@
+use std::{fs::File, io::Read};
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Config {
+    pub ip: String,
+    pub port: String,
+    pub password: String,
+    pub server_log: bool,
+    pub jobs: Vec<CronJob>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CronJob {
+    pub name: String,
+    pub commands: Vec<String>,
+    pub cron: String,
+}
+
+impl Config {
+    pub fn from_file(path: String) -> Self {
+        let mut config_file = File::open(path).expect("Failed to open config file.");
+        let mut config = String::default();
+        config_file.read_to_string(&mut config).unwrap();
+        serde_yaml::from_str::<Config>(&config).unwrap()
+    }
+}

src/cron.rs

@@ -0,0 +1,33 @@
+use std::sync::Arc;
+
+use crate::logger::log;
+use tokio::sync::Mutex;
+use tokio_cron_scheduler::{Job, JobScheduler};
+
+use crate::{config::Config, rcon::Rcon, rcon_web::WebRcon};
+
+pub async fn create_scheduler(config: Config, web_rcon: Arc<Mutex<WebRcon>>) {
+    let scheduler = JobScheduler::new().await.unwrap();
+    for job in config.jobs {
+        let web_rcon_clone = web_rcon.clone();
+        let job_clone = job.clone();
+        scheduler
+            .add(
+                Job::new_async(job.cron.as_str(), move |_uuid, _l| {
+                    let web_rcon = web_rcon_clone.clone();
+                    let job = job_clone.clone();
+                    Box::pin(async move {
+                        log("CRON_JOB", "cyan", &format!("Running Job: {}", job.name));
+                        for command in job.commands {
+                            web_rcon.lock().await.execute(&command).await;
+                        }
+                    })
+                })
+                .unwrap(),
+            )
+            .await
+            .unwrap();
+    }
+
+    scheduler.start().await.unwrap();
+}

src/logger.rs

@@ -0,0 +1,9 @@
+use chrono::Local;
+use colored::Colorize;
+
+pub fn log(kind: &str, color: &str, message: &str) {
+    let time = Local::now();
+    let time = time.format("%m-%d %H:%M:%S");
+
+    println!("[{} {}]: {}", time, kind.color(color), message);
+}

src/main.rs

@@ -1,3 +1,76 @@
-fn main() {
-    println!("Hello, world!");
+mod args;
+mod config;
+mod cron;
+mod logger;
+mod rcon;
+mod rcon_pure;
+mod rcon_web;
+
+use args::Args;
+use clap::Parser;
+use config::Config;
+use cron::create_scheduler;
+use logger::log;
+use rcon::Rcon;
+use rcon_web::WebRcon;
+use std::sync::Arc;
+use tokio::sync::Mutex;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+
+    let config = Config::from_file(args.config);
+
+    let web_rcon = Arc::new(Mutex::new(
+        WebRcon::connect(&config.ip, &config.port, &config.password).await,
+    ));
+
+    log("INFO", "bright_black", "Connected.");
+
+    let receiver = { web_rcon.lock().await.receiver.clone() };
+
+    create_scheduler(config.clone(), web_rcon.clone()).await;
+
+    tokio::spawn(async move {
+        loop {
+            let mut input = String::default();
+            std::io::stdin().read_line(&mut input).unwrap();
+            web_rcon.lock().await.execute(&input).await;
+        }
+    });
+
+    while let Ok(data) = receiver.recv().await {
+        if config.server_log {
+            match data {
+                rcon::RconMessage::Chat {
+                    Channel,
+                    Message,
+                    UserId,
+                    Username,
+                    ..
+                } => {
+                    log(
+                        "CHAT",
+                        "green",
+                        &format!("<{}[{}]>({}) {}", Username, UserId, Channel, Message),
+                    );
+                }
+                rcon::RconMessage::Generic { message } => {
+                    log("INFO", "bright_black", &message);
+                }
+                rcon::RconMessage::Warning { message } => {
+                    log("WARN", "red", &message);
+                }
+                rcon::RconMessage::Error { message } => {
+                    log("ERROR", "red", &message);
+                }
+                rcon::RconMessage::Disconnected => {
+                    std::process::exit(0);
+                }
+            }
+        }
+    }
+
+    Ok(())
 }

src/rcon.rs

@@ -0,0 +1,52 @@
+use async_trait::async_trait;
+use serde::{Deserialize, Serialize};
+
+use crate::rcon_web::WebRconMessage;
+
+#[async_trait]
+pub trait Rcon {
+    async fn connect(ip: &str, port: &str, password: &str) -> Self;
+    async fn execute(&mut self, command: &str);
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(untagged)]
+#[allow(non_snake_case)]
+pub enum RconMessage {
+    Generic {
+        message: String,
+    },
+    Warning {
+        message: String,
+    },
+    Error {
+        message: String,
+    },
+    Chat {
+        Channel: u64,
+        Message: String,
+        UserId: String,
+        Username: String,
+        Color: String,
+        Time: u64,
+    },
+    Disconnected,
+}
+
+impl RconMessage {
+    pub fn from_webrcon_message(message: WebRconMessage) -> Self {
+        match &*message.Type {
+            "Chat" => serde_json::from_str(&message.Message).unwrap(),
+            "Generic" => RconMessage::Generic {
+                message: message.Message,
+            },
+            "Warning" => RconMessage::Warning {
+                message: message.Message,
+            },
+            "Error" => RconMessage::Error {
+                message: message.Message,
+            },
+            _ => panic!("Cannot parse WebRconMessage. {}", message.Type),
+        }
+    }
+}

src/rcon_pure.rs

@@ -0,0 +1,38 @@
+use bytes::{Buf, BufMut};
+use encoding_rs::UTF_8;
+
+#[derive(Debug, Clone)]
+pub struct RconPacket {
+    pub id: i32,
+    pub packet_type: i32,
+    pub body: String,
+}
+
+#[allow(dead_code)]
+impl RconPacket {
+    pub fn from_bytes(bytes: &[u8]) -> Self {
+        let mut bytes = bytes;
+        let id = bytes.get_i32_le();
+        let packet_type = bytes.get_i32_le();
+        let (string, _, _) = UTF_8.decode(&bytes[..bytes.len() - 2]);
+        let body = string.to_string();
+
+        Self {
+            id,
+            packet_type,
+            body,
+        }
+    }
+
+    pub fn to_bytes(&self) -> Vec<u8> {
+        let mut buf = Vec::new();
+
+        buf.put_i32_le(8 + self.body.as_bytes().len() as i32);
+        buf.put_i32_le(self.id);
+        buf.put_i32_le(self.packet_type);
+        buf.put(self.body.as_bytes());
+        buf.put_u8(0);
+
+        buf
+    }
+}

src/rcon_web.rs

@@ -0,0 +1,86 @@
+use async_channel::{Receiver, Sender};
+use async_trait::async_trait;
+use futures_util::{
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt,
+};
+use serde::{Deserialize, Serialize};
+use tokio::net::TcpStream;
+use tokio_tungstenite::{
+    connect_async, tungstenite::protocol::Message, MaybeTlsStream, WebSocketStream,
+};
+
+use crate::rcon::{Rcon, RconMessage};
+
+#[allow(non_snake_case)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WebRconMessage {
+    pub Message: String,
+    pub Identifier: i32,
+    pub Type: String,
+    pub Stacktrace: Option<String>,
+}
+
+#[derive(Debug)]
+pub struct WebRcon {
+    pub receiver: Receiver<RconMessage>,
+    write_stream: SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>,
+}
+
+#[async_trait]
+impl Rcon for WebRcon {
+    async fn connect(ip: &str, port: &str, password: &str) -> Self {
+        let url = url::Url::parse(&format!("ws://{}:{}/{}", ip, port, password)).unwrap();
+
+        let (ws_stream, _) = connect_async(url).await.expect("Failed to connect");
+        let (write, read) = ws_stream.split();
+
+        let (sender, receiver) = async_channel::unbounded();
+        tokio::spawn(WebRcon::receive_data(read, sender));
+
+        Self {
+            receiver,
+            write_stream: write,
+        }
+    }
+
+    async fn execute(&mut self, command: &str) {
+        let message = WebRconMessage {
+            Message: command.to_string(),
+            Identifier: 10,
+            Type: String::default(),
+            Stacktrace: None,
+        };
+
+        self.write_stream
+            .send(serde_json::to_string(&message).unwrap().into())
+            .await
+            .unwrap();
+    }
+}
+
+impl WebRcon {
+    async fn receive_data(
+        stream_read: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
+        sender: Sender<RconMessage>,
+    ) {
+        stream_read
+            .for_each(|message| async {
+                if let Ok(message) = message {
+                    if let Ok(data) = message.into_text() {
+                        let message: WebRconMessage = serde_json::from_str(&data).unwrap();
+
+                        sender
+                            .send(RconMessage::from_webrcon_message(message))
+                            .await
+                            .unwrap();
+                    } else {
+                        sender.send(RconMessage::Disconnected).await.unwrap();
+                    }
+                } else {
+                    sender.send(RconMessage::Disconnected).await.unwrap();
+                }
+            })
+            .await;
+    }
+}