From ab02ac1cf656ea8e09877727bb84abaa12eaff09 Mon Sep 17 00:00:00 2001 From: ByteDice Date: Sun, 2 Mar 2025 12:18:49 +0100 Subject: [PATCH] put ASSISTANT_OWNERS in the ctx.data() --- README.md | 4 ++-- src/main.rs | 30 ++++++++++++++++-------------- src/messages.rs | 18 +++++------------- src/websocket.rs | 12 ++++++------ 4 files changed, 29 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 08e4de0..53dfde2 100644 --- a/README.md +++ b/README.md @@ -38,10 +38,10 @@ You can install Python modules by running `$ pip install {module}` or `$ python | `ASSISTANT_R_NAME` | The username of the Reddit bot/account. | | `ASSISTANT_R_PASS` | The password for the Reddit bot/account. | | `ASSISTANT_OWNERS` | (OPTIONAL) A list of Discord user IDs that "own" the bot. Separate each ID with a single comma and **no** spaces. This will allow the specified user IDs to run root commands such as `/stop`, it will also DM these users when *certain* errors occur. | -| `ASSISTANT_BK_WEEK_MODS` | (OPTIONAL) Same format as `ASSISTANT_OWNERS` but for | +| `ASSISTANT_BK_WEEK_MODS` | (OPTIONAL) Same format as `ASSISTANT_OWNERS` but for people who are allowed to use the `/bk_week` commands. | ### Required permissions: -**These are automatically set if you use the [official invite link](https://discord.com/oauth2/authorize?client_id=1212127255795335208&permissions=84992&integration_type=0&scope=bot).** +**These are automatically set if you use the [official invite link](https://discord.com/oauth2/authorize?client_id=1212127255795335208&permissions=84992&integration_type=0&scope=bot) or an invite link with the permissions integer set to `84992`.** * Send Messages * Read Message History * View Channels diff --git a/src/main.rs b/src/main.rs index 59ff336..ae3c98b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -58,6 +58,7 @@ type Context<'a> = poise::Context<'a, Data, Error>; struct Data { + owners: Vec, ball_prompts: [Vec; 2], reddit_data: Mutex>, discord_data: Mutex>, @@ -74,6 +75,13 @@ async fn main() { let args = ::parse(); let args_str = serde_json::to_string(&args).expect("Error serializing args to JSON"); + let own_env = std::env::var("ASSISTANT_OWNERS").expect("Missing ASSISTANT_OWNERS env var!"); + let own_vec_str: Vec = own_env.split(",").map(String::from).collect(); + let own_vec_u64: Vec = own_vec_str + .iter() + .filter_map(|s| Some(s.parse::().expect("Failed to parse ASSISTANT_OWNERS. Invalid syntax."))) + .collect(); + if args.test { println!("----- USING TEST BOT -----"); } if args.dev { println!("----- DEV MODE ENABLED -----"); } if args.dev && args.wipe { println!("----- \"DON'T WORRY ABOUT IT\" MODE ENABLED -----"); } @@ -88,7 +96,7 @@ async fn main() { else if args.rs && ! args.py { println!("----- RUST ONLY MODE -----"); rs_println!("ARGS: {}", args_str); - start(args).await; + start(args, own_vec_u64.clone()).await; process::exit(0); } else if args.py && args.rs { @@ -103,8 +111,8 @@ async fn main() { let rust = thread::spawn(move || { rt.block_on(async { - websocket::start(rust_args.clone()).await; - start(rust_args).await; + websocket::start(rust_args.clone(), own_vec_u64.clone()).await; + start(rust_args, own_vec_u64).await; }); }); @@ -125,8 +133,8 @@ async fn main() { } -async fn start(args: Args) { - let data = gen_data(args.clone()).await; +async fn start(args: Args, owners: Vec) { + let data = gen_data(args.clone(), owners).await; let mut bot = gen_bot(data, args).await; rs_println!("Starting bot..."); @@ -134,7 +142,7 @@ async fn start(args: Args) { } -async fn gen_data(args: Args) -> Data { +async fn gen_data(args: Args, owners: Vec) -> 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(); let bk_mods_str = std::fs::read_to_string("./data/bk_mods.json").unwrap(); @@ -144,6 +152,7 @@ async fn gen_data(args: Args) -> Data { let bk_mods: Value = serde_json::from_str(&bk_mods_str).unwrap(); let data = Data { + owners, ball_prompts: [ball_classic, ball_quirk], bk_mods_json: bk_mods, reddit_data: None.into(), @@ -174,14 +183,7 @@ async fn gen_bot(data: Data, args: Args) -> Client { let token_end_len = token[peek_len..].len(); rs_println!("Token: {}{}", token_peek, "*".repeat(token_end_len)); - let own_env = std::env::var("ASSISTANT_OWNERS").expect("Missing ASSISTANT_OWNERS env var!"); - let own_vec_str: Vec = own_env.split(",").map(String::from).collect(); - let own_vec_u64: Vec = own_vec_str - .iter() - .filter_map(|s| Some(s.parse::().expect("Failed to parse ASSISTANT_OWNERS. Invalid syntax."))) - .collect(); - - let own: HashSet = own_vec_u64.into_iter().map(UserId::from).collect(); + let own: HashSet = data.owners.clone().into_iter().map(UserId::from).collect(); let framework = poise::Framework::builder() .options(poise::FrameworkOptions { diff --git a/src/messages.rs b/src/messages.rs index c45d5cd..97f4931 100644 --- a/src/messages.rs +++ b/src/messages.rs @@ -178,25 +178,17 @@ pub async fn http_edit_msg( } -pub async fn send_dm(msg: String, args: Args) { - let uids = env::var("ASSISTANT_OWNERS"); - if uids.is_err() { return; } - - let uids_split: Vec = uids.unwrap() - .split(',') - .map(|s| s.to_string()) - .collect(); - +pub async fn send_dm(msg: String, args: Args, owners: Vec) { let token: String = - if !args.test { std::env::var("ASSISTANT_TOKEN") .expect("Missing ASSISTANT_TOKEN env var!") } - else { std::env::var("ASSISTANT_TOKEN_TEST").expect("Missing ASSISTANT_TOKEN_TEST env var!") }; + if !args.test { env::var("ASSISTANT_TOKEN") .expect("Missing ASSISTANT_TOKEN env var!") } + else { env::var("ASSISTANT_TOKEN_TEST").expect("Missing ASSISTANT_TOKEN_TEST env var!") }; let http = Http::new(&token); let c_msg = CreateMessage::new().content(msg); - for uid in uids_split { - let user = UserId::new(uid.parse::().unwrap()); + for uid in owners { + let user = UserId::new(uid); let _ = user.dm(http.as_ref(), c_msg.clone()).await; } } diff --git a/src/websocket.rs b/src/websocket.rs index 8fe5775..a97d554 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -98,17 +98,17 @@ async fn receive_response() -> Option { } -pub async fn start(args: Args) { +pub async fn start(args: Args, owners: Vec) { rs_println!("Starting local websocket..."); 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)); + tokio::spawn(handle_connections(listener, args, owners)); } -async fn handle_connections(listener: TcpListener, args: Args) { +async fn handle_connections(listener: TcpListener, args: Args, owners: Vec) { while let Ok((stream, _)) = listener.accept().await { let ws_stream = accept_async(stream).await.unwrap(); let (sender, receiver) = ws_stream.split(); @@ -120,13 +120,13 @@ async fn handle_connections(listener: TcpListener, args: Args) { set_receiver(receiver_arc.clone()).await; while let Some(Ok(msg)) = receiver_arc.lock().await.as_mut().unwrap().next().await { - handle_message(msg, args.clone()).await; + handle_message(msg, args.clone(), owners.clone()).await; } } } -async fn handle_message(msg: tungstenite::protocol::Message, args: Args) { +async fn handle_message(msg: tungstenite::protocol::Message, args: Args, owners: Vec) { match msg { tungstenite::Message::Text(text) => { rs_println!("Received from Python: {}", text); @@ -134,7 +134,7 @@ async fn handle_message(msg: tungstenite::protocol::Message, args: Args) { if text.starts_with("json:") { let t_json: Value = serde_json::from_str(&text[5..]).unwrap(); if t_json.get("error").is_some() { - send_dm("Unknown internal Python error occurred!".to_string(), args).await; + send_dm("Unknown internal Python error occurred!".to_string(), args, owners).await; } }