From 1c7170bb81f58a89a56a671fbaa84b3f652ef1bd Mon Sep 17 00:00:00 2001 From: Byte Dice Date: Sat, 8 Feb 2025 15:50:33 +0100 Subject: [PATCH] Updated arg system & made websocked somehow work??? --- Cargo.toml | 2 ++ src/bk_week_cmds.rs | 1 - src/cmds.rs | 3 +- src/macros.rs | 13 ++++++++ src/main.rs | 65 +++++++++++++++++++++----------------- src/python.rs | 7 ++-- src/python/bot.py | 4 +-- src/python/main.py | 9 +++--- src/python/py_websocket.py | 7 ++-- src/websocket.rs | 18 ++++++----- 10 files changed, 77 insertions(+), 52 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 5b6175f..a732d9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,10 +6,12 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +clap = { version = "4.5.28", features = ["derive"] } futures = "0.3.31" poise = "0.6.1" pyo3 = "0.23.4" rand = "0.9.0" +serde = "1.0.217" serde_json = "1.0.138" tokio = { version = "1.43.0", features = ["rt-multi-thread"] } tokio-tungstenite = "0.26.1" diff --git a/src/bk_week_cmds.rs b/src/bk_week_cmds.rs index 521900d..01737ab 100644 --- a/src/bk_week_cmds.rs +++ b/src/bk_week_cmds.rs @@ -3,7 +3,6 @@ use crate::messages::send_msg; use std::fs; - #[poise::command(slash_command, prefix_command)] pub async fn bk_week_help( ctx: Context<'_>, diff --git a/src/cmds.rs b/src/cmds.rs index 1fd2120..c11e8e1 100644 --- a/src/cmds.rs +++ b/src/cmds.rs @@ -25,8 +25,7 @@ pub async fn stop( #[description = "Type \"i want to stop the bot now\" to confirm."] confirmation: Option, ) -> Result<(), Error> { - let dev_enabled = ctx.data().dev; - let should_stop = dev_enabled + let should_stop = ctx.data().args.dev || confirmation.unwrap_or_else(|| "".to_string()).to_lowercase() == "i want to stop the bot now"; let is_creator = ctx.author().id == UserId::new(ctx.data().creator_id); diff --git a/src/macros.rs b/src/macros.rs index f981ade..819226c 100644 --- a/src/macros.rs +++ b/src/macros.rs @@ -22,3 +22,16 @@ macro_rules! rs_errln { process::exit(-1); }; } + + +#[macro_export] +macro_rules! errln { + ($($arg:tt)*) => { + println!("{}ERROR{} - {}", + "\x1b[41m", + "\x1b[0m", + format!($($arg)*) + ); + process::exit(-1); + }; +} \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index e99330f..904f00f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ #![warn(unused_extern_crates)] mod cmds; +#[allow(unused_variables)] mod bk_week_cmds; mod events; mod messages; @@ -8,21 +9,35 @@ mod python; mod macros; mod websocket; -use std::env; use std::process; use std::thread; -use std::fs; -use tokio::runtime::Runtime; -use poise::serenity_prelude::Client; +use clap::Parser; use poise::serenity_prelude as serenity; +use poise::serenity_prelude::Client; +use serde::Serialize; use serde_json::Value; +use tokio::runtime::Runtime; +use serde_json; + + +#[derive(Parser, Serialize, Clone)] +struct Args { + #[arg(short = 'p', long, default_value = "2920", help = "Sets the port number, e.g 2200.")] + port: u16, + #[arg(long, help = "Runs only the Python part of the program.")] + py: bool, + #[arg(long, help = "Runs only the Rust part of the program.")] + rs: bool, + #[arg(short = 'd', long, help = "Enables dev mode. Dev mode shows more debug info and turns of certain security measures.")] + dev: bool +} struct Data { - dev: bool, ball_prompts: [Vec; 2], creator_id: u64, reddit_data: Option, + args: Args // TODO: schedules } type Error = Box; @@ -31,38 +46,32 @@ type Context<'a> = poise::Context<'a, Data, Error>; #[tokio::main] async fn main() { - let args: Vec = env::args().collect(); + let args = ::parse(); + let args_str = serde_json::to_string(&args).expect("Error serializing args to JSON"); - if args.contains(&"--h".to_string()) || args.contains(&"--help".to_string()) { - let help = fs::read_to_string("./help.txt").unwrap_or_else(|_| "No help.txt file found.".to_string()); - println!("HELP MENU:\n{}", help); - process::exit(1); - } - if args.contains(&"--dev".to_string()) { println!("----- DEV MODE ENABLED -----"); } + if args.dev { println!("----- DEV MODE ENABLED -----"); rs_println!("ARGS: {}", args_str); } - if args.contains(&"--py".to_string()) - && !args.contains(&"--rs".to_string()) - { + if args.py && !args.rs { println!("----- PYTHON ONLY MODE -----"); - let _ = python::start(args); + let _ = python::start(args_str); process::exit(0); } - else if args.contains(&"--rs".to_string()) - && ! args.contains(&"--py".to_string()) - { + else if args.rs && ! args.py { println!("----- RUST ONLY MODE -----"); start(args).await; process::exit(0); } + else if args.py && args.rs { + errln!("Invalid arguments: Arguments cannot include both --rs and --py."); + } let rt = Runtime::new().unwrap(); - let rust_args = args.clone(); - let python_args = args.clone(); + let python_args = args_str; let rust = thread::spawn(move || { rt.block_on(async { - websocket::start(rust_args.clone()).await; - start(rust_args).await; + websocket::start(args.clone()).await; + start(args).await; }); }); @@ -75,9 +84,7 @@ async fn main() { } -async fn start(args: Vec) { - rs_println!("ARGS: {:?}", &args[1..]); - +async fn start(args: Args) { let data = gen_data(args); let mut bot = gen_bot(data).await; @@ -86,7 +93,7 @@ async fn start(args: Vec) { } -fn gen_data(args: Vec) -> Data { +fn gen_data(args: Args) -> Data { let ball_classic_str = std::fs::read_to_string("./data/8-ball_classic.txt").unwrap(); let ball_quirk_str = std::fs::read_to_string("./data/8-ball_quirky.txt").unwrap(); @@ -94,10 +101,10 @@ fn gen_data(args: Vec) -> Data { let ball_quirk: Vec = ball_quirk_str .lines().map(String::from).collect(); return Data { - dev: args.contains(&"--dev".to_string()), ball_prompts: [ball_classic, ball_quirk], creator_id: 697149665166229614, - reddit_data: None + reddit_data: None, + args }; } diff --git a/src/python.rs b/src/python.rs index f55dc4a..730ebc6 100644 --- a/src/python.rs +++ b/src/python.rs @@ -8,13 +8,14 @@ use pyo3::prelude::*; use pyo3::types::PyList; -pub fn start(args: Vec) -> PyResult<()> { +pub fn start(args: String) -> PyResult<()> { rs_println!("Running Python program..."); let path = concat!(env!("CARGO_MANIFEST_DIR"), "\\src\\python"); let code = get_code(&(path.to_owned() + "\\main.py")); - let app_path = CString::new(format!("args = {:?}\n{}", args, code)).unwrap(); + let py_args = args.replace(":true", ":True").replace(":false", ":False"); + let app_path = CString::new(format!("args = {}\n{}", py_args, code)).unwrap(); pyo3::prepare_freethreaded_python(); @@ -30,7 +31,7 @@ pub fn start(args: Vec) -> PyResult<()> { return app.call0(py); }); - println!("py: {}", from_python?); + if from_python.is_err() { println!("py: {:?}", from_python); } return Ok(()); } diff --git a/src/python/bot.py b/src/python/bot.py index 30b8b33..bbd4ac0 100644 --- a/src/python/bot.py +++ b/src/python/bot.py @@ -6,7 +6,7 @@ import os from macros import * class Bot: - args: list[str] = ["NO_RUST", "--dev", "--py"] + args: dict = {"NO_RUST": True, "dev": True, "py": True, "port": 2920} password: str = os.environ.get("ASSISTANT_R_PASS") secret: str = os.environ.get("ASSISTANT_R_TOKEN") @@ -26,5 +26,5 @@ class Bot: data_f: TextIOWrapper = None data: dict = {} - def set_args(self, args: list[str]): + def set_args(self, args: dict): self.args = args \ No newline at end of file diff --git a/src/python/main.py b/src/python/main.py index f90b787..3c092ad 100644 --- a/src/python/main.py +++ b/src/python/main.py @@ -15,15 +15,16 @@ def main(): except NameError: py_print("No command args found from Rust. Don't worry though, we have backup in place.") - py_print("ARGS:", str(bot.args[1:])) + if bot.args["dev"]: + py_print("ARGS:", str(bot.args)) py_print("Reading data...") data.read_data(bot) - if "--py" not in bot.args: + if True: #bot.args["py"]: py_print("Connecting to local websocket...") - asyncio.run(py_websocket.websocket_client()) + asyncio.run(py_websocket.websocket_client(bot)) - py_websocket.send_message("[Connection test] Hello from Python!") + #py_websocket.send_message("[Connection test] Hello from Python!") main() \ No newline at end of file diff --git a/src/python/py_websocket.py b/src/python/py_websocket.py index 0c6995f..42820d9 100644 --- a/src/python/py_websocket.py +++ b/src/python/py_websocket.py @@ -1,5 +1,6 @@ import websockets from macros import py_print +import bot as botPy ws_global = None @@ -9,11 +10,11 @@ async def send_message(message: str): await ws_global.send(message) -async def websocket_client(): +async def websocket_client(bot: botPy.Bot): global ws_global - async with websockets.connect("ws://127.0.0.1:9001") as ws: + async with websockets.connect(f"ws://127.0.0.1:{bot.args["port"]}") as ws: ws_global = ws - py_print("Connected webSocket server on ws://127.0.0.1:9001") + py_print(f"Connected webSocket server on ws://127.0.0.1:{bot.args["port"]}") while True: response = await ws.recv() diff --git a/src/websocket.rs b/src/websocket.rs index 06b4b0e..dfe4090 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -6,6 +6,7 @@ use futures::StreamExt; use std::sync::Arc; use crate::rs_println; +use crate::Args; type Sender = Arc, tungstenite::Message>>>>; @@ -32,16 +33,17 @@ pub async fn send_msg(msg: &str) { } -pub async fn start(args: Vec) { +pub async fn start(args: Args) { rs_println!("Starting local websocket..."); - let listener = TcpListener::bind("127.0.0.1:9001").await.unwrap(); - rs_println!("WebSocket server running on ws://127.0.0.1:9001"); + let ip = format!("127.0.0.1:{}", args.port); + let listener = TcpListener::bind(&ip).await.unwrap(); + rs_println!("WebSocket server running on ws://{}", ip); tokio::spawn(handle_connections(listener, args)); } -async fn handle_connections(listener: TcpListener, args: Vec) { +async fn handle_connections(listener: TcpListener, args: Args) { while let Ok((stream, _)) = listener.accept().await { let ws_stream = accept_async(stream).await.unwrap(); let (sender, mut receiver) = ws_stream.split(); @@ -50,13 +52,13 @@ async fn handle_connections(listener: TcpListener, args: Vec) { set_sender(sender_arc.clone()).await; while let Some(Ok(msg)) = receiver.next().await { - handle_message(msg, &args).await; + handle_message(msg, args.clone()).await; } } } -async fn handle_message(msg: tungstenite::protocol::Message, args: &[String]) { +async fn handle_message(msg: tungstenite::protocol::Message, args: Args) { match msg { tungstenite::Message::Text(text) => { rs_println!("Received from Python: {}", text); @@ -70,12 +72,12 @@ async fn handle_message(msg: tungstenite::protocol::Message, args: &[String]) { } } tungstenite::Message::Binary(bytes) => { - if args.contains(&"--dev".to_string()) { + if args.dev { rs_println!("[Binary] from Python: {:?}", bytes); } } _ => { - if args.contains(&"--dev".to_string()) { + if args.dev { rs_println!("Received from Python: [UNKNOWN / OTHER]"); } }