put ASSISTANT_OWNERS in the ctx.data()

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

src/main.rs

@@ -58,6 +58,7 @@ type Context<'a> = poise::Context<'a, Data, Error>;
 
 
 struct Data {
+  owners: Vec<u64>,
   ball_prompts: [Vec<String>; 2],
   reddit_data: Mutex<Option<Value>>,
   discord_data: Mutex<Option<Value>>,
@@ -74,6 +75,13 @@ async fn main() {
   let args = <Args as clap::Parser>::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<String> = own_env.split(",").map(String::from).collect();
+  let own_vec_u64: Vec<u64> = own_vec_str
+    .iter()
+    .filter_map(|s| Some(s.parse::<u64>().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;
+      websocket::start(rust_args.clone(), own_vec_u64.clone()).await;
-      start(rust_args).await;
+      start(rust_args, own_vec_u64).await;
     });
   });
 
@@ -125,8 +133,8 @@ async fn main() {
 }
 
 
-async fn start(args: Args) {
+async fn start(args: Args, owners: Vec<u64>) {
-  let data = gen_data(args.clone()).await;
+  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<u64>) -> 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: HashSet<UserId> = data.owners.clone().into_iter().map(UserId::from).collect();
-  let own_vec_str: Vec<String> = own_env.split(",").map(String::from).collect();
-  let own_vec_u64: Vec<u64> = own_vec_str
-    .iter()
-    .filter_map(|s| Some(s.parse::<u64>().expect("Failed to parse ASSISTANT_OWNERS. Invalid syntax.")))
-    .collect();
-
-  let own: HashSet<UserId> = own_vec_u64.into_iter().map(UserId::from).collect();
 
   let framework = poise::Framework::builder()
     .options(poise::FrameworkOptions {

src/messages.rs

@@ -178,25 +178,17 @@ pub async fn http_edit_msg(
 }
 
 
-pub async fn send_dm(msg: String, args: Args) {
+pub async fn send_dm(msg: String, args: Args, owners: Vec<u64>) {
-  let uids = env::var("ASSISTANT_OWNERS");
-  if uids.is_err() { return; }
-  
-  let uids_split: Vec<String> = uids.unwrap()
-    .split(',')
-    .map(|s| s.to_string())
-    .collect();
-
   let token: String =
-    if !args.test { std::env::var("ASSISTANT_TOKEN")     .expect("Missing ASSISTANT_TOKEN env var!") }
+    if !args.test { env::var("ASSISTANT_TOKEN")     .expect("Missing ASSISTANT_TOKEN env var!") }
-    else          { std::env::var("ASSISTANT_TOKEN_TEST").expect("Missing ASSISTANT_TOKEN_TEST 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 {
+  for uid in owners {
-    let user = UserId::new(uid.parse::<u64>().unwrap());
+    let user = UserId::new(uid);
     let _ = user.dm(http.as_ref(), c_msg.clone()).await;
   }
 }

src/websocket.rs

@@ -98,17 +98,17 @@ async fn receive_response() -> Option<Value> {
 }
 
 
-pub async fn start(args: Args) {
+pub async fn start(args: Args, owners: Vec<u64>) {
   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<u64>) {
   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<u64>) {
   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;
         }
       }